def file_based_input_fn_builder(input_file, seq_length, fewshot_num_classes,
                                fewshot_num_examples_per_class, drop_remainder):
  """Creates an `input_fn` closure to be passed to tf.Estimator."""

  # Add one for the 'query' example.
  fewshot_batch = fewshot_num_classes * fewshot_num_examples_per_class + 1
  name_to_features = {
      "input_ids": tf.FixedLenSequenceFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenSequenceFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenSequenceFeature([seq_length], tf.int64),
      "guid": tf.FixedLenSequenceFeature([], tf.string),
  }

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

    # 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)
      shape = tf.shape(example[name])
      # sequence_examples come with dynamic/unknown dimension which we reshape
      # to explicit dimension for the fewshot "batch" size.
      example[name] = tf.reshape(t, tf.concat([[fewshot_batch], shape[1:]], 0))

    return example

  def input_fn(params):
    """The actual input function."""
    d = tf.data.TFRecordDataset(input_file)
    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=params["batch_size"],
            drop_remainder=drop_remainder))

    return d

  return input_fn
Creates an `input_fn` closure to be passed to tf.Estimator.
language/relation_learning/models/bert_fewshot_classifier.py