def load_data():
    """Load data.

    Load GLUE/MRPC dataset, and convert the dictionary format to
    (features, label), where `features` is a tuple of all input sentences.
    """
    feature_names = ("sentence1", "sentence2")

    def split_features(x):
        # GLUE comes with dictionary data, we convert it to a uniform format
        # (features, label), where features is a tuple consisting of all
        # features. This format is necessary for using KerasNLP preprocessors.
        features = tuple([x[name] for name in feature_names])
        label = x["label"]
        return (features, label)

    train_ds, test_ds, validation_ds = tfds.load(
        "glue/mrpc",
        split=["train", "test", "validation"],
    )

    train_ds = (
        train_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)
        .batch(FLAGS.batch_size)
        .prefetch(tf.data.AUTOTUNE)
    )
    test_ds = (
        test_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)
        .batch(FLAGS.batch_size)
        .prefetch(tf.data.AUTOTUNE)
    )
    validation_ds = (
        validation_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)
        .batch(FLAGS.batch_size)
        .prefetch(tf.data.AUTOTUNE)
    )
    return train_ds, test_ds, validation_ds
Load data.

    Load GLUE/MRPC dataset, and convert the dictionary format to
    (features, label), where `features` is a tuple of all input sentences.
    
benchmarks/model_benchmark/bert_benchmark.py