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_short_pipeline_dataset.get_dataset(is_train)

  dataset = dataset.map(
      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)

    # Compute the batch size of each bucket such that the total number of
    # context words in each batch is bounded by FLAGS.total_batch_size during
    # training.
    bucket_boundaries = list(range(128, FLAGS.max_context_len, 128))
    bucket_batch_sizes = [
        int(FLAGS.total_batch_size / b) for b in bucket_boundaries
    ]
    bucket_batch_sizes.append(
        int(FLAGS.total_batch_size / FLAGS.max_context_len))
    dataset = dataset.apply(
        tf.data.experimental.bucket_by_sequence_length(
            element_length_func=lambda d: d["context_len"],
            bucket_boundaries=bucket_boundaries,
            bucket_batch_sizes=bucket_batch_sizes,
            padded_shapes=dataset.output_shapes))
  else:
    # Don't do any bucketing during evals to avoid biased sampling.
    dataset = dataset.padded_batch(
        batch_size=int(FLAGS.total_batch_size / FLAGS.max_context_len),
        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["answer_start"], d["answer_end"])))
  dataset = dataset.prefetch(10)
  return dataset
An input function satisfying the tf.estimator API.
language/question_answering/decatt_docreader/models/nq_short_pipeline_model.py