def vocab_size(vocab_type='subword',
               vocab_file=None,
               vocab_dir=None,
               n_reserved_ids=0):
  """Returns the size of the vocabulary (number of symbols used).

  This function can be used to set the size of the final layers of a model that
  needs to predict symbols from a given vocabulary. More precisely, if this
  function returns N then the last layer size should be set to at least N (it
  can be more). Note that this function does take reserved IDs into account.

  Args:
    vocab_type: Type of vocabulary, one of: 'subword', 'sentencepiece', 'char'.
    vocab_file: Name of the vocabulary file.
    vocab_dir: Directory which contains the vocabulary file.
    n_reserved_ids: An int, offset added so 0, ..., n_reserved_ids-1 are unused.

  Returns:
    An integer, the number of symbols used (including reserved IDs).
  """
  vocab = _get_vocab(vocab_type, vocab_file, vocab_dir)
  return vocab.vocab_size + n_reserved_ids
Returns the size of the vocabulary (number of symbols used).

  This function can be used to set the size of the final layers of a model that
  needs to predict symbols from a given vocabulary. More precisely, if this
  function returns N then the last layer size should be set to at least N (it
  can be more). Note that this function does take reserved IDs into account.

  Args:
    vocab_type: Type of vocabulary, one of: 'subword', 'sentencepiece', 'char'.
    vocab_file: Name of the vocabulary file.
    vocab_dir: Directory which contains the vocabulary file.
    n_reserved_ids: An int, offset added so 0, ..., n_reserved_ids-1 are unused.

  Returns:
    An integer, the number of symbols used (including reserved IDs).
  
trax/data/tf_inputs.py