def input_function(is_train, embeddings):
  """An input function satisfying the tf.estimator API."""
  # A dataset with keys `question`, `context`, `answer_start`, and `answer_end`.
  dataset = nq_long_dataset.get_dataset(is_train)

  dataset = dataset.map(
      functools.partial(
          preprocess_mapper, lookup_table=embeddings.get_lookup_table()),
      num_parallel_calls=12)

  if is_train:
    dataset = dataset.repeat(FLAGS.num_epochs)
    dataset = dataset.shuffle(10000)

  dataset = dataset.padded_batch(
      batch_size=FLAGS.batch_size, padded_shapes=dataset.output_shapes)

  # The tf.estimator API expects the dataset to be a (features, labels) pair.
  dataset = dataset.map(lambda d: (d, d["long_answer_indices"]))

  dataset = dataset.prefetch(1)
  return dataset
An input function satisfying the tf.estimator API.
language/question_answering/decatt_docreader/models/nq_long_model.py