def tokenize(stream,
             keys=None,
             vocab_type='subword',
             vocab_file=None,
             vocab_dir=None,
             n_reserved_ids=0):
  """Tokenize examples from the stream.

  This function assumes that `stream` generates either strings or tuples/dicts
  containing strings at some `keys`. This function maps these strings to
  numpy arrays of integers -- the tokenized version of each string.

  Args:
    stream: A python generator yielding strings, tuples or dicts.
    keys: which keys of the tuple/dict to tokenize (by default: all)
    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;
      This is common for example when reserving the 0 for padding and 1 for EOS,
      but it's only needed if these symbols are not already included (and thus
      reserved) in the vocab_file.

  Yields:
    Examples from stream with strings at `keys` replaced by np.arrays of
    integers -- the tokenized version of these strings.
  """
  vocab = _get_vocab(vocab_type, vocab_file, vocab_dir)
  for example in stream:
    if isinstance(example, (list, tuple)):
      new_example = []
      for i, x in enumerate(example):
        if keys is None or i in keys:
          new_example.append(np.array(vocab.encode(x)) + n_reserved_ids)
        else:
          new_example.append(x)
      output = tuple(new_example)
      yield output
    elif isinstance(example, dict):
      new_example = {}
      for k in example:
        if keys is None or k in keys:
          new_example[k] = np.array(vocab.encode(example[k])) + n_reserved_ids
        else:
          new_example[k] = example[k]
      yield new_example
    else:
      output = np.array(vocab.encode(example)) + n_reserved_ids
      yield output
Tokenize examples from the stream.

  This function assumes that `stream` generates either strings or tuples/dicts
  containing strings at some `keys`. This function maps these strings to
  numpy arrays of integers -- the tokenized version of each string.

  Args:
    stream: A python generator yielding strings, tuples or dicts.
    keys: which keys of the tuple/dict to tokenize (by default: all)
    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;
      This is common for example when reserving the 0 for padding and 1 for EOS,
      but it's only needed if these symbols are not already included (and thus
      reserved) in the vocab_file.

  Yields:
    Examples from stream with strings at `keys` replaced by np.arrays of
    integers -- the tokenized version of these strings.
  
trax/data/tf_inputs.py