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

  input_file = input_file.split(",")

  expanded_files = []
  for infile in input_file:
    try:
      sharded_files = tf.io.gfile.glob(infile)
      expanded_files.append(sharded_files)
    except tf.errors.OpError:
      expanded_files.append(infile)

  name_to_features = {
      "context_input_ids": tf.FixedLenFeature([FLAGS.max_seq_length], tf.int64)
  }

  cand_size = 25
  max_target_len = 100
  for i in range(cand_size):
    name_to_features["comp_input_ids" + str(i)] = tf.FixedLenFeature(
        [max_target_len], tf.int64)

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

    example_obj = build_record_example(raw_example, cand_size)

    example = {}
    example["input_ids"] = example_obj.input_ids
    example["input_mask"] = example_obj.input_mask
    example["segment_ids"] = example_obj.segment_ids
    example["label"] = example_obj.label

    example["input_ids"].set_shape([cand_size + 1, FLAGS.max_seq_length])
    example["input_mask"].set_shape([cand_size + 1, FLAGS.max_seq_length])
    example["segment_ids"].set_shape([cand_size + 1, FLAGS.max_seq_length])
    example["label"].set_shape([1])

    # 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()):  # pylint: disable=g-builtin-op
      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"]

    if len(expanded_files) == 1:
      d = tf.data.TFRecordDataset(expanded_files[0])
      if is_training:
        d = d.repeat()
        d = d.shuffle(buffer_size=256)
    else:
      assert False, "Only one data file allowed"

    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/conpono/evals/run_record.py