def input_fn_builder(input_file, seq_length, is_training, drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  vocab = tf.gfile.GFile(FLAGS.vocab_file).read().split("\n")
  cls_id = vocab.index("[CLS]")
  q_id = vocab.index("[Q]")

  name_to_features = {
      "unique_ids": tf.FixedLenFeature([], tf.int64),
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "start_positions": tf.FixedLenFeature([], tf.int64),
      "end_positions": tf.FixedLenFeature([], tf.int64),
      "answer_types": tf.FixedLenFeature([], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    if not is_training:
      # This zeroes out the question portion of the input and then adde back
      # the first two tokens.
      example["input_ids"] = example["input_ids"] * example["segment_ids"]
      example["input_ids"] = (
          example["input_ids"] +
          cls_id * tf.one_hot(0, seq_length, dtype=tf.int32) +
          q_id * tf.one_hot(1, seq_length, dtype=tf.int32))

    example["segment_ids"] = (
        example["segment_ids"] + tf.to_int32(
            tf.logical_and(
                tf.range(seq_length) >= example["start_positions"],
                tf.range(seq_length) <= example["end_positions"])))

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
Creates an `input_fn` closure to be passed to TPUEstimator.
language/capwap/synthetic/generate_questions.py