import tensorflow as tf

@gin.configurable(module='trax.data', denylist=['dataset', 'training'])
def c4_preprocess(dataset,
                  training,
                  max_target_length=-1,
                  tokenization=None,
                  spm_path=None):
  """Pre-processing function for C4 dataset."""
  del training

  def unicode_decode_chars(features, targets):
    targets = tf.strings.unicode_decode(features['text'], 'UTF-8')
    targets = tf.cast(targets, tf.int64)
    features['targets'] = targets
    features['inputs'] = targets
    return (features, targets)

  def spc_tokenize(tokenizer, features, targets):
    del targets
    tokenized_text = tokenizer.tokenize(features['text'])
    features['targets'] = tf.cast(tokenized_text, tf.int64)
