import tensorflow as tf

  return t5_data().SentencePieceVocabulary(sentencepiece_model_file=path,
                                           extra_ids=extra_ids)


# Makes the function accessible in gin configs, even with all args denylisted.
@gin.configurable(module='trax.data', denylist=['dataset', 'training'])
def cifar10_no_augmentation_preprocess(dataset, training):
  del training

  def cast_image(features, targets):
    features['image'] = tf.cast(features['image'], tf.float32) / 255.0
    return features, targets

  dataset = dataset.map(cast_image)
  return dataset


def _cifar_augment_image(image):
  """Image augmentation suitable for CIFAR-10/100.

