def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder, num_choices):
  """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 = {}
  for i in range(50):
    name_to_features["input_ids" + str(i)] = tf.FixedLenFeature([seq_length],
                                                                tf.int64)
    name_to_features["input_mask" + str(i)] = tf.FixedLenFeature([seq_length],
                                                                 tf.int64)
    name_to_features["segment_ids" + str(i)] = tf.FixedLenFeature([seq_length],
                                                                  tf.int64)
  name_to_features["label_types"] = tf.FixedLenFeature([4], 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)

    # total of 32 examples
    # 4 labels / shuffled
    # random samples from 23 + 23 last for distractors
    num_sampled = 14
    same_doc_idxs = tf.random.shuffle(tf.range(4, 27))[:num_sampled]
    rand_doc_idxs = tf.random.shuffle(tf.range(27, 50))[:num_sampled]
    batch_indexes = tf.concat([tf.range(4), same_doc_idxs, rand_doc_idxs],
                              axis=0)
    batch_indexes = tf.random.shuffle(batch_indexes)
    # At this point, we have shuffled the indexes and sampled them such that
    # we still have the index of 4 targets, 14 sampled from the same doc
    # and 14 sampled from different docs. But these are just indexes.

    # Here we need to grab the inputs according to the indexes above
    # We stack all the inputs so we can gather on the matrix
    input_id_stack, input_mask_stack, segment_id_stack = [], [], []
    for i in range(50):
      input_id_stack.append(example["input_ids" + str(i)])
      input_mask_stack.append(example["input_mask" + str(i)])
      segment_id_stack.append(example["segment_ids" + str(i)])
    input_id_stack = tf.stack(input_id_stack)
    input_mask_stack = tf.stack(input_mask_stack)
    segment_id_stack = tf.stack(segment_id_stack)

    input_ids = tf.gather(input_id_stack, batch_indexes)
    input_masks = tf.gather(input_mask_stack, batch_indexes)
    segment_ids = tf.gather(segment_id_stack, batch_indexes)

    # Note that we override the name of the input (eg. input_ids5)
    # So we replace the input with the shuffled and sampled input
    # We only set num_choices of them since those will be used.
    for i in range(num_choices):
      example["input_ids" + str(i)] = input_ids[i]
      example["input_mask" + str(i)] = input_masks[i]
      example["segment_ids" + str(i)] = segment_ids[i]

    # Note that for inputs num_choices-50 will not be used so we must purge them
    for i in range(num_choices, 50):
      del example["input_ids" + str(i)]
      del example["input_mask" + str(i)]
      del example["segment_ids" + str(i)]

    label_idx = []
    for i in range(4):
      label_idx.append(tf.where(tf.equal(batch_indexes, tf.constant(i)))[0])
    label_idx = tf.reshape(tf.concat(label_idx, axis=0), [-1])
    label_idx = tf.scatter_nd(
        tf.reshape(example["label_types"], [4, 1]), label_idx, [8])
    example["labels"] = label_idx

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

    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]
      wiki_pct = 0.02222
      dset_weights = [wiki_pct, 1 - wiki_pct]
      d = tf.data.experimental.sample_from_datasets(dataset_list, dset_weights)
      # 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)

    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/cpc/run_cpc.py