import tensorflow as tf

@gin.configurable(module='trax.data', denylist=['dataset', 'training'])
def add_eos_to_output_features(dataset,
                               training,
                               output_features='targets',
                               eos=1):
  """Adds `EOS` to all features in `output_features`."""
  del training
  if not isinstance(output_features, (list, tuple)):
    output_features = [output_features]

  def add_eos(x):
    for output_feature in output_features:
      x[output_feature] = tf.concat([x[output_feature], [eos]], axis=0)
    return x

  return dataset.map(add_eos)


@gin.configurable(module='trax.data', denylist=['dataset', 'training'])
def generic_text_dataset_preprocess_fn(dataset,
                                       training=True,
                                       text_preprocess_fns=None,
                                       token_preprocess_fns=None,
                                       spm_path=None,
