import tensorflow as tf

                [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
            ),
