def file_based_input_fn_builder(input_file, is_training, drop_remainder,
                                add_masking):
  """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 = {
      "sents":
          tf.FixedLenFeature([FLAGS.max_para_length * FLAGS.max_sent_length],
                             tf.int64)
  }

  def _decode_record(record, name_to_features, vocab_table):
    """Decodes a record to a TensorFlow example."""
    target_example = tf.parse_single_example(record[0], name_to_features)
    target_example = tf.reshape(target_example["sents"],
                                [FLAGS.max_para_length, FLAGS.max_sent_length])
    # distractor_examples = []
    # for rec in record[1:]:
    #   distractor_examples.append(
    #       tf.reshape(
    #           tf.parse_single_example(rec, name_to_features)["sents"],
    #           [FLAGS.max_para_length, FLAGS.max_sent_length]))
    # This is an unfortunate hack but is necessary to get around a TF error.
    dist0 = tf.reshape(
        tf.parse_single_example(record[1], name_to_features)["sents"],
        [FLAGS.max_para_length, FLAGS.max_sent_length])
    dist1 = tf.reshape(
        tf.parse_single_example(record[2], name_to_features)["sents"],
        [FLAGS.max_para_length, FLAGS.max_sent_length])
    dist2 = tf.reshape(
        tf.parse_single_example(record[3], name_to_features)["sents"],
        [FLAGS.max_para_length, FLAGS.max_sent_length])
    dist3 = tf.reshape(
        tf.parse_single_example(record[4], name_to_features)["sents"],
        [FLAGS.max_para_length, FLAGS.max_sent_length])

    inputs_obj = build_bert_inputs(target_example)

    distractor_obj = build_distractors([dist0, dist1, dist2, dist3],
                                       inputs_obj.context)

    example = {}
    example["input_ids"] = tf.concat(
        [inputs_obj.input_ids, distractor_obj.input_ids], axis=0)
    example["input_mask"] = tf.concat(
        [inputs_obj.input_mask, distractor_obj.input_mask], axis=0)
    example["segment_ids"] = tf.concat(
        [inputs_obj.segment_ids, distractor_obj.segment_ids], axis=0)
    example["label_types"] = inputs_obj.label_types

    # Add masking:
    if add_masking:
      mask_rate = FLAGS.mask_rate
      max_predictions_per_seq = int(math.ceil(FLAGS.max_seq_length * mask_rate))
      cls_token = "[CLS]"
      sep_token = "[SEP]"
      mask_token = "[MASK]"
      # pad_token = "[PAD]"
      mask_blacklist = tf.constant([cls_token, sep_token])  # , pad_token])
      mask_blacklist_ids = tf.to_int32(vocab_table.lookup(mask_blacklist))
      mask_token_id = tf.to_int32(vocab_table.lookup(tf.constant(mask_token)))
      input_ids = tf.to_int32(example["input_ids"])

      def call_sample_mask_indices(x):
        return ip.sample_mask_indices(x, mask_rate, mask_blacklist_ids,
                                      max_predictions_per_seq)

      mask_indices = tf.map_fn(
          call_sample_mask_indices, input_ids, dtype=tf.int32)

      def call_get_target_tokens(x):
        input_len = tf.shape(input_ids)[-1]
        x_input_id = x[:input_len]
        x_mask_indices = x[input_len:]
        return ip.get_target_tokens_for_apply(x_input_id, x_mask_indices)

      map_input = tf.concat([input_ids, mask_indices], -1)
      target_token_ids = tf.map_fn(call_get_target_tokens, map_input)

      def call_apply_masking(x):
        input_len = tf.shape(input_ids)[-1]
        mask_idx_len = tf.shape(mask_indices)[-1]
        x_input_id = x[:input_len]
        x_mask_indices = x[input_len:input_len + mask_idx_len]
        x_target_token_ids = x[input_len + mask_idx_len:]
        return ip.apply_masking(x_input_id, x_target_token_ids, x_mask_indices,
                                mask_token_id, 1000)

      map_input2 = tf.concat([input_ids, mask_indices, target_token_ids], -1)
      token_ids_masked = tf.map_fn(call_apply_masking, tf.to_int64(map_input2))
      target_token_weights = tf.ones_like(target_token_ids, dtype=tf.float32)
      pad_targets = tf.where(
          tf.equal(target_token_ids, 0),
          tf.ones_like(target_token_ids, dtype=tf.float32),
          tf.zeros_like(target_token_ids, dtype=tf.float32))
      target_token_weights = target_token_weights - pad_targets
      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

      # Set shape explicitly for TPU
      example["target_token_weights"].set_shape(
          [FLAGS.num_choices, max_predictions_per_seq])
      example["target_token_ids"].set_shape(
          [FLAGS.num_choices, max_predictions_per_seq])
      example["mask_indices"].set_shape(
          [FLAGS.num_choices, max_predictions_per_seq])

    # Set shape explicitly for TPU
    example["input_ids"].set_shape([FLAGS.num_choices, FLAGS.max_seq_length])
    example["input_mask"].set_shape([FLAGS.num_choices, FLAGS.max_seq_length])
    example["segment_ids"].set_shape([FLAGS.num_choices, FLAGS.max_seq_length])
    example["label_types"].set_shape([4])

    example["label_ids"] = tf.scatter_nd(
        tf.reshape(example["label_types"], [4, 1]), tf.range(4), [8])

    # 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"]

    vocab_table = contrib_lookup.index_table_from_file(FLAGS.vocab_file)

    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:
      dataset_list = [
          tf.data.TFRecordDataset(expanded_files[i])
          for i in range(len(expanded_files))
      ]
      if is_training:
        dataset_list = [d.repeat() for d in dataset_list]
      dset_weights = [FLAGS.dataset_one_weight, 1 - FLAGS.dataset_one_weight]
      d = tf.data.experimental.sample_from_datasets(dataset_list, dset_weights)

      # Note that sample_from_datasets() inserts randomness into the training
      # An alternative would be to use choose_from_datasets() but then the
      # order must be stated explicitly which is less intitive for unbalanced
      # datasets. Example below:
      #
      # choice_dataset = tf.data.Dataset.range(len(dataset_list)).repeat()
      # d = tf.data.experimental.choose_from_datasets(dataset_list,
      #                                               choice_dataset)

      if is_training:
        d = d.shuffle(buffer_size=256)

    # The window size will be for selecting negative samples
    # It equals the number of documents to sample from -1
    d = d.window(size=FLAGS.data_window_size, shift=FLAGS.data_window_shift)
    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/cpc/run_cc_cpc.py