import tensorflow as tf

            "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),
        }

        # For training, we want a lot of parallel reading and shuffling.
        # For eval, we want no shuffling and parallel reading doesn't matter.
        if is_training:
            d = tf.data.Dataset.from_tensor_slices(tf.constant(input_files))
            d = d.repeat()
            d = d.shuffle(buffer_size=len(input_files))

            # `cycle_length` is the number of parallel files that get read.
            cycle_length = min(num_cpu_threads, len(input_files))

            # `sloppy` mode means that the interleaving is not exact. This adds
            # even more randomness to the training pipeline.
            d = d.apply(
                tf.contrib.data.parallel_interleave(
                    tf.data.TFRecordDataset,
                    sloppy=is_training,
                    cycle_length=cycle_length,
