def get_dataset(is_train):
  """Gets a tf.data.Dataset representing the NQ data."""
  if is_train:
    data_pattern = FLAGS.nq_short_pipeline_train_pattern
  else:
    data_pattern = FLAGS.nq_short_pipeline_eval_pattern
  data_files = tf.gfile.Glob(data_pattern)
  assert data_files

  def _load_records(filenames):
    return tf.data.TFRecordDataset(filenames, buffer_size=16 * 1024 * 1024)

  if is_train:
    # During training, read from all files in parallel to improve the speed of
    # the input pipeline.
    dataset = tf.data.Dataset.from_tensor_slices(tf.constant(data_files))
    dataset = dataset.apply(
        tf.contrib.data.shuffle_and_repeat(buffer_size=len(data_files)))
    dataset = dataset.apply(
        tf.contrib.data.parallel_interleave(
            _load_records, sloppy=is_train, cycle_length=len(data_files)))
    dataset = dataset.shuffle(buffer_size=1000)
  else:
    dataset = _load_records(data_files)
  dataset = dataset.map(parse_example, num_parallel_calls=6)
  return dataset
Gets a tf.data.Dataset representing the NQ data.
language/question_answering/decatt_docreader/datasets/nq_short_pipeline_dataset.py