def _get_vocab(vocab_type='subword', vocab_file=None, vocab_dir=None,
               extra_ids=0):
  """Gets the vocabulary object for tokenization; see tokenize for details."""
  if vocab_type not in [
      'char', 'subword', 'sentencepiece', 'bert', 'bert-lowercase'
  ]:
    raise ValueError(
        'vocab_type must be "subword", "char", "sentencepiece", "bert" or "bert-lowercase" '
        f'but got {vocab_type}')

  if vocab_type == 'char':
    # Note that we set num_reserved_ids=0 below. We could instead pass
    # the value n_reserved_ids from tokenize here -- ByteTextEncoder does
    # exactly the same thing as tokenize above, ie., adds num_reserved_ids.
    return text_encoder.ByteTextEncoder(num_reserved_ids=0)

  vocab_dir = vocab_dir or 'gs://trax-ml/vocabs/'
  path = os.path.join(vocab_dir, vocab_file)

  if vocab_type == 'subword':
    return text_encoder.SubwordTextEncoder(path)

  if vocab_type == 'bert':
    return text_encoder.BertEncoder(path, do_lower_case=False)

  if vocab_type == 'bert-lowercase':
    return text_encoder.BertEncoder(path, do_lower_case=True)

  assert vocab_type == 'sentencepiece'
  return t5_data().SentencePieceVocabulary(sentencepiece_model_file=path,
                                           extra_ids=extra_ids)
Gets the vocabulary object for tokenization; see tokenize for details.
trax/data/tf_inputs.py