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

  name_to_features = {
      # Image, annotation and choice ids.
      "img_id": tf.FixedLenFeature([], tf.int64),
      "annot_id": tf.FixedLenFeature([], tf.int64),

      # JPEG image.
      "image": tf.FixedLenFeature([], tf.string),

      # Text input.
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),

      # Label to predict for this example.
      "label": tf.FixedLenFeature([], tf.int64),
  }

  if FLAGS.has_choice_id:
    name_to_features["choice_id"] = tf.FixedLenFeature([], tf.int64)

  if FLAGS.use_bboxes:
    name_to_features.update({
        # Bounding boxes metadata.
        "bbox_pos": tf.FixedLenFeature([FLAGS.max_num_bboxes, 4], tf.int64),
        "bbox_idx": tf.FixedLenFeature([seq_length], tf.int64),
    })

  if FLAGS.mask_lm_loss:
    name_to_features.update({
        # Masked LM ids.
        "masked_lm_positions":
            tf.FixedLenFeature([FLAGS.max_predictions_per_seq], tf.int64),
        "masked_lm_ids":
            tf.FixedLenFeature([FLAGS.max_predictions_per_seq], tf.int64),
        "masked_lm_weights":
            tf.FixedLenFeature([FLAGS.max_predictions_per_seq], tf.float32),
    })

  def parse_bounding_box(height, width, image_float, box):
    """Decodes an image, resizes and scales RGB values to be in [0, 1]."""
    image_bbox = tf.image.crop_to_bounding_box(
        image_float,
        offset_height=box[0],
        offset_width=box[1],
        target_height=box[2],
        target_width=box[3])
    image_resized = tf.image.resize_image_with_pad(image_bbox, height, width)
    image = tf.reshape(image_resized, [height, width, 3])
    return image

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

    # Convert image to float tensor.
    image = example["image"]
    image_decoded = tf.image.decode_jpeg(image, channels=3)
    image_decoded.set_shape([None, None, 3])
    image_float = tf.image.convert_image_dtype(image_decoded, tf.float32)
    image_resized = tf.image.resize_image_with_pad(
        image_float, IMG_HEIGHT, IMG_WIDTH)
    example["image"] = tf.reshape(image_resized, [IMG_HEIGHT, IMG_WIDTH, 3])

    # Get bboxes.
    if FLAGS.use_bboxes:
      example["bbox_pos"] = tf.to_int32(example["bbox_pos"])
      bboxes = []
      for idx in range(FLAGS.max_num_bboxes):
        bboxes.append(
            parse_bounding_box(IMG_HEIGHT, IMG_WIDTH, image_float,
                               example["bbox_pos"][idx, :]))
      example["bboxes"] = tf.stack(bboxes)

      if FLAGS.use_bbox_position:
        # Resized bboxes.
        y, x, bbox_height, bbox_width = tf.unstack(example["bbox_pos"], axis=1)
        orig_height, orig_width = modeling.get_shape_list(image_float)[:2]
        example["bbox_pos"] = tf.cast(
            tf.stack([
                IMG_HEIGHT * y / orig_height,
                IMG_WIDTH * x / orig_width,
                IMG_HEIGHT * bbox_height / orig_height,
                IMG_WIDTH * bbox_width / orig_width
            ], 1),
            dtype=tf.int32)

    # 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

    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.Dataset.list_files(input_file, shuffle=True)
    d = d.apply(
        tf.data.experimental.parallel_interleave(
            tf.data.TFRecordDataset, cycle_length=20, sloppy=is_training
        )
    )
    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))

    d = d.prefetch(1)

    return d

  return input_fn
Creates an `input_fn` closure to be passed to TPUEstimator.
language/question_answering/b2t2/run_b2t2.py