import tensorflow as tf

    @classmethod
    def from_tfrecord_files(cls, input_files, **kwargs) -> tf.data.Dataset:
        dataset = utils.read_tfrecord_files(input_files, **kwargs)
        d = cls(examples=None, **kwargs)
        # parse example
        features = {
            d.input_ids: tf.io.VarLenFeature(tf.int64),
            d.token_type_ids: tf.io.VarLenFeature(tf.int64),
            d.attention_mask: tf.io.VarLenFeature(tf.int64),
            d.labels: tf.io.VarLenFeature(tf.int64),
        }
        dataset = dataset.map(
            lambda x: tf.io.parse_example(x, features),
            num_parallel_calls=utils.AUTOTUNE,
