code: def pretokenized_tfrecord_dataset(filenames,
                                    text2self,
                                    eos_included,
                                    repeat,
                                    batch_size,
                                    sequence_length):
    
    dataset = tf.data.TFRecordDataset(filenames, buffer_size=64 * 1024 * 1024)
    if repeat:
      dataset = dataset.repeat()
    keys = ["targets"] if text2self else ["inputs", "targets"]
    def decode_example(serialized_example):
      """Return a dict of Tensors from a serialized tensorflow.Example."""
      data_fields = {}
      data_items_to_decoders = {}
      for k in keys:
        data_fields[k] = tf.VarLenFeature(tf.int64)
        data_items_to_decoders[k] = tf.contrib.slim.tfexample_decoder.Tensor(k)
      decoder = tf.contrib.slim.tfexample_decoder.TFExampleDecoder(
          data_fields, data_items_to_decoders)
      decode_items = list(sorted(data_items_to_decoders))
      decoded = decoder.decode(serialized_example, items=decode_items)
      if not eos_included:
        decoded = [tf.concat([v, [1]], 0) for v in decoded]
      return dict(zip(decode_items, decoded))
    dataset = dataset.map(decode_example,
                          num_parallel_calls=tf.data.experimental.AUTOTUNE)
    return pack_and_batch(dataset, batch_size, sequence_length)
docs: Reads tensor2tensor-style data files.
  
    The dataset is defined by sets of TFRecord files of TFExample protos.
    There should be a "targets" feature (a 1d tensor of integers)
    If not text2self, there should also be an "inputs" feature.
    Other features get ignored.
  
    eos_included specifies whether the inputs and targets were written with an
    EOS token, as in tensor2tensor
  
    Args:
      filenames: a list of strings
      text2self: a boolean
      eos_included: a boolean
      repeat: a boolean
      batch_size: an integer
      sequence_length: an integer
    Returns:
      A tf.data.Dataset of batches