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

def pretokenized_tfrecord_dataset(filenames,
                                  text2self,
                                  eos_included,
                                  repeat,
                                  batch_size,
                                  sequence_length):
  """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
  """
  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)