import tensorflow as tf

            token_type_ids.append(e.token_type_ids)
            attention_mask.append(e.attention_mask)
            labels.append(e.label_ids)

        # parse examples to dataset
        def _to_dataset(x, dtype=tf.int32):
            x = tf.ragged.constant(x, dtype=dtype)
            d = tf.data.Dataset.from_tensor_slices(x)
            d = d.map(lambda x: x)
            return d

        dataset = tf.data.Dataset.zip(
