def encode_dataset(dataset, vocabulary):
  """Encode from strings to token ids.

  Args:
    dataset: a tf.data.Dataset with string values.
    vocabulary: a mesh_tensorflow.transformer.Vocabulary
  Returns:
    a tf.data.Dataset with integer-vector values ending in EOS=1
  """
  def encode(features):
    return {k: vocabulary.encode_tf(v) for k, v in features.items()}
  return dataset.map(encode, num_parallel_calls=tf.data.experimental.AUTOTUNE)
Encode from strings to token ids.

  Args:
    dataset: a tf.data.Dataset with string values.
    vocabulary: a mesh_tensorflow.transformer.Vocabulary
  Returns:
    a tf.data.Dataset with integer-vector values ending in EOS=1