def pretokenized_tfds_dataset(dataset_name=gin.REQUIRED,
                              text2self=gin.REQUIRED,
                              tfds_data_dir=gin.REQUIRED,
                              dataset_split=gin.REQUIRED,
                              batch_size=gin.REQUIRED,
                              sequence_length=gin.REQUIRED,
                              vocabulary=None):
  """Reads a tensorflow_datasets dataset.

  Args:
    dataset_name: a string
    text2self: a boolean
    tfds_data_dir: a boolean
    dataset_split: a string
    batch_size: an integer
    sequence_length: an integer
    vocabulary: ignored
  Returns:
    a tf.data.Dataset of batches
  """
  del vocabulary
  dataset = tfds.load(
      dataset_name,
      split=dataset_split,
      as_supervised=True,
      data_dir=tfds_data_dir)
  if dataset_split == "train":
    dataset = dataset.repeat()
    dataset = dataset.shuffle(1000)
  def shift_and_append_eos(t):
    # tfds encoder does not reserve an EOS token, so we need to shift
    # in order to do so.  We also append EOS=1.
    return tf.concat([t + 1, [1]], 0)
  def feature_map(inputs, targets):
    if text2self:
      return {"targets": shift_and_append_eos(targets)}
    else:
      return {"inputs": shift_and_append_eos(inputs),
              "targets": shift_and_append_eos(targets)}
  dataset = dataset.map(feature_map,
                        num_parallel_calls=tf.data.experimental.AUTOTUNE)
  return pack_and_batch(dataset, batch_size, sequence_length)
Reads a tensorflow_datasets dataset.

  Args:
    dataset_name: a string
    text2self: a boolean
    tfds_data_dir: a boolean
    dataset_split: a string
    batch_size: an integer
    sequence_length: an integer
    vocabulary: ignored
  Returns:
    a tf.data.Dataset of batches