import tensorflow as tf

    input_files, max_seq_length, max_predictions_per_seq, is_training, num_cpu_threads=4
):
    """Creates an `input_fn` closure to be passed to TPUEstimator."""

    def input_fn(params):
        """The actual input function."""
        batch_size = params["batch_size"]

        name_to_features = {
            "input_ids": tf.FixedLenFeature([max_seq_length], tf.int64),
            "input_mask": tf.FixedLenFeature([max_seq_length], tf.int64),
            "segment_ids": tf.FixedLenFeature([max_seq_length], tf.int64),
            "masked_lm_positions": tf.FixedLenFeature(
                [max_predictions_per_seq], tf.int64
            ),
            "masked_lm_ids": tf.FixedLenFeature([max_predictions_per_seq], tf.int64),
            "masked_lm_weights": tf.FixedLenFeature(
                [max_predictions_per_seq], tf.float32
            ),
            "next_sentence_labels": tf.FixedLenFeature([1], tf.int64),
