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

  try:
    input_file = tf.io.gfile.glob(input_file)
  except tf.errors.OpError:
    pass  # if it's not a sharded file just keep it as is

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "order": tf.FixedLenFeature([5], tf.int64),
  }

  def _decode_record(record, name_to_features, vocab_table):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)
    example["decoder_input"] = tf.identity(example["order"])[:4]
    example["decoder_input"] = tf.concat([[-1], example["decoder_input"]], 0)
    example["decoder_input"] = tf.reshape(example["decoder_input"], (5, 1))
    example["decoder_input"] = tf.to_float(example["decoder_input"])

    if add_masking:
      # Apply masking.
      mask_rate = 0.15
    else:
      mask_rate = 0.01
    max_predictions_per_seq = int(math.ceil(seq_length * mask_rate))
    cls_token = "[CLS]"
    sep_token = "[SEP]"
    mask_token = "[MASK]"
    pad_token = "[PAD]"

    # Don't allow masking of these tokens.
    mask_blacklist = tf.constant([cls_token, sep_token, pad_token])
    mask_blacklist_ids = vocab_table.lookup(mask_blacklist)
    mask_token_id = vocab_table.lookup(tf.constant(mask_token))

    mask_indices = ip.sample_mask_indices(example["input_ids"], mask_rate,
                                          mask_blacklist_ids,
                                          max_predictions_per_seq)

    token_ids_masked, target_token_ids = ip.apply_masking(
        example["input_ids"], mask_indices, mask_token_id, vocab_table.size())
    target_token_weights = tf.ones_like(target_token_ids, dtype=tf.float32)

    mask_indices = ip.pad_to_length(mask_indices, max_predictions_per_seq)
    target_token_ids = ip.pad_to_length(target_token_ids,
                                        max_predictions_per_seq)
    target_token_weights = ip.pad_to_length(target_token_weights,
                                            max_predictions_per_seq)

    # Get sep token indecies
    sep_token_id = vocab_table.lookup(tf.constant(_SEP_TOKEN))
    sep_positions = tf.where(tf.equal(example["input_ids"], sep_token_id))
    sep_positions = tf.reshape(sep_positions, [5, 1])
    example["sep_positions"] = sep_positions

    example["target_token_weights"] = target_token_weights
    example["target_token_ids"] = target_token_ids
    example["input_ids"] = token_ids_masked
    example["mask_indices"] = mask_indices

    # 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.
    vocab_table = contrib_lookup.index_table_from_file(FLAGS.vocab_file)

    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=256)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features, vocab_table
                                         ),
            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/reconstruct/run_paragraph_reconstruct.py