def input_fn(is_train, name, params):
  """An input function satisfying the tf.estimator API."""
  dataset = orqa_dataset.get_dataset(
      data_root=params["data_root"],
      name=name,
      split="train" if is_train else "dev")
  if is_train:
    dataset = dataset.repeat()

  def _extract_labels(d):
    return d, d.pop("answers")

  dataset = dataset.map(_extract_labels)
  dataset = dataset.prefetch(10)
  return dataset
An input function satisfying the tf.estimator API.
language/orqa/models/orqa_model.py