Repository: keras-team/keras-nlp

Issue:
Add presets for GPT-2
This is actually underway in https://github.com/keras-team/keras-nlp/pull/366, but will need quite a rebase for the latest changes we have made on BERT.

 * Contribute `gpt2_presets.py` and `gpt2_presets_test.py` files modeled on BERT.
 * Add a `from_preset()` method to `keras_nlp.models.GPT2` with appropriate documentation.
 * Mirror the additional preset testing (without weights downloads) on `bert_models_test.py` to `gpt_models_test.py`.

This can come before or after https://github.com/keras-team/keras-nlp/issues/462 (adding a tokenizer). But if we do this second, we should include a `from_preset()` method for the tokenizer.

Discussion:
abheesht17: > Thanks for the great work, but can you confirm why the tests are failing?

```
/home/runner/work/_temp/93cb0ed8-bd8c-40ef-a3e2-3e39de33365b.sh: line 1:  1778 Killed                  pytest --cov=keras_nlp --cov-report xml:coverage.xml
keras_nlp/models/gpt2/gpt2_models_test.py::GPT2Test::test_valid_call_presets 
Error: Process completed with exit code 137.
```

I think it is unable to test all four presets because of the size of the models. I've tested all 4 of them on Colab.

Patch:
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
index e2b843e5..2242aa44 100644
--- a/.github/workflows/nightly.yml
+++ b/.github/workflows/nightly.yml
@@ -31,7 +31,9 @@ jobs:
       - name: Install dependencies
         run: |
           pip install -r requirements-nightly.txt --progress-bar off
-          pip install -e "." --progress-bar off
+          # --no-deps so pip will not try to install tf and tf-text stable
+          # versions after installing the nightly versions above.
+          pip install --no-deps -e "." --progress-bar off
       - name: Test with pytest
         run: |
           pytest --cov=keras_nlp --cov-report xml:coverage.xml
diff --git a/examples/glue_benchmark/README.md b/examples/glue_benchmark/README.md
new file mode 100644
index 00000000..a8d14998
--- /dev/null
+++ b/examples/glue_benchmark/README.md
@@ -0,0 +1,98 @@
+# GLUE Finetuning Script
+
+This script is written to help you evaluate your model on GLUE benchmarking.
+It provides the functionalities below:
+
+- Load and preprocess GLUE data.
+- Finetuning your Keras text classification model. 
+- Generate GLUE submission files.
+
+To use the script, you need to change the code to load your pretrained model,
+and run the command below:
+
+```shell
+python glue.py --task_name="mrpc" --batch_size=32 \
+    --submission_file_path="glue_submissions/"
+```
+
+By default the script finetunes on the tiniest BERT model we have available 
+(this will be fast but not top performing).
+
+To make a real GLUE leaderboard submission, you need to call the finetuning on 
+all tasks, then enter the submission directory then zip the submission files:
+```shell
+for task in cola sst2 mrpc rte mnli_matched mnli_mismatched ax qnli qqp; do
+  python glue.py --task_name="$task" --submission_directory="glue_submissions/"
+done
+cd glue_submissions
+zip -r submission.zip *.tsv
+```
+
+GLUE submission requires the `submission.zip` contains `.tsv` file for all 
+tasks, otherwise it will be a failed submission. An empty `.tsv` will also fail 
+because it checks the content. If you only want to evaluate on certain tasks, 
+you can download the sample submission, and put the `.tsv` files for tasks you 
+don't run inside your submission file. For example if you don't want to 
+run the `ax` task, then you can do:
+
+```
+curl -O https://gluebenchmark.com/assets/CBOW.zip
+unzip CBOW.zip -d sample_submissions
+cp sample_submissions/AX.tsv glue_submissions
+```
+
+## How to Use the Script
+
+To use this script on your model, you need to do 3 things:
+
+1. Implement your custom preprocessing in `preprocess_fn()`.
+2. Load your pretrained model.
+3. Make the finetune model with your model.
+
+Code needing customization is wrapped between comment
+`Custom code block starts` and 
+`Custom code block ends`. See instructions on each step below.
+
+### Custom Preprocessing
+
+In all GLUE dataset, each record comes with features of one or two sentences, 
+and one label. In the script, we load GLUE dataset in the format 
+`(features, labels)`,  where `features` is a tuple of either 1 sentence or 2
+sentences. Your need to write custom preprocessing logic to convert to data
+to the required input of your model. For example, in the current script 
+(finetuning for KerasNLP BERT), it is doing:
+
+```python
+bert_preprocessor = keras_nlp.models.BertPreprocessor.from_preset(
+    "bert_tiny_uncased_en"
+)
+def preprocess_fn(feature, label):
+    return bert_preprocessor(feature), label
+```
+It uses the `BertPreprocessor` to convert input feature formats.
+
+### Load Pretrained Model
+
+As long as it is a Keras model, you can use it with this script. 
+
+### Make the Finetuning Model
+
+Users need to make a classification model based on your pretrained model for 
+evaluation purposes. For example, [`BertClassifier`](https://github.com/keras-team/keras-nlp/blob/24c969b7f6eb5074b67d513b504c91122b0b3185/keras_nlp/models/bert/bert_tasks.py#L28) takes a `Bert` model as backbone, 
+and adds a dense layer on top of it. Please pay attention that different model 
+could use different classifier structure, e.g., in [RoBERTa](https://github.com/huggingface/transformers/blob/94b3f544a1f5e04b78d87a2ae32a7ac252e22e31/src/transformers/models/roberta/modeling_roberta.py#L1437-L1456), 
+it has 2 dense layers. If you are using pretrained model from an OSS package, 
+please find the correct classifier. If you use a custom model, you can start 
+experimenting with a simple dense layer, and adjust the structure based on 
+its performance.
+
+## Flags Table
+
+| Flags Name                 	| Explanation                                     	| Default 	|
+|----------------------------	|-------------------------------------------------	|---------	|
+| task_name                  	| The name of the GLUE task to finetune on.       	| "mrpc"  	|
+| batch_size                 	| Data batch size                                 	| 32      	|
+| epochs                     	| Number of epochs to run finetuning.             	| 2       	|
+| learning_rate              	| The optimizer's learning rate                   	| 5e-5    	|
+| submission_directory       	| The file path to save the glue submission file. 	| None    	|
+| finetuning_model_save_path 	| The path to save the finetuning model.          	| None    	|
diff --git a/examples/glue_benchmark/glue.py b/examples/glue_benchmark/glue.py
new file mode 100644
index 00000000..d909aa7c
--- /dev/null
+++ b/examples/glue_benchmark/glue.py
@@ -0,0 +1,236 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import csv
+import os
+
+import numpy as np
+import tensorflow as tf
+import tensorflow_datasets as tfds
+from absl import app
+from absl import flags
+from tensorflow import keras
+
+import keras_nlp
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string(
+    "task_name",
+    "mrpc",
+    "The name of the GLUE task to finetune on.",
+)
+
+flags.DEFINE_integer(
+    "batch_size",
+    32,
+    "Batch size of data.",
+)
+
+flags.DEFINE_integer(
+    "epochs",
+    2,
+    "Number of epochs to run finetuning.",
+)
+
+flags.DEFINE_float(
+    "learning_rate",
+    5e-5,
+    "Learning rate",
+)
+
+flags.DEFINE_string(
+    "submission_directory",
+    None,
+    "The directory to save the glue submission file.",
+)
+
+
+flags.DEFINE_string(
+    "finetuning_model_save_path",
+    None,
+    "The path to save the finetuning model. If None, the model is not saved.",
+)
+
+
+def load_data(task_name):
+    """
+    Load GLUE dataset.
+
+    Load GLUE dataset, and convert the dictionary format to (features, label),
+    where features is a tuple of all input sentences.
+    """
+    if task_name in ("cola", "sst2"):
+        feature_names = ("sentence",)
+    elif task_name in ("mrpc", "stsb", "rte", "wnli"):
+        feature_names = ("sentence1", "sentence2")
+    elif task_name in ("mnli", "mnli_matched", "mnli_mismatched", "ax"):
+        feature_names = ("premise", "hypothesis")
+    elif task_name in "qnli":
+        feature_names = ("question", "sentence")
+    elif task_name in "qqp":
+        feature_names = ("question1", "question2")
+    else:
+        raise ValueError(f"Unknown task_name {task_name}.")
+
+    test_suffix = ""
+    if task_name in ("mnli", "mnli_matched"):
+        # For "mnli", just run default to "mnli_matched".
+        task_name = "mnli"
+        test_suffix = "_matched"
+    elif task_name in ("mnli_mismatched"):
+        task_name = "mnli"
+        test_suffix = "_mismatched"
+
+    def split_features(x):
+        # GLUE comes with dictonary data, we convert it to a uniform format
+        # (features, label), where features is a tuple consisting of all
+        # features.
+        features = tuple([x[name] for name in feature_names])
+        label = x["label"]
+        return (features, label)
+
+    if task_name == "ax":
+        # AX is trained and evaluated on MNLI, and has its own test split.
+        train_ds, validation_ds = tfds.load(
+            "glue/mnli",
+            split=["train", "validation_matched"],
+        )
+        test_ds = tfds.load(
+            "glue/ax",
+            split="test",
+        )
+    else:
+        train_ds, test_ds, validation_ds = tfds.load(
+            f"glue/{task_name}",
+            split=["train", "test" + test_suffix, "validation" + test_suffix],
+        )
+    train_ds = train_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)
+    test_ds = test_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)
+    validation_ds = validation_ds.map(
+        split_features, num_parallel_calls=tf.data.AUTOTUNE
+    )
+    return train_ds, test_ds, validation_ds
+
+
+def preprocess_data(preprocess_fn, dataset):
+    """Run `proprocess_fn` on input dataset then batch & prefetch."""
+    return (
+        dataset.map(preprocess_fn)
+        .batch(FLAGS.batch_size)
+        .prefetch(tf.data.AUTOTUNE)
+    )
+
+
+def generate_submission_files(finetuning_model, test_ds):
+    """Generate GLUE leaderboard submission files."""
+    filenames = {
+        "cola": "CoLA.tsv",
+        "sst2": "SST-2.tsv",
+        "mrpc": "MRPC.tsv",
+        "qqp": "QQP.tsv",
+        "stsb": "STS-B.tsv",
+        "mnli_matched": "MNLI-m.tsv",
+        "mnli_mismatched": "MNLI-mm.tsv",
+        "qnli": "QNLI.tsv",
+        "rte": "RTE.tsv",
+        "wnli": "WNLI.tsv",
+        "ax": "AX.tsv",
+    }
+
+    labelnames = {
+        "mnli_matched": ["entailment", "neutral", "contradiction"],
+        "mnli_mismatched": ["entailment", "neutral", "contradiction"],
+        "ax": ["entailment", "neutral", "contradiction"],
+        "qnli": ["entailment", "not_entailment"],
+        "rte": ["entailment", "not_entailment"],
+    }
+    if not os.path.exists(FLAGS.submission_directory):
+        os.makedirs(FLAGS.submission_directory)
+    filename = FLAGS.submission_directory + "/" + filenames[FLAGS.task_name]
+
+    labelname = labelnames.get(FLAGS.task_name)
+    with tf.io.gfile.GFile(filename, "w") as f:
+        # GLUE requires a format of index + tab + prediction.
+        writer = csv.writer(f, delimiter="\t")
+        # Write the required headline for GLUE.
+        writer.writerow(["index", "prediction"])
+        predictions = finetuning_model.predict(test_ds)
+        predictions = np.argmax(predictions, -1)
+        for idx, pred in enumerate(predictions):
+            if labelname:
+                pred_value = labelname[int(pred)]
+            else:
+                pred_value = pred
+            writer.writerow([idx, pred_value])
+
+
+def main(_):
+    train_ds, test_ds, val_ds = load_data(FLAGS.task_name)
+
+    # ----- Custom code block starts -----
+    bert_model = keras_nlp.models.Bert.from_preset("bert_tiny_uncased_en")
+    bert_preprocessor = keras_nlp.models.BertPreprocessor.from_preset(
+        "bert_tiny_uncased_en"
+    )
+
+    # Users should change this function to implement the preprocessing required
+    # by the model.
+    def preprocess_fn(feature, label):
+        return bert_preprocessor(feature), label
+
+    # ----- Custom code block ends -----
+
+    train_ds = preprocess_data(preprocess_fn, train_ds)
+    val_ds = preprocess_data(preprocess_fn, val_ds)
+    test_ds = preprocess_data(preprocess_fn, test_ds)
+
+    loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+    metrics = [keras.metrics.SparseCategoricalAccuracy()]
+    if FLAGS.task_name == "stsb":
+        num_classes = 1
+        loss = keras.losses.MeanSquaredError()
+        metrics = [keras.losses.MeanSquaredError()]
+    elif FLAGS.task_name in ("mnli", "mnli_mismatched", "mnli_matched", "ax"):
+        num_classes = 3
+    else:
+        num_classes = 2
+
+    # ----- Custom code block starts -----
+    # Users should change this `BertClassifier` to your own classifier.
+    # Commonly the classifier is simply your model + several dense layers,
+    # please refer to "Make the Finetuning Model" section in README for
+    # detailed instructions.
+    finetuning_model = keras_nlp.models.BertClassifier(
+        backbone=bert_model,
+        num_classes=num_classes,
+    )
+    # ----- Custom code block ends -----
+
+    finetuning_model.compile(
+        optimizer=tf.keras.optimizers.Adam(FLAGS.learning_rate),
+        loss=loss,
+        metrics=metrics,
+    )
+
+    finetuning_model.fit(train_ds, validation_data=val_ds, epochs=FLAGS.epochs)
+
+    if FLAGS.submission_directory:
+        generate_submission_files(finetuning_model, test_ds)
+
+    if FLAGS.finetuning_model_save_path:
+        finetuning_model.save(FLAGS.finetuning_model_save_path)
+
+
+if __name__ == "__main__":
+    app.run(main)
diff --git a/keras_nlp/__init__.py b/keras_nlp/__init__.py
index 67ab65ba..73e7e2c5 100644
--- a/keras_nlp/__init__.py
+++ b/keras_nlp/__init__.py
@@ -18,4 +18,4 @@ from keras_nlp import models
 from keras_nlp import tokenizers
 from keras_nlp import utils
 
-__version__ = "0.3.0"
+__version__ = "0.4.0"
diff --git a/keras_nlp/models/bert/bert_preprocessing.py b/keras_nlp/models/bert/bert_preprocessing.py
index 126563f2..fcfdc67f 100644
--- a/keras_nlp/models/bert/bert_preprocessing.py
+++ b/keras_nlp/models/bert/bert_preprocessing.py
@@ -21,6 +21,7 @@ from tensorflow import keras
 from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
 from keras_nlp.models.bert.bert_presets import backbone_presets
 from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
+from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 
@@ -167,15 +168,27 @@ class BertPreprocessor(keras.layers.Layer):
 
     This preprocessing layer will do three things:
 
-     - Tokenize any number of inputs using the `tokenizer`.
+     - Tokenize any number of input segments using the `tokenizer`.
      - Pack the inputs together using a `keras_nlp.layers.MultiSegmentPacker`.
        with the appropriate `"[CLS]"`, `"[SEP]"` and `"[PAD]"` tokens.
      - Construct a dictionary of with keys `"token_ids"`, `"segment_ids"`,
        `"padding_mask"`, that can be passed directly to a BERT model.
 
-    This layer will accept either a tuple of (possibly batched) inputs, or a
-    single input tensor. If a single tensor is passed, it will be packed
-    equivalently to a tuple with a single element.
+    This layer can be used directly with `tf.data.Dataset.map` to preprocess
+    string data in the `(x, y, sample_weight)` format used by
+    `keras.Model.fit`.
+
+    The call method of this layer accepts three arguments, `x`, `y`, and
+    `sample_weights`. `x` should be either a (possibly batched) string tensor,
+    or a tuple of (possibly batched) string tensors. `y` and `sample_weights`
+    are both optional, can have any format, and will be passed through
+    unaltered.
+
+    Special care should be taken when using `tf.data` to map over an unlabeled
+    tuple of string segments. `tf.data` will unpack this tuple directly into the
+    call arguments of this layer, rather than forward all argument to `x`. To
+    handle this case, it is recommended to  explicitly call the layer, e.g.
+    `ds.map(lambda seg1, seg2: preprocessor(x=(seg1, seg2)))`.
 
     Args:
         tokenizer: A `keras_nlp.models.BertTokenizer` instance.
@@ -211,12 +224,9 @@ class BertPreprocessor(keras.layers.Layer):
     features = ["The quick brown fox jumped.", "I forgot my homework."]
     labels = [0, 1]
     ds = tf.data.Dataset.from_tensor_slices((features, labels))
-    ds = ds.map(
-        lambda x, y: (preprocessor(x), y),
-        num_parallel_calls=tf.data.AUTOTUNE,
-    )
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
 
-    # Map a dataset to preprocess a multiple sentences.
+    # Map a dataset to preprocess multiple sentence pairs.
     first_sentences = ["The quick brown fox jumped.", "Call me Ishmael."]
     second_sentences = ["The fox tripped.", "Oh look, a whale."]
     labels = [1, 1]
@@ -225,10 +235,19 @@ class BertPreprocessor(keras.layers.Layer):
             (first_sentences, second_sentences), labels
         )
     )
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Map a dataset to preprocess unlabeled sentence pairs.
+    first_sentences = ["The quick brown fox jumped.", "Call me Ishmael."]
+    second_sentences = ["The fox tripped.", "Oh look, a whale."]
+    ds = tf.data.Dataset.from_tensor_slices((first_sentences, second_sentences))
+    # Watch out for tf.data's default unpacking of tuples here!
+    # Best to invoke the `preprocessor` directly in this case.
     ds = ds.map(
-        lambda x, y: (preprocessor(x), y),
+        lambda s1, s2: preprocessor(x=(s1, s2)),
         num_parallel_calls=tf.data.AUTOTUNE,
-    )    ```
+    )
+    ```
     """
 
     def __init__(
@@ -270,17 +289,18 @@ class BertPreprocessor(keras.layers.Layer):
             config["tokenizer"] = keras.layers.deserialize(config["tokenizer"])
         return cls(**config)
 
-    def call(self, inputs):
-        if not isinstance(inputs, (list, tuple)):
-            inputs = [inputs]
+    def call(self, x, y=None, sample_weight=None):
+        if not isinstance(x, (list, tuple)):
+            x = [x]
 
-        inputs = [self.tokenizer(x) for x in inputs]
-        token_ids, segment_ids = self.packer(inputs)
-        return {
+        x = [self.tokenizer(segment) for segment in x]
+        token_ids, segment_ids = self.packer(x)
+        x = {
             "token_ids": token_ids,
             "segment_ids": segment_ids,
             "padding_mask": token_ids != self.tokenizer.pad_token_id,
         }
+        return pack_x_y_sample_weight(x, y, sample_weight)
 
     @classproperty
     def presets(cls):
diff --git a/keras_nlp/models/bert/bert_presets.py b/keras_nlp/models/bert/bert_presets.py
index 728444ca..9a0487df 100644
--- a/keras_nlp/models/bert/bert_presets.py
+++ b/keras_nlp/models/bert/bert_presets.py
@@ -220,3 +220,34 @@ backbone_presets = {
         "vocabulary_hash": "bb6ca9b42e790e5cd986bbb16444d0e0",
     },
 }
+
+classifier_presets = {
+    "bert_tiny_uncased_en_sst2": {
+        "config": {
+            "backbone": {
+                "class_name": "keras_nlp>Bert",
+                "config": {
+                    "vocabulary_size": 30522,
+                    "hidden_dim": 128,
+                    "intermediate_dim": 512,
+                    "num_layers": 2,
+                    "num_heads": 2,
+                    "max_sequence_length": 512,
+                    "num_segments": 2,
+                    "dropout": 0.1,
+                },
+            },
+            "num_classes": 2,
+        },
+        "preprocessor_config": {
+            "lowercase": True,
+        },
+        "description": (
+            "bert_tiny_uncased_en backbone fine-tuned on the glue/sst2 dataset."
+        ),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_tiny_uncased_en_sst2/model.h5",
+        "weights_hash": "1f9c2d59f9e229e08f3fbd44239cfb0b",
+        "vocabulary_url": "https://storage.googleapis.com/keras-nlp/models/bert_tiny_uncased_en_sst2/vocab.txt",
+        "vocabulary_hash": "64800d5d8528ce344256daf115d4965e",
+    }
+}
diff --git a/keras_nlp/models/bert/bert_tasks.py b/keras_nlp/models/bert/bert_tasks.py
index f89d8c29..a01f5cc1 100644
--- a/keras_nlp/models/bert/bert_tasks.py
+++ b/keras_nlp/models/bert/bert_tasks.py
@@ -14,15 +14,21 @@
 """BERT task specific models and heads."""
 
 import copy
+import os
 
 from tensorflow import keras
 
 from keras_nlp.models.bert.bert_models import Bert
 from keras_nlp.models.bert.bert_models import bert_kernel_initializer
 from keras_nlp.models.bert.bert_presets import backbone_presets
+from keras_nlp.models.bert.bert_presets import classifier_presets
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 
+CLASSIFIER_PRESET_NAMES = ", ".join(
+    list(backbone_presets) + list(classifier_presets)
+)
+
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
 class BertClassifier(keras.Model):
@@ -110,11 +116,15 @@ class BertClassifier(keras.Model):
         return cls(**config)
 
     @classproperty
-    def presets(cls):
+    def backbone_presets(cls):
         return copy.deepcopy(backbone_presets)
 
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy({**backbone_presets, **classifier_presets})
+
     @classmethod
-    @format_docstring(names=", ".join(backbone_presets))
+    @format_docstring(names=CLASSIFIER_PRESET_NAMES)
     def from_preset(
         cls,
         preset,
@@ -158,15 +168,31 @@ class BertClassifier(keras.Model):
         output = classifier(input_data)
         ```
         """
-        # Check if preset is backbone-only model
-        if preset in Bert.presets:
-            backbone = Bert.from_preset(preset, load_weights)
-            return cls(backbone, **kwargs)
-
-        # Otherwise must be one of class presets
-        # Currently no classifier-level presets, so must throw.
         if preset not in cls.presets:
             raise ValueError(
                 "`preset` must be one of "
                 f"""{", ".join(cls.presets)}. Received: {preset}."""
             )
+
+        # Check if preset is backbone-only model
+        if preset in cls.backbone_presets:
+            backbone = Bert.from_preset(preset, load_weights)
+            return cls(backbone, **kwargs)
+
+        # Otherwise must be one of class presets
+        metadata = cls.presets[preset]
+        config = metadata["config"]
+        model = cls.from_config({**config, **kwargs})
+
+        if not load_weights:
+            return model
+
+        weights = keras.utils.get_file(
+            "model.h5",
+            metadata["weights_url"],
+            cache_subdir=os.path.join("models", preset),
+            file_hash=metadata["weights_hash"],
+        )
+
+        model.load_weights(weights)
+        return model
diff --git a/keras_nlp/models/gpt2/gpt2_models.py b/keras_nlp/models/gpt2/gpt2_models.py
index 025581dc..58cdae0a 100644
--- a/keras_nlp/models/gpt2/gpt2_models.py
+++ b/keras_nlp/models/gpt2/gpt2_models.py
@@ -14,11 +14,17 @@
 
 """GPT-2 backbone models."""
 
+import copy
+import os
+
 import tensorflow as tf
 from tensorflow import keras
 
 from keras_nlp.layers import PositionEmbedding
 from keras_nlp.layers import TransformerDecoder
+from keras_nlp.models.gpt2.gpt2_presets import backbone_presets
+from keras_nlp.utils.python_utils import classproperty
+from keras_nlp.utils.python_utils import format_docstring
 
 
 def _gpt_2_kernel_initializer(stddev=0.02):
@@ -186,11 +192,64 @@ class GPT2(keras.Model):
     def from_config(cls, config):
         return cls(**config)
 
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy(backbone_presets)
+
     @classmethod
+    @format_docstring(names=", ".join(backbone_presets))
     def from_preset(
         cls,
         preset,
         load_weights=True,
         **kwargs,
     ):
-        raise NotImplementedError
+        """Instantiate GPT-2 model from preset architecture and weights.
+
+        Args:
+            preset: string. Must be one of {{names}}.
+            load_weights: Whether to load pre-trained weights into model.
+                Defaults to `True`.
+
+        Examples:
+        ```python
+        input_data = {
+            "token_ids": tf.random.uniform(
+                shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
+            ),
+            "padding_mask": tf.constant(
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+            ),
+        }
+
+        # Load architecture and weights from preset
+        model = GPT2.from_preset("gpt2_base")
+        output = model(input_data)
+
+        # Load randomly initalized model from preset architecture
+        model = GPT2.from_preset("gpt2_base", load_weights=False)
+        output = model(input_data)
+        ```
+        """
+
+        if preset not in cls.presets:
+            raise ValueError(
+                "`preset` must be one of "
+                f"""{", ".join(cls.presets)}. Received: {preset}."""
+            )
+        metadata = cls.presets[preset]
+        config = metadata["config"]
+        model = cls.from_config({**config, **kwargs})
+
+        if not load_weights:
+            return model
+
+        weights = keras.utils.get_file(
+            "model.h5",
+            metadata["weights_url"],
+            cache_subdir=os.path.join("models", preset),
+            file_hash=metadata["weights_hash"],
+        )
+
+        model.load_weights(weights)
+        return model
diff --git a/keras_nlp/models/gpt2/gpt2_preprocessing.py b/keras_nlp/models/gpt2/gpt2_preprocessing.py
index 01c39b93..b2a76069 100644
--- a/keras_nlp/models/gpt2/gpt2_preprocessing.py
+++ b/keras_nlp/models/gpt2/gpt2_preprocessing.py
@@ -13,10 +13,15 @@
 # limitations under the License.
 """GPT-2 preprocessing layers."""
 
+import copy
+import os
+
 from tensorflow import keras
 
+from keras_nlp.models.gpt2.gpt2_presets import backbone_presets
 from keras_nlp.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 from keras_nlp.utils.python_utils import classproperty
+from keras_nlp.utils.python_utils import format_docstring
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
@@ -43,8 +48,7 @@ class GPT2Tokenizer(BytePairTokenizer):
         merges: string or list, contains the merge rule. If it is a string,
             it should be the file path to merge rules. The merge rule file
             should have one merge rule per line. Every merge rule contains
-            merge entities separated by a space. Please refer to this example:
-            https://storage.googleapis.com/keras-nlp/models/roberta_base/merges.txt.
+            merge entities separated by a space.
 
     Examples:
 
@@ -113,12 +117,59 @@ class GPT2Tokenizer(BytePairTokenizer):
 
     @classproperty
     def presets(cls):
-        raise NotImplementedError
+        return copy.deepcopy(backbone_presets)
 
     @classmethod
+    @format_docstring(names=", ".join(backbone_presets))
     def from_preset(
         cls,
         preset,
         **kwargs,
     ):
-        raise NotImplementedError
+        """Instantiate a GPT-2 tokenizer from preset vocabulary and merge rules.
+
+        Args:
+            preset: string. Must be one of {{names}}.
+
+        Examples:
+        ```python
+        # Load a preset tokenizer.
+        tokenizer = keras_nlp.models.GPT2Tokenizer.from_preset(
+            "gpt2_base",
+        )
+        # Tokenize some input.
+        tokenizer("The quick brown fox tripped.")
+        # Detokenize some input.
+        tokenizer.detokenize([5, 6, 7, 8, 9])
+        ```
+        """
+
+        if preset not in cls.presets:
+            raise ValueError(
+                "`preset` must be one of "
+                f"""{", ".join(cls.presets)}. Received: {preset}."""
+            )
+        metadata = cls.presets[preset]
+
+        vocabulary = keras.utils.get_file(
+            "vocab.json",
+            metadata["vocabulary_url"],
+            cache_subdir=os.path.join("models", preset),
+            file_hash=metadata["vocabulary_hash"],
+        )
+        merges = keras.utils.get_file(
+            "merges.txt",
+            metadata["merges_url"],
+            cache_subdir=os.path.join("models", preset),
+            file_hash=metadata["merges_hash"],
+        )
+
+        config = metadata["preprocessor_config"]
+        config.update(
+            {
+                "vocabulary": vocabulary,
+                "merges": merges,
+            },
+        )
+
+        return cls.from_config({**config, **kwargs})
diff --git a/keras_nlp/models/gpt2/gpt2_presets.py b/keras_nlp/models/gpt2/gpt2_presets.py
new file mode 100644
index 00000000..e42762c5
--- /dev/null
+++ b/keras_nlp/models/gpt2/gpt2_presets.py
@@ -0,0 +1,102 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Metadata for loading pretrained model weights.
+backbone_presets = {
+    "gpt2_base": {
+        "config": {
+            "vocabulary_size": 50257,
+            "num_layers": 12,
+            "num_heads": 12,
+            "hidden_dim": 768,
+            "intermediate_dim": 3072,
+            "dropout": 0.1,
+            "max_sequence_length": 1024,
+        },
+        "preprocessor_config": {},
+        "description": (
+            "Base size of GPT-2 with 124M parameters. Trained on WebText."
+        ),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/gpt2_base/model.h5",
+        "weights_hash": "f4ea6e1b214516dd7de452461ee6e16e",
+        "vocabulary_url": "https://storage.googleapis.com/keras-nlp/models/gpt2_base/vocab.json",
+        "vocabulary_hash": "dffec25a898b1f5e569bec4dffd7e5c0",
+        "merges_url": "https://storage.googleapis.com/keras-nlp/models/gpt2_base/merges.txt",
+        "merges_hash": "75a37753dd7a28a2c5df80c28bf06e4e",
+    },
+    "gpt2_medium": {
+        "config": {
+            "vocabulary_size": 50257,
+            "num_layers": 24,
+            "num_heads": 16,
+            "hidden_dim": 1024,
+            "intermediate_dim": 4096,
+            "dropout": 0.1,
+            "max_sequence_length": 1024,
+        },
+        "preprocessor_config": {},
+        "description": (
+            "Medium size of GPT-2 with 355M parameters. Trained on WebText."
+        ),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/gpt2_medium/model.h5",
+        "weights_hash": "580ff9b79c04fc90e6d6f47e975c5afe",
+        "vocabulary_url": "https://storage.googleapis.com/keras-nlp/models/gpt2_medium/vocab.json",
+        "vocabulary_hash": "dffec25a898b1f5e569bec4dffd7e5c0",
+        "merges_url": "https://storage.googleapis.com/keras-nlp/models/gpt2_medium/merges.txt",
+        "merges_hash": "75a37753dd7a28a2c5df80c28bf06e4e",
+    },
+    "gpt2_large": {
+        "config": {
+            "vocabulary_size": 50257,
+            "num_layers": 36,
+            "num_heads": 20,
+            "hidden_dim": 1280,
+            "intermediate_dim": 5120,
+            "dropout": 0.1,
+            "max_sequence_length": 1024,
+        },
+        "preprocessor_config": {},
+        "description": (
+            "Large size of GPT-2 with 774M parameters. Trained on WebText."
+        ),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/gpt2_large/model.h5",
+        "weights_hash": "67957cb3dfc9e965960dabe068811e1a",
+        "vocabulary_url": "https://storage.googleapis.com/keras-nlp/models/gpt2_large/vocab.json",
+        "vocabulary_hash": "dffec25a898b1f5e569bec4dffd7e5c0",
+        "merges_url": "https://storage.googleapis.com/keras-nlp/models/gpt2_large/merges.txt",
+        "merges_hash": "75a37753dd7a28a2c5df80c28bf06e4e",
+    },
+    "gpt2_extra_large": {
+        "config": {
+            "vocabulary_size": 50257,
+            "num_layers": 48,
+            "num_heads": 25,
+            "hidden_dim": 1600,
+            "intermediate_dim": 6400,
+            "dropout": 0.1,
+            "max_sequence_length": 1024,
+        },
+        "preprocessor_config": {},
+        "description": (
+            "Extra large size of GPT-2 with 1558M parameters. "
+            "Trained on WebText."
+        ),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/gpt2_extra_large/model.h5",
+        "weights_hash": "d093c1ee0d9705d845c0190909aa2917",
+        "vocabulary_url": "https://storage.googleapis.com/keras-nlp/models/gpt2_extra_large/vocab.json",
+        "vocabulary_hash": "dffec25a898b1f5e569bec4dffd7e5c0",
+        "merges_url": "https://storage.googleapis.com/keras-nlp/models/gpt2_extra_large/merges.txt",
+        "merges_hash": "75a37753dd7a28a2c5df80c28bf06e4e",
+    },
+}
diff --git a/keras_nlp/utils/keras_utils.py b/keras_nlp/utils/keras_utils.py
index 36d7126c..d9a0b245 100644
--- a/keras_nlp/utils/keras_utils.py
+++ b/keras_nlp/utils/keras_utils.py
@@ -27,3 +27,21 @@ def clone_initializer(initializer):
         return initializer
     config = initializer.get_config()
     return initializer.__class__.from_config(config)
+
+
+def pack_x_y_sample_weight(x, y=None, sample_weight=None):
+    """Packs user-provided data into a tuple.
+
+    This is a temporary copy of `keras.utils.pack_x_y_sample_weight` while we
+    wait for the a change to the upstream version to propagate to a stable
+    release. See https://github.com/keras-team/keras-nlp/issues/492
+    """
+    if y is None:
+        if not isinstance(x, (list, tuple)):
+            return x
+        else:
+            return (x,)
+    elif sample_weight is None:
+        return (x, y)
+    else:
+        return (x, y, sample_weight)
diff --git a/keras_nlp/utils/pipeline_model.py b/keras_nlp/utils/pipeline_model.py
new file mode 100644
index 00000000..c3579f08
--- /dev/null
+++ b/keras_nlp/utils/pipeline_model.py
@@ -0,0 +1,286 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A base class for models including preprocessing."""
+
+import functools
+import math
+
+import numpy as np
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
+
+try:
+    import pandas as pd
+except ImportError:
+    pd = None
+
+
+def _get_tensor_types():
+    if pd is None:
+        return (tf.Tensor, np.ndarray)
+    else:
+        return (tf.Tensor, np.ndarray, pd.Series, pd.DataFrame)
+
+
+def _convert_inputs_to_dataset(
+    x=None,
+    y=None,
+    sample_weight=None,
+    batch_size=None,
+):
+    """Convert inputs to a `tf.data.Dataset`.
+
+    This is a stand in for the `TensorLikeDataAdapter` in core Keras.
+    """
+    if isinstance(x, tf.data.Dataset):
+        if y is not None:
+            raise ValueError(
+                "When `x` is a `tf.data.Dataset`, please do not provide "
+                f"`y`. Received: `type(y)={type(y)}`."
+            )
+        if sample_weight is not None:
+            raise ValueError(
+                "When `x` is a `tf.data.Dataset`, please do not provide "
+                "`sample_weight`. Received: "
+                f"`type(sample_weight)={type(sample_weight)}`."
+            )
+        if batch_size is not None:
+            raise ValueError(
+                "When `x` is a `tf.data.Dataset`, please do not provide "
+                "`batch_size`. Received: "
+                f"`type(batch_size)={type(batch_size)}`."
+            )
+        return x
+
+    inputs = keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
+    return tf.data.Dataset.from_tensor_slices(inputs).batch(batch_size or 32)
+
+
+def _train_validation_split(arrays, validation_split):
+    """Split arrays into train and validation subsets in deterministic order.
+
+    This is copied directly from core Keras.
+    """
+
+    def _can_split(t):
+        tensor_types = _get_tensor_types()
+        return isinstance(t, tensor_types) or t is None
+
+    flat_arrays = tf.nest.flatten(arrays)
+    unsplitable = [type(t) for t in flat_arrays if not _can_split(t)]
+    if unsplitable:
+        raise ValueError(
+            "`validation_split` is only supported for Tensors or NumPy "
+            "arrays, found following types in the input: {}".format(unsplitable)
+        )
+
+    if all(t is None for t in flat_arrays):
+        return arrays, arrays
+
+    first_non_none = None
+    for t in flat_arrays:
+        if t is not None:
+            first_non_none = t
+            break
+
+    # Assumes all arrays have the same batch shape or are `None`.
+    batch_dim = int(first_non_none.shape[0])
+    split_at = int(math.floor(batch_dim * (1.0 - validation_split)))
+
+    if split_at == 0 or split_at == batch_dim:
+        raise ValueError(
+            "Training data contains {batch_dim} samples, which is not "
+            "sufficient to split it into a validation and training set as "
+            "specified by `validation_split={validation_split}`. Either "
+            "provide more data, or a different value for the "
+            "`validation_split` argument.".format(
+                batch_dim=batch_dim, validation_split=validation_split
+            )
+        )
+
+    def _split(t, start, end):
+        if t is None:
+            return t
+        return t[start:end]
+
+    train_arrays = tf.nest.map_structure(
+        functools.partial(_split, start=0, end=split_at), arrays
+    )
+    val_arrays = tf.nest.map_structure(
+        functools.partial(_split, start=split_at, end=batch_dim), arrays
+    )
+
+    return train_arrays, val_arrays
+
+
+class PipelineModel(keras.Model):
+    """A model which allows automatically applying preprocessing."""
+
+    def __init__(self, *args, include_preprocessing=True, **kwargs):
+        # Workaround for https://github.com/keras-team/keras/issues/17270
+        # Reset any attempt to overwrite this classes base class to this class
+        # can continue to be used for functional and non-functional models.
+        PipelineModel.__bases__ = (keras.Model,)
+        super().__init__(*args, **kwargs)
+        self.include_preprocessing = include_preprocessing
+
+    def preprocess_features(self, x):
+        """An overridable function which preprocesses features."""
+        return x
+
+    def preprocess_labels(self, y):
+        """An overridable function which preprocesses labels."""
+        return y
+
+    def preprocess_samples(self, x, y=None, sample_weight=None):
+        """An overridable function which preprocesses entire samples."""
+        x = self.preprocess_features(x)
+        if y is not None:
+            y = self.preprocess_labels(y)
+        return pack_x_y_sample_weight(x, y, sample_weight)
+
+    # ========================================================================
+    # Below are overrides to keras.Model methods to apply the functions above.
+    # ========================================================================
+    def fit(
+        self,
+        x=None,
+        y=None,
+        batch_size=None,
+        sample_weight=None,
+        validation_data=None,
+        validation_split=None,
+        **kwargs,
+    ):
+        if validation_split and validation_data is None:
+            (x, y, sample_weight), validation_data = _train_validation_split(
+                (x, y, sample_weight), validation_split=validation_split
+            )
+
+        x = _convert_inputs_to_dataset(x, y, sample_weight, batch_size)
+        if self.include_preprocessing:
+            x = x.map(
+                self.preprocess_samples, num_parallel_calls=tf.data.AUTOTUNE
+            ).prefetch(tf.data.AUTOTUNE)
+
+        if validation_data is not None:
+            if not isinstance(validation_data, tf.data.Dataset):
+                (vx, vy, vsw) = keras.utils.unpack_x_y_sample_weight(
+                    validation_data
+                )
+                validation_data = _convert_inputs_to_dataset(
+                    vx, vy, vsw, batch_size
+                )
+
+        return super().fit(
+            x=x,
+            y=None,
+            batch_size=None,
+            sample_weight=None,
+            validation_data=validation_data,
+            **kwargs,
+        )
+
+    def evaluate(
+        self,
+        x=None,
+        y=None,
+        batch_size=None,
+        sample_weight=None,
+        **kwargs,
+    ):
+        # During `fit()`, `keras.Model` attempts to cache the validation
+        # dataset and ignores the values for `x`, `y`, and `sample_weight`.
+        # We don't want that behavior here, as the validation dataset still
+        # needs preprocessing.
+        kwargs.pop("_use_cached_eval_dataset", None)
+        x = _convert_inputs_to_dataset(x, y, sample_weight, batch_size)
+        if self.include_preprocessing:
+            x = x.map(
+                self.preprocess_samples, num_parallel_calls=tf.data.AUTOTUNE
+            ).prefetch(tf.data.AUTOTUNE)
+        return super().evaluate(
+            x=x,
+            y=None,
+            batch_size=None,
+            **kwargs,
+        )
+
+    def predict(
+        self,
+        x=None,
+        batch_size=None,
+        **kwargs,
+    ):
+        x = _convert_inputs_to_dataset(x, None, None, batch_size)
+        if self.include_preprocessing:
+            x = x.map(
+                self.preprocess_samples, num_parallel_calls=tf.data.AUTOTUNE
+            ).prefetch(tf.data.AUTOTUNE)
+
+        return super().predict(
+            x=x,
+            batch_size=None,
+            **kwargs,
+        )
+
+    def train_on_batch(
+        self,
+        x,
+        y=None,
+        sample_weight=None,
+        **kwargs,
+    ):
+        if self.include_preprocessing:
+            data = self.preprocess_samples(x, y, sample_weight)
+            x, y, sample_weight = tf.keras.utils.unpack_x_y_sample_weight(data)
+        return super().train_on_batch(
+            x=x,
+            y=y,
+            sample_weight=sample_weight,
+            **kwargs,
+        )
+
+    def test_on_batch(
+        self,
+        x,
+        y=None,
+        sample_weight=None,
+        **kwargs,
+    ):
+        if self.include_preprocessing:
+            data = self.preprocess_samples(x, y, sample_weight)
+            x, y, sample_weight = tf.keras.utils.unpack_x_y_sample_weight(data)
+        return super().test_on_batch(
+            x=x,
+            y=y,
+            sample_weight=sample_weight,
+            **kwargs,
+        )
+
+    def predict_on_batch(
+        self,
+        x,
+        **kwargs,
+    ):
+        if self.include_preprocessing:
+            data = self.preprocess_samples(x)
+            x, _, _ = tf.keras.utils.unpack_x_y_sample_weight(data)
+        return super().predict_on_batch(
+            x=x,
+            **kwargs,
+        )
diff --git a/requirements.txt b/requirements.txt
index 5db761ea..d733898b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,6 @@
 # Core deps.
-tensorflow==2.10
-tensorflow-text==2.10
+tensorflow~=2.11.0
+tensorflow-text~=2.11.0
 
 # Common deps.
 -r requirements-common.txt
diff --git a/setup.py b/setup.py
index cfb153fc..797b0084 100644
--- a/setup.py
+++ b/setup.py
@@ -29,7 +29,7 @@ setup(
     ),
     long_description=README,
     long_description_content_type="text/markdown",
-    version="0.3.0",
+    version="0.4.0",
     url="https://github.com/keras-team/keras-nlp",
     author="Keras team",
     author_email="keras-nlp@google.com",
@@ -60,9 +60,16 @@ setup(
             "keras-tuner",
         ],
     },
+    # Supported Python versions
+    python_requires=">=3.7",
     classifiers=[
-        "Programming Language :: Python",
+        "Development Status :: 3 - Alpha",
+        "Programming Language :: Python :: 3",
         "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3 :: Only",
         "Operating System :: Unix",
         "Operating System :: Microsoft :: Windows",
         "Operating System :: MacOS",
diff --git a/shell/format.sh b/shell/format.sh
index 2eee93f9..117b08c8 100755
--- a/shell/format.sh
+++ b/shell/format.sh
@@ -1,7 +1,7 @@
 #!/bin/bash -e
 
 base_dir=$(dirname $(dirname $0))
-targets="${base_dir}/*.py ${base_dir}/examples/ ${base_dir}/keras_nlp/"
+targets="${base_dir}/*.py ${base_dir}/examples/ ${base_dir}/keras_nlp/ ${base_dir}/tools/checkpoint_conversion/*.py"
 
 isort --sp "${base_dir}/setup.cfg" --sl ${targets}
 black --line-length 80 ${targets}
diff --git a/tools/checkpoint_conversion/checkpoint_conversion_utils.py b/tools/checkpoint_conversion/checkpoint_conversion_utils.py
new file mode 100644
index 00000000..2301b8db
--- /dev/null
+++ b/tools/checkpoint_conversion/checkpoint_conversion_utils.py
@@ -0,0 +1,22 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import hashlib
+
+
+def get_md5_checksum(file_path):
+    md5_hash = hashlib.md5()
+    with open(file_path, "rb") as f:
+        for byte_block in iter(lambda: f.read(4096), b""):
+            md5_hash.update(byte_block)
+    return md5_hash.hexdigest()
diff --git a/tools/checkpoint_conversion/convert_gpt2_checkpoints.py b/tools/checkpoint_conversion/convert_gpt2_checkpoints.py
new file mode 100644
index 00000000..6ba60ad1
--- /dev/null
+++ b/tools/checkpoint_conversion/convert_gpt2_checkpoints.py
@@ -0,0 +1,278 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import os
+
+import numpy as np
+import requests
+import tensorflow as tf
+import transformers
+from absl import app
+from absl import flags
+from checkpoint_conversion_utils import get_md5_checksum
+
+import keras_nlp
+
+PRESET_MAP = {
+    "gpt2_base": ("124M", "gpt2"),
+    "gpt2_medium": ("355M", "gpt2-medium"),
+    "gpt2_large": ("774M", "gpt2-large"),
+    "gpt2_extra_large": ("1558M", "gpt2-xl"),
+}
+
+DOWNLOAD_SCRIPT_URL = (
+    "https://raw.githubusercontent.com/openai/gpt-2/master/download_model.py"
+)
+
+EXTRACT_DIR = "./models/{}"
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string(
+    "preset", None, f'Must be one of {",".join(PRESET_MAP.keys())}'
+)
+
+
+def download_model(preset, num_params):
+    print("-> Download original weights.")
+    response = requests.get(DOWNLOAD_SCRIPT_URL)
+    open("download_model.py", "wb").write(response.content)
+
+    os.system(f"python download_model.py {num_params}")
+
+
+def convert_checkpoints(preset, num_params):
+    print("\n-> Convert original weights to KerasNLP format.")
+    # GPT-2 paths.
+    extract_dir = EXTRACT_DIR.format(num_params)
+    checkpoint_path = os.path.join(extract_dir, "model.ckpt")
+    config_path = os.path.join(extract_dir, "hparams.json")
+
+    with open(config_path, "r") as f:
+        cfg = json.load(f)
+    print("Config:", cfg)
+
+    print("Original weights:")
+    vars = tf.train.list_variables(checkpoint_path)
+    weights = {}
+    for name, shape in vars:
+        print(name, shape)
+        weight = tf.train.load_variable(checkpoint_path, name)
+        weights[name] = weight
+
+    keras_nlp_model = keras_nlp.models.GPT2.from_preset(
+        preset,
+        load_weights=False,
+    )
+
+    keras_nlp_model.get_layer("token_embedding").embeddings.assign(
+        weights["model/wte"]
+    )
+    keras_nlp_model.get_layer("position_embedding").position_embeddings.assign(
+        weights["model/wpe"]
+    )
+
+    range_1 = (0, cfg["n_embd"])
+    range_2 = (cfg["n_embd"], 2 * cfg["n_embd"])
+    range_3 = (2 * cfg["n_embd"], 3 * cfg["n_embd"])
+
+    for i in range(keras_nlp_model.num_layers):
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._self_attention_layer._query_dense.kernel.assign(
+            weights[f"model/h{i}/attn/c_attn/w"][
+                0, :, range_1[0] : range_1[1]
+            ].reshape((cfg["n_embd"], cfg["n_head"], -1))
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._self_attention_layer._query_dense.bias.assign(
+            weights[f"model/h{i}/attn/c_attn/b"][
+                range_1[0] : range_1[1]
+            ].reshape((cfg["n_head"], -1))
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._self_attention_layer._key_dense.kernel.assign(
+            weights[f"model/h{i}/attn/c_attn/w"][
+                0, :, range_2[0] : range_2[1]
+            ].reshape((cfg["n_embd"], cfg["n_head"], -1))
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._self_attention_layer._key_dense.bias.assign(
+            weights[f"model/h{i}/attn/c_attn/b"][
+                range_2[0] : range_2[1]
+            ].reshape((cfg["n_head"], -1))
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._self_attention_layer._value_dense.kernel.assign(
+            weights[f"model/h{i}/attn/c_attn/w"][
+                0, :, range_3[0] : range_3[1]
+            ].reshape((cfg["n_embd"], cfg["n_head"], -1))
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._self_attention_layer._value_dense.bias.assign(
+            weights[f"model/h{i}/attn/c_attn/b"][
+                range_3[0] : range_3[1]
+            ].reshape((cfg["n_head"], -1))
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._self_attention_layer._output_dense.kernel.assign(
+            weights[f"model/h{i}/attn/c_proj/w"][0].reshape(
+                (cfg["n_head"], -1, cfg["n_embd"])
+            )
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._self_attention_layer._output_dense.bias.assign(
+            weights[f"model/h{i}/attn/c_proj/b"]
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._self_attention_layernorm.gamma.assign(weights[f"model/h{i}/ln_1/g"])
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._self_attention_layernorm.beta.assign(weights[f"model/h{i}/ln_1/b"])
+
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._feedforward_intermediate_dense.kernel.assign(
+            weights[f"model/h{i}/mlp/c_fc/w"][0]
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._feedforward_intermediate_dense.bias.assign(
+            weights[f"model/h{i}/mlp/c_fc/b"]
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._feedforward_output_dense.kernel.assign(
+            weights[f"model/h{i}/mlp/c_proj/w"][0]
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._feedforward_output_dense.bias.assign(
+            weights[f"model/h{i}/mlp/c_proj/b"]
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._feedforward_layernorm.gamma.assign(weights[f"model/h{i}/ln_2/g"])
+        keras_nlp_model.get_layer(
+            f"transformer_layer_{i}"
+        )._feedforward_layernorm.beta.assign(weights[f"model/h{i}/ln_2/b"])
+
+    keras_nlp_model.get_layer("layer_norm").gamma.assign(
+        weights["model/ln_f/g"]
+    )
+
+    keras_nlp_model.get_layer("layer_norm").beta.assign(weights["model/ln_f/b"])
+
+    # Save the model.
+    print(f"\n-> Save KerasNLP model weights to `{preset}.h5`.")
+    keras_nlp_model.save_weights(f"{preset}.h5")
+
+    return keras_nlp_model
+
+
+def define_tokenizer(preset, num_params, hf_model_name):
+    print("\n-> Define the tokenizers.")
+    extract_dir = extract_dir = EXTRACT_DIR.format(num_params)
+    merges_path = os.path.join(extract_dir, "vocab.bpe")
+    vocab_path = os.path.join(extract_dir, "encoder.json")
+
+    keras_nlp_tokenizer = keras_nlp.models.GPT2Tokenizer(
+        vocabulary=vocab_path,
+        merges=merges_path,
+    )
+    hf_tokenizer = transformers.AutoTokenizer.from_pretrained(hf_model_name)
+
+    print("\n-> Print MD5 checksum of the vocab files.")
+    print(f"`{vocab_path}` md5sum: ", get_md5_checksum(vocab_path))
+    print(f"`{merges_path}` md5sum: ", get_md5_checksum(merges_path))
+
+    return keras_nlp_tokenizer, hf_tokenizer
+
+
+def check_output(
+    preset,
+    keras_nlp_model,
+    keras_nlp_tokenizer,
+    hf_model,
+    hf_tokenizer,
+):
+    print("\n-> Check the outputs.")
+    input_str = ["the quick brown fox ran, galloped and jumped."]
+
+    # KerasNLP
+    token_ids = keras_nlp_tokenizer(input_str)
+    padding_mask = token_ids != 0
+
+    keras_nlp_inputs = {
+        "token_ids": token_ids.to_tensor(),
+        "padding_mask": padding_mask.to_tensor(),
+    }
+    keras_nlp_output = keras_nlp_model.predict(keras_nlp_inputs)
+
+    # HF
+    hf_inputs = hf_tokenizer(input_str, return_tensors="pt")
+    hf_output = hf_model(**hf_inputs).last_hidden_state
+
+    print("KerasNLP output:", keras_nlp_output[0, 0, :10])
+    print("HF output:", hf_output[0, 0, :10])
+    print("Difference:", np.mean(keras_nlp_output - hf_output.detach().numpy()))
+
+    # Show the MD5 checksum of the model weights.
+    print("Model md5sum: ", get_md5_checksum(f"./{preset}.h5"))
+
+    return keras_nlp_output
+
+
+def main(_):
+    assert (
+        FLAGS.preset in PRESET_MAP.keys()
+    ), f'Invalid preset {FLAGS.preset}. Must be one of {",".join(PRESET_MAP.keys())}'
+    num_params = PRESET_MAP[FLAGS.preset][0]
+    hf_model_name = PRESET_MAP[FLAGS.preset][1]
+
+    download_model(FLAGS.preset, num_params)
+
+    keras_nlp_model = convert_checkpoints(FLAGS.preset, num_params)
+
+    print("\n-> Load HF model.")
+    hf_model = transformers.AutoModel.from_pretrained(hf_model_name)
+    hf_model.eval()
+
+    keras_nlp_tokenizer, hf_tokenizer = define_tokenizer(
+        FLAGS.preset, num_params, hf_model_name
+    )
+
+    check_output(
+        FLAGS.preset,
+        keras_nlp_model,
+        keras_nlp_tokenizer,
+        hf_model,
+        hf_tokenizer,
+    )
+
+
+if __name__ == "__main__":
+    flags.mark_flag_as_required("preset")
+    app.run(main)
diff --git a/tools/checkpoint_conversion/requirements.txt b/tools/checkpoint_conversion/requirements.txt
new file mode 100644
index 00000000..4f492ddc
--- /dev/null
+++ b/tools/checkpoint_conversion/requirements.txt
@@ -0,0 +1,2 @@
+torch
+transformers
diff --git a/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb b/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
new file mode 100644
index 00000000..66e19078
--- /dev/null
+++ b/tools/checkpoint_training/bert_tiny_uncased_en_sst2_training.ipynb
@@ -0,0 +1,4059 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "view-in-github"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/jbischof/keras-nlp/blob/classifier_preset/bert_tiny_uncased_en_sst2_training.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "2h__HbR944VZ",
+        "outputId": "ebbee015-5caf-453b-d246-42ce6f2d846d",
+        "vscode": {
+          "languageId": "python"
+        }
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\u001b[K     |████████████████████████████████| 578.0 MB 16 kB/s \n",
+            "\u001b[K     |████████████████████████████████| 5.8 MB 72.0 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 5.9 MB 68.7 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 438 kB 88.6 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 1.7 MB 80.4 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 5.9 MB 61.9 MB/s \n",
+            "\u001b[?25h  Building wheel for keras-nlp (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
+          ]
+        }
+      ],
+      "source": [
+        "!pip install -q git+https://github.com/mattdangerw/keras-nlp.git@bert-pipeline tensorflow==2.10 --upgrade"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "rlmPYqUn7YJm",
+        "vscode": {
+          "languageId": "python"
+        }
+      },
+      "outputs": [],
+      "source": [
+        "import keras_nlp\n",
+        "import tensorflow as tf\n",
+        "from tensorflow import keras\n",
+        "import tensorflow_datasets as tfds"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_hE-veD7utTT"
+      },
+      "source": [
+        "# Data"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 498,
+          "referenced_widgets": [
+            "7dac43c361074f0f998d8c150aa922b9",
+            "02ecb50a721f450b940172d3d35a4abe",
+            "3977c5ac1af645b0b2c1e5886ac037ad",
+            "e9a881d7fcb84bafb6a063009bf4dd54",
+            "efdf1618f9224961b840bf58fb763876",
+            "071de690987f46d5b11d30fd416d0903",
+            "75083b32cd854cdd9de6490b7f828140",
+            "964b0fedfdbe4b4a990d3b72d68b4112",
+            "2308d8d3d7064c62ae470850e9b06a29",
+            "77c0fa586d4d4dcd982513c827799adb",
+            "a70a3e92a37247eeaad91ef1e140aa54",
+            "6bed029127e44fefa0ccd576edb855e9",
+            "289e60ba017345b1b61adeecc24bc2cb",
+            "bb9d3ced705148b6a8c3b71498bd9d81",
+            "4579a6614cd64cdcab0ef6c615cad8bf",
+            "5b7a72d055cd4cb68a4a9706bd276e1f",
+            "25e82eaf0b9849ada515c712b359b21d",
+            "1aca96855cbe4b41a050bdb59f97a903",
+            "078e2f25095d40cc9a96ad3488030022",
+            "16dc3ec494f6456794b43d33a95ef042",
+            "8c3528b24aba4002b71168fd0b8beb44",
+            "d9d4d67205d340a8acb919b7c80a842a",
+            "71a33b25dece4caca90a91160ea2382f",
+            "c84704566a0a45249c6973c8e57e3856",
+            "44ff2d0794b74097af53a8116da941a7",
+            "b22d3db2d68046a1afbca22b13ab7152",
+            "acbe1228b9bb42da9fc19e57e4dc6e74",
+            "06fa78192fd74020962d609ac59af6d5",
+            "4de362ac37224f968d2f6a4f2336b327",
+            "fac49204e83e49c9920c29a61382d9b7",
+            "8bf597ef7f7a4ce1a66cbfc1aa74a216",
+            "b96bbee1b3254bc397828295b8b6528a",
+            "4d7f2052da7f43f2a41ebfaaf562ae77",
+            "c96bafeb19454961aa95b819bec135df",
+            "879fa9af242b4c698c15747315a91ba9",
+            "cda35a02449442e48f3b6131547e1628",
+            "9b65c0c17d3d49589feed767a4ced04c",
+            "09d1f164a0e941438f06a3feb58bdd30",
+            "27d5831d71b547c081cb2dec9a42b300",
+            "14df05ca9bc642d9ad588db644ab7cf7",
+            "8fe22f392c204fdeb8c8b34b9c4bdc12",
+            "97717222d177460bad6c626004f3dd9d",
+            "78d934805a364c8989efba79bead117d",
+            "5dd04d5d68a54166b432bdf53dad0792",
+            "96d397cf27714dca8a8545c3ceb13d73",
+            "91a8aba5cca943fcb2f3f3353450aa78",
+            "10e9f95c33dd4b7b88150d6a7825e973",
+            "d929b9d20b1a4b00bb5902c0883b4034",
+            "c1976cfcd5c4438bb769ad5a77b0b265",
+            "42117d78043a4336bc9c90c88741c15a",
+            "b28ba82c231545e0b509a93ba0470449",
+            "d410bad163c24c1aa5e00d54090b05db",
+            "9eb1c4b7b7e94f8599705b2d1d509497",
+            "e9113eab130140f9b9367a61b1e4b5c0",
+            "454dfc6558ae4a1b83621ea1352fc051",
+            "1f5f58f9400b48a991a26786568c75af",
+            "d9a622959c0d4c45b0b1932d7c85893d",
+            "6e2dd84de8b3483e9f18142fa30177eb",
+            "6dc9fbdebed24e35bccaf0e92a6a491e",
+            "d147c549823e407eb843750f0ec45cbc",
+            "758811c0d8284c76984fc2ac974d835a",
+            "4e06b19ebe63441ebb6a9b1cd57e2d21",
+            "69a3fcf1ec174cbaa11f93fa1dded564",
+            "2522eb90704c4fd6be3a0223dd16442e",
+            "b0999cce50384ccfa1fb6dedcf6bc1d3",
+            "7f65a8e8ab6a411286d11c0125b316d0",
+            "cd67f31d575f44d9944a5019b3e38c19",
+            "cb1357c15a4c4a28a6198a61af82848b",
+            "0a9d89c682624b208b399ea9d8c26b0b",
+            "130007f46b794a4c9d283aa5fb333d04",
+            "a7fb5284dab64287b108a848c7203dd5",
+            "e70cb8396ab341c1bb0d60d2ebdd708a",
+            "d5020ccf7aff400988f0f0e6d1ff3557",
+            "8a0f69f4dc4040559001183e68b73aff",
+            "6f7051dbe4ad467bada9e50157141145",
+            "90e15c8359274c5a85a7cf427cd9ca22",
+            "b3086aaf2d8b491fbb5569d6226bfd29",
+            "52fca5d8bd9b4229bd64ebce81043160",
+            "8388b866783e4d9eb11093a99eaa3116",
+            "5df67222f0f74e1c9b57a59dbb1d3d98",
+            "04a3c2356d10466fb1683d202208f66d",
+            "335fe992357d4ca298312a55466c1d41",
+            "e5d9c734fb3e412aabc03e5790b12382",
+            "07fd35fa2e634720b064998b96e0e60e",
+            "ba64c4461ad14990afe4c1df45ad8d7d",
+            "30b1448d901a4191ab5bd62e89028e75",
+            "cbd21b3e1e5f4e318940805ac4d142ec",
+            "20ffd17f760c45788040797926b7fc35",
+            "a0d66c335f174159a7b2df88cb4b607f",
+            "58dfc96028a948c7af85f4df518bc171",
+            "436d985f48104913933079fdb44fde6f",
+            "aca24255c81943aeb92a036311e53cef",
+            "c7a31a4a0e4d4b27a93e48950acfe5e5",
+            "b65aff687d4242fa94ede609a48ad847",
+            "1fc728a0824144d1a60399e3d8529ea9",
+            "5634ff298df447dd92f4b34cf382b7e0",
+            "4b2ed16c315e413188360923ec7e19eb",
+            "80053cb69a2a4b2c8ad3c038fc3e2c39",
+            "7c9acf3891a040f5af892b86e72c803c",
+            "b4a60897f73642eeb692cafa7997167c",
+            "c4d8f2f2b9604745ac03fe5f4da7eeee",
+            "87e1ae4c68394e23ba0ed21f8fffb5a6",
+            "670f7df971d14fabbffcfd2535e90005",
+            "148f6232f6db44919f04a78c70aba8c1",
+            "66071ddc16014e9a8604f37fe1441b33",
+            "88b26fe32ba64f9e90c10b6c40381635",
+            "4679a29c8e2c4fa686713038e1c4686e",
+            "b9d49ffe543342c694905d2dee72b880",
+            "93c4b10a8f074c2c95cee65850ab2040",
+            "275c1f7417e247ce99519f0ef43738c9"
+          ]
+        },
+        "id": "e3Q1wpT6mh5o",
+        "outputId": "2537d487-6dc3-4072-a03d-0e5f1879e1e1",
+        "vscode": {
+          "languageId": "python"
+        }
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Downloading and preparing dataset 7.09 MiB (download: 7.09 MiB, generated: 7.22 MiB, total: 14.31 MiB) to ~/tensorflow_datasets/glue/sst2/2.0.0...\n"
+          ]
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "7dac43c361074f0f998d8c150aa922b9",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "Dl Completed...: 0 url [00:00, ? url/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "6bed029127e44fefa0ccd576edb855e9",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "Dl Size...: 0 MiB [00:00, ? MiB/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "71a33b25dece4caca90a91160ea2382f",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "Extraction completed...: 0 file [00:00, ? file/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "c96bafeb19454961aa95b819bec135df",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "Generating splits...:   0%|          | 0/3 [00:00<?, ? splits/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "96d397cf27714dca8a8545c3ceb13d73",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "Generating train examples...:   0%|          | 0/67349 [00:00<?, ? examples/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "1f5f58f9400b48a991a26786568c75af",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "Shuffling ~/tensorflow_datasets/glue/sst2/2.0.0.incompleteQH5CX4/glue-train.tfrecord*...:   0%|          | 0/6…"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "cd67f31d575f44d9944a5019b3e38c19",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "Generating validation examples...:   0%|          | 0/872 [00:00<?, ? examples/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "52fca5d8bd9b4229bd64ebce81043160",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "Shuffling ~/tensorflow_datasets/glue/sst2/2.0.0.incompleteQH5CX4/glue-validation.tfrecord*...:   0%|          …"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "a0d66c335f174159a7b2df88cb4b607f",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "Generating test examples...:   0%|          | 0/1821 [00:00<?, ? examples/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "b4a60897f73642eeb692cafa7997167c",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "Shuffling ~/tensorflow_datasets/glue/sst2/2.0.0.incompleteQH5CX4/glue-test.tfrecord*...:   0%|          | 0/18…"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Dataset glue downloaded and prepared to ~/tensorflow_datasets/glue/sst2/2.0.0. Subsequent calls will reuse this data.\n"
+          ]
+        },
+        {
+          "data": {
+            "text/plain": [
+              "(<tf.Tensor: shape=(16,), dtype=string, numpy=\n",
+              " array([b'for the uninitiated plays better on video with the sound ',\n",
+              "        b'like a giant commercial for universal studios , where much of the action takes place ',\n",
+              "        b'company once again dazzle and delight us ',\n",
+              "        b\"'s no surprise that as a director washington demands and receives excellent performances , from himself and from newcomer derek luke \",\n",
+              "        b', this cross-cultural soap opera is painfully formulaic and stilted . ',\n",
+              "        b\", the film is n't nearly as downbeat as it sounds , but strikes a tone that 's alternately melancholic , hopeful and strangely funny . \",\n",
+              "        b'only masochistic moviegoers need apply . ',\n",
+              "        b'convince almost everyone that it was put on the screen , just for them ',\n",
+              "        b\"like the english patient and the unbearable lightness of being , the hours is one of those reputedly `` unfilmable '' novels that has bucked the odds to emerge as an exquisite motion picture in its own right . \",\n",
+              "        b'his supple understanding of the role ',\n",
+              "        b'revelatory nor truly edgy -- merely crassly flamboyant and comedically labored . ',\n",
+              "        b'teaming ', b\"'s not very well shot or composed or edited \",\n",
+              "        b'of the startling intimacy ', b'draggy ',\n",
+              "        b'admire this film for its harsh objectivity and refusal to seek our tears , our sympathies . '],\n",
+              "       dtype=object)>,\n",
+              " <tf.Tensor: shape=(16,), dtype=int64, numpy=array([0, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 0, 1])>)"
+            ]
+          },
+          "execution_count": 3,
+          "metadata": {},
+          "output_type": "execute_result"
+        }
+      ],
+      "source": [
+        "train_ds, valid_ds = tfds.load(\n",
+        "    \"glue/sst2\",\n",
+        "    split=[\"train\", \"validation\"],\n",
+        "    batch_size=16,\n",
+        ")\n",
+        "\n",
+        "def split_features(x):\n",
+        "    # GLUE comes with dictonary data, we convert it to a uniform format\n",
+        "    # (features, label), where features is a tuple consisting of all\n",
+        "    # features.\n",
+        "    features = x[\"sentence\"]\n",
+        "    label = x[\"label\"]\n",
+        "    return (features, label)\n",
+        " \n",
+        "train_ds = train_ds.map(\n",
+        "    split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
+        "    ).prefetch(tf.data.AUTOTUNE)\n",
+        "valid_ds = valid_ds.map(\n",
+        "        split_features, num_parallel_calls=tf.data.AUTOTUNE\n",
+        "    ).prefetch(tf.data.AUTOTUNE)\n",
+        "\n",
+        "# Look first training set batch\n",
+        "# The format is (string_tensor, label_tensor)\n",
+        "train_ds.take(1).get_single_element()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "_cj9caiYt_jb",
+        "outputId": "0f09fccb-97c5-44c0-c156-94f182bc5c05",
+        "vscode": {
+          "languageId": "python"
+        }
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Epoch 1/2\n",
+            "4210/4210 [==============================] - 740s 173ms/step - loss: 0.3747 - sparse_categorical_accuracy: 0.8295 - val_loss: 0.4244 - val_sparse_categorical_accuracy: 0.8165\n",
+            "Epoch 2/2\n",
+            "4210/4210 [==============================] - 710s 169ms/step - loss: 0.2378 - sparse_categorical_accuracy: 0.9054 - val_loss: 0.4463 - val_sparse_categorical_accuracy: 0.8234\n"
+          ]
+        },
+        {
+          "data": {
+            "text/plain": [
+              "<keras.callbacks.History at 0x7f6bbd0923d0>"
+            ]
+          },
+          "execution_count": 11,
+          "metadata": {},
+          "output_type": "execute_result"
+        }
+      ],
+      "source": [
+        "classifier = keras_nlp.models.BertClassifier.from_preset(\n",
+        "    \"bert_tiny_uncased_en\",\n",
+        "    num_classes=2,\n",
+        ")\n",
+        "classifier.compile(\n",
+        "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+        "    optimizer=keras.optimizers.experimental.AdamW(5e-5),\n",
+        "    metrics=keras.metrics.SparseCategoricalAccuracy(),\n",
+        "    jit_compile=True,\n",
+        ")\n",
+        "classifier.fit(\n",
+        "    train_ds,\n",
+        "    validation_data=valid_ds,\n",
+        "    epochs=2,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fJcxn1si0dKV",
+        "vscode": {
+          "languageId": "python"
+        }
+      },
+      "outputs": [],
+      "source": [
+        "classifier.save_weights('/content/model.h5')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "jFSJUTZEoUws",
+        "outputId": "93cd3576-0c42-43e0-8fd0-c4d9a42d3a87",
+        "vscode": {
+          "languageId": "python"
+        }
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "-rw-r--r-- 1 root root 17M Nov 28 19:58 /content/model.h5\n",
+            "1f9c2d59f9e229e08f3fbd44239cfb0b  /content/model.h5\n"
+          ]
+        }
+      ],
+      "source": [
+        "!ls -lh /content/model.h5\n",
+        "!md5sum /content/model.h5"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Wk6aVbYyUHBX",
+        "vscode": {
+          "languageId": "python"
+        }
+      },
+      "outputs": [],
+      "source": [
+        "classifier.load_weights('/content/model.h5')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "MJqmTfYtwt8V",
+        "outputId": "726b2a10-c650-4059-e15e-3173001ad444",
+        "vscode": {
+          "languageId": "python"
+        }
+      },
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "{'backbone': {'class_name': 'keras_nlp>Bert',\n",
+              "  'config': {'vocabulary_size': 30522,\n",
+              "   'hidden_dim': 128,\n",
+              "   'intermediate_dim': 512,\n",
+              "   'num_layers': 2,\n",
+              "   'num_heads': 2,\n",
+              "   'max_sequence_length': 512,\n",
+              "   'num_segments': 2,\n",
+              "   'dropout': 0.1,\n",
+              "   'name': 'backbone',\n",
+              "   'trainable': True}},\n",
+              " 'preprocessor': {'class_name': 'keras_nlp>BertPreprocessor',\n",
+              "  'config': {'name': 'bert_preprocessor_1',\n",
+              "   'trainable': True,\n",
+              "   'dtype': 'float32',\n",
+              "   'tokenizer': {'class_name': 'keras_nlp>BertTokenizer',\n",
+              "    'config': {'name': 'bert_tokenizer_1',\n",
+              "     'trainable': True,\n",
+              "     'dtype': 'int32',\n",
+              "     'vocabulary': [],\n",
+              "     'sequence_length': None,\n",
+              "     'lowercase': True,\n",
+              "     'strip_accents': False,\n",
+              "     'split': True,\n",
+              "     'suffix_indicator': '##',\n",
+              "     'oov_token': '[UNK]'}},\n",
+              "   'sequence_length': 512,\n",
+              "   'truncate': 'round_robin'}},\n",
+              " 'num_classes': 2,\n",
+              " 'name': 'bert_classifier_1',\n",
+              " 'trainable': True}"
+            ]
+          },
+          "execution_count": 23,
+          "metadata": {},
+          "output_type": "execute_result"
+        }
+      ],
+      "source": [
+        "config = classifier.get_config()\n",
+        "# Zero out vocab so readable\n",
+        "config[\"preprocessor\"][\"config\"][\"tokenizer\"][\"config\"][\"vocabulary\"] = []\n",
+        "config"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "z113U3yutzY8",
+        "outputId": "763edfc4-fcd1-4216-9afb-276771258f6a",
+        "vscode": {
+          "languageId": "python"
+        }
+      },
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "{'backbone': {'class_name': 'keras_nlp>Bert',\n",
+              "  'config': {'vocabulary_size': 30522,\n",
+              "   'hidden_dim': 128,\n",
+              "   'intermediate_dim': 512,\n",
+              "   'num_layers': 2,\n",
+              "   'num_heads': 2,\n",
+              "   'max_sequence_length': 512,\n",
+              "   'num_segments': 2,\n",
+              "   'dropout': 0.1,\n",
+              "   'name': 'backbone',\n",
+              "   'trainable': True}},\n",
+              " 'num_classes': 2,\n",
+              " 'name': 'bert_classifier_1',\n",
+              " 'trainable': True}"
+            ]
+          },
+          "execution_count": 15,
+          "metadata": {},
+          "output_type": "execute_result"
+        }
+      ],
+      "source": [
+        "config.pop(\"preprocessor\")\n",
+        "config"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "include_colab_link": true,
+      "provenance": []
+    },
+    "gpuClass": "premium",
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "widgets": {
+      "application/vnd.jupyter.widget-state+json": {
+        "02ecb50a721f450b940172d3d35a4abe": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_071de690987f46d5b11d30fd416d0903",
+            "placeholder": "​",
+            "style": "IPY_MODEL_75083b32cd854cdd9de6490b7f828140",
+            "value": "Dl Completed...: 100%"
+          }
+        },
+        "04a3c2356d10466fb1683d202208f66d": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_cbd21b3e1e5f4e318940805ac4d142ec",
+            "placeholder": "​",
+            "style": "IPY_MODEL_20ffd17f760c45788040797926b7fc35",
+            "value": " 0/872 [00:00&lt;?, ? examples/s]"
+          }
+        },
+        "06fa78192fd74020962d609ac59af6d5": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "071de690987f46d5b11d30fd416d0903": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "078e2f25095d40cc9a96ad3488030022": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": "20px"
+          }
+        },
+        "07fd35fa2e634720b064998b96e0e60e": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "09d1f164a0e941438f06a3feb58bdd30": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": "hidden",
+            "width": null
+          }
+        },
+        "0a9d89c682624b208b399ea9d8c26b0b": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_8a0f69f4dc4040559001183e68b73aff",
+            "max": 872,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_6f7051dbe4ad467bada9e50157141145",
+            "value": 872
+          }
+        },
+        "10e9f95c33dd4b7b88150d6a7825e973": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_d410bad163c24c1aa5e00d54090b05db",
+            "max": 67349,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_9eb1c4b7b7e94f8599705b2d1d509497",
+            "value": 67349
+          }
+        },
+        "130007f46b794a4c9d283aa5fb333d04": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_90e15c8359274c5a85a7cf427cd9ca22",
+            "placeholder": "​",
+            "style": "IPY_MODEL_b3086aaf2d8b491fbb5569d6226bfd29",
+            "value": " 634/872 [00:00&lt;00:00, 6335.34 examples/s]"
+          }
+        },
+        "148f6232f6db44919f04a78c70aba8c1": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": "hidden",
+            "width": null
+          }
+        },
+        "14df05ca9bc642d9ad588db644ab7cf7": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "16dc3ec494f6456794b43d33a95ef042": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "1aca96855cbe4b41a050bdb59f97a903": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "1f5f58f9400b48a991a26786568c75af": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_d9a622959c0d4c45b0b1932d7c85893d",
+              "IPY_MODEL_6e2dd84de8b3483e9f18142fa30177eb",
+              "IPY_MODEL_6dc9fbdebed24e35bccaf0e92a6a491e"
+            ],
+            "layout": "IPY_MODEL_d147c549823e407eb843750f0ec45cbc"
+          }
+        },
+        "1fc728a0824144d1a60399e3d8529ea9": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "20ffd17f760c45788040797926b7fc35": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "2308d8d3d7064c62ae470850e9b06a29": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "2522eb90704c4fd6be3a0223dd16442e": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "25e82eaf0b9849ada515c712b359b21d": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "275c1f7417e247ce99519f0ef43738c9": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "27d5831d71b547c081cb2dec9a42b300": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "289e60ba017345b1b61adeecc24bc2cb": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_25e82eaf0b9849ada515c712b359b21d",
+            "placeholder": "​",
+            "style": "IPY_MODEL_1aca96855cbe4b41a050bdb59f97a903",
+            "value": "Dl Size...: 100%"
+          }
+        },
+        "30b1448d901a4191ab5bd62e89028e75": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "335fe992357d4ca298312a55466c1d41": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": "hidden",
+            "width": null
+          }
+        },
+        "3977c5ac1af645b0b2c1e5886ac037ad": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "success",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_964b0fedfdbe4b4a990d3b72d68b4112",
+            "max": 1,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_2308d8d3d7064c62ae470850e9b06a29",
+            "value": 1
+          }
+        },
+        "42117d78043a4336bc9c90c88741c15a": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "436d985f48104913933079fdb44fde6f": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_5634ff298df447dd92f4b34cf382b7e0",
+            "max": 1821,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_4b2ed16c315e413188360923ec7e19eb",
+            "value": 1821
+          }
+        },
+        "44ff2d0794b74097af53a8116da941a7": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "success",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_fac49204e83e49c9920c29a61382d9b7",
+            "max": 1,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_8bf597ef7f7a4ce1a66cbfc1aa74a216",
+            "value": 1
+          }
+        },
+        "454dfc6558ae4a1b83621ea1352fc051": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "4579a6614cd64cdcab0ef6c615cad8bf": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_8c3528b24aba4002b71168fd0b8beb44",
+            "placeholder": "​",
+            "style": "IPY_MODEL_d9d4d67205d340a8acb919b7c80a842a",
+            "value": " 7/7 [00:01&lt;00:00,  3.67 MiB/s]"
+          }
+        },
+        "4679a29c8e2c4fa686713038e1c4686e": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "4b2ed16c315e413188360923ec7e19eb": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "4d7f2052da7f43f2a41ebfaaf562ae77": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "4de362ac37224f968d2f6a4f2336b327": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "4e06b19ebe63441ebb6a9b1cd57e2d21": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "52fca5d8bd9b4229bd64ebce81043160": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_8388b866783e4d9eb11093a99eaa3116",
+              "IPY_MODEL_5df67222f0f74e1c9b57a59dbb1d3d98",
+              "IPY_MODEL_04a3c2356d10466fb1683d202208f66d"
+            ],
+            "layout": "IPY_MODEL_335fe992357d4ca298312a55466c1d41"
+          }
+        },
+        "5634ff298df447dd92f4b34cf382b7e0": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "58dfc96028a948c7af85f4df518bc171": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_b65aff687d4242fa94ede609a48ad847",
+            "placeholder": "​",
+            "style": "IPY_MODEL_1fc728a0824144d1a60399e3d8529ea9",
+            "value": "Generating test examples...:  77%"
+          }
+        },
+        "5b7a72d055cd4cb68a4a9706bd276e1f": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "5dd04d5d68a54166b432bdf53dad0792": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "5df67222f0f74e1c9b57a59dbb1d3d98": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_ba64c4461ad14990afe4c1df45ad8d7d",
+            "max": 872,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_30b1448d901a4191ab5bd62e89028e75",
+            "value": 872
+          }
+        },
+        "66071ddc16014e9a8604f37fe1441b33": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "670f7df971d14fabbffcfd2535e90005": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_93c4b10a8f074c2c95cee65850ab2040",
+            "placeholder": "​",
+            "style": "IPY_MODEL_275c1f7417e247ce99519f0ef43738c9",
+            "value": " 0/1821 [00:00&lt;?, ? examples/s]"
+          }
+        },
+        "69a3fcf1ec174cbaa11f93fa1dded564": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "6bed029127e44fefa0ccd576edb855e9": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_289e60ba017345b1b61adeecc24bc2cb",
+              "IPY_MODEL_bb9d3ced705148b6a8c3b71498bd9d81",
+              "IPY_MODEL_4579a6614cd64cdcab0ef6c615cad8bf"
+            ],
+            "layout": "IPY_MODEL_5b7a72d055cd4cb68a4a9706bd276e1f"
+          }
+        },
+        "6dc9fbdebed24e35bccaf0e92a6a491e": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_b0999cce50384ccfa1fb6dedcf6bc1d3",
+            "placeholder": "​",
+            "style": "IPY_MODEL_7f65a8e8ab6a411286d11c0125b316d0",
+            "value": " 52039/67349 [00:00&lt;00:00, 292448.11 examples/s]"
+          }
+        },
+        "6e2dd84de8b3483e9f18142fa30177eb": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_69a3fcf1ec174cbaa11f93fa1dded564",
+            "max": 67349,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_2522eb90704c4fd6be3a0223dd16442e",
+            "value": 67349
+          }
+        },
+        "6f7051dbe4ad467bada9e50157141145": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "71a33b25dece4caca90a91160ea2382f": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_c84704566a0a45249c6973c8e57e3856",
+              "IPY_MODEL_44ff2d0794b74097af53a8116da941a7",
+              "IPY_MODEL_b22d3db2d68046a1afbca22b13ab7152"
+            ],
+            "layout": "IPY_MODEL_acbe1228b9bb42da9fc19e57e4dc6e74"
+          }
+        },
+        "75083b32cd854cdd9de6490b7f828140": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "758811c0d8284c76984fc2ac974d835a": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "77c0fa586d4d4dcd982513c827799adb": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "78d934805a364c8989efba79bead117d": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "7c9acf3891a040f5af892b86e72c803c": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "7dac43c361074f0f998d8c150aa922b9": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_02ecb50a721f450b940172d3d35a4abe",
+              "IPY_MODEL_3977c5ac1af645b0b2c1e5886ac037ad",
+              "IPY_MODEL_e9a881d7fcb84bafb6a063009bf4dd54"
+            ],
+            "layout": "IPY_MODEL_efdf1618f9224961b840bf58fb763876"
+          }
+        },
+        "7f65a8e8ab6a411286d11c0125b316d0": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "80053cb69a2a4b2c8ad3c038fc3e2c39": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "8388b866783e4d9eb11093a99eaa3116": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_e5d9c734fb3e412aabc03e5790b12382",
+            "placeholder": "​",
+            "style": "IPY_MODEL_07fd35fa2e634720b064998b96e0e60e",
+            "value": "Shuffling ~/tensorflow_datasets/glue/sst2/2.0.0.incompleteQH5CX4/glue-validation.tfrecord*...:   0%"
+          }
+        },
+        "879fa9af242b4c698c15747315a91ba9": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_27d5831d71b547c081cb2dec9a42b300",
+            "placeholder": "​",
+            "style": "IPY_MODEL_14df05ca9bc642d9ad588db644ab7cf7",
+            "value": "Generating splits...: 100%"
+          }
+        },
+        "87e1ae4c68394e23ba0ed21f8fffb5a6": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_4679a29c8e2c4fa686713038e1c4686e",
+            "max": 1821,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_b9d49ffe543342c694905d2dee72b880",
+            "value": 1821
+          }
+        },
+        "88b26fe32ba64f9e90c10b6c40381635": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "8a0f69f4dc4040559001183e68b73aff": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "8bf597ef7f7a4ce1a66cbfc1aa74a216": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "8c3528b24aba4002b71168fd0b8beb44": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "8fe22f392c204fdeb8c8b34b9c4bdc12": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "90e15c8359274c5a85a7cf427cd9ca22": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "91a8aba5cca943fcb2f3f3353450aa78": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_42117d78043a4336bc9c90c88741c15a",
+            "placeholder": "​",
+            "style": "IPY_MODEL_b28ba82c231545e0b509a93ba0470449",
+            "value": "Generating train examples...: 100%"
+          }
+        },
+        "93c4b10a8f074c2c95cee65850ab2040": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "964b0fedfdbe4b4a990d3b72d68b4112": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": "20px"
+          }
+        },
+        "96d397cf27714dca8a8545c3ceb13d73": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_91a8aba5cca943fcb2f3f3353450aa78",
+              "IPY_MODEL_10e9f95c33dd4b7b88150d6a7825e973",
+              "IPY_MODEL_d929b9d20b1a4b00bb5902c0883b4034"
+            ],
+            "layout": "IPY_MODEL_c1976cfcd5c4438bb769ad5a77b0b265"
+          }
+        },
+        "97717222d177460bad6c626004f3dd9d": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "9b65c0c17d3d49589feed767a4ced04c": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_78d934805a364c8989efba79bead117d",
+            "placeholder": "​",
+            "style": "IPY_MODEL_5dd04d5d68a54166b432bdf53dad0792",
+            "value": " 3/3 [00:09&lt;00:00,  2.22s/ splits]"
+          }
+        },
+        "9eb1c4b7b7e94f8599705b2d1d509497": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "a0d66c335f174159a7b2df88cb4b607f": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_58dfc96028a948c7af85f4df518bc171",
+              "IPY_MODEL_436d985f48104913933079fdb44fde6f",
+              "IPY_MODEL_aca24255c81943aeb92a036311e53cef"
+            ],
+            "layout": "IPY_MODEL_c7a31a4a0e4d4b27a93e48950acfe5e5"
+          }
+        },
+        "a70a3e92a37247eeaad91ef1e140aa54": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "a7fb5284dab64287b108a848c7203dd5": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": "hidden",
+            "width": null
+          }
+        },
+        "aca24255c81943aeb92a036311e53cef": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_80053cb69a2a4b2c8ad3c038fc3e2c39",
+            "placeholder": "​",
+            "style": "IPY_MODEL_7c9acf3891a040f5af892b86e72c803c",
+            "value": " 1402/1821 [00:00&lt;00:00, 7169.58 examples/s]"
+          }
+        },
+        "acbe1228b9bb42da9fc19e57e4dc6e74": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "b0999cce50384ccfa1fb6dedcf6bc1d3": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "b22d3db2d68046a1afbca22b13ab7152": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_b96bbee1b3254bc397828295b8b6528a",
+            "placeholder": "​",
+            "style": "IPY_MODEL_4d7f2052da7f43f2a41ebfaaf562ae77",
+            "value": " 1/1 [00:01&lt;00:00,  1.28s/ file]"
+          }
+        },
+        "b28ba82c231545e0b509a93ba0470449": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "b3086aaf2d8b491fbb5569d6226bfd29": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "b4a60897f73642eeb692cafa7997167c": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_c4d8f2f2b9604745ac03fe5f4da7eeee",
+              "IPY_MODEL_87e1ae4c68394e23ba0ed21f8fffb5a6",
+              "IPY_MODEL_670f7df971d14fabbffcfd2535e90005"
+            ],
+            "layout": "IPY_MODEL_148f6232f6db44919f04a78c70aba8c1"
+          }
+        },
+        "b65aff687d4242fa94ede609a48ad847": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "b96bbee1b3254bc397828295b8b6528a": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "b9d49ffe543342c694905d2dee72b880": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "ProgressStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "ba64c4461ad14990afe4c1df45ad8d7d": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "bb9d3ced705148b6a8c3b71498bd9d81": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "success",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_078e2f25095d40cc9a96ad3488030022",
+            "max": 1,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_16dc3ec494f6456794b43d33a95ef042",
+            "value": 1
+          }
+        },
+        "c1976cfcd5c4438bb769ad5a77b0b265": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": "hidden",
+            "width": null
+          }
+        },
+        "c4d8f2f2b9604745ac03fe5f4da7eeee": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_66071ddc16014e9a8604f37fe1441b33",
+            "placeholder": "​",
+            "style": "IPY_MODEL_88b26fe32ba64f9e90c10b6c40381635",
+            "value": "Shuffling ~/tensorflow_datasets/glue/sst2/2.0.0.incompleteQH5CX4/glue-test.tfrecord*...:   0%"
+          }
+        },
+        "c7a31a4a0e4d4b27a93e48950acfe5e5": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": "hidden",
+            "width": null
+          }
+        },
+        "c84704566a0a45249c6973c8e57e3856": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_06fa78192fd74020962d609ac59af6d5",
+            "placeholder": "​",
+            "style": "IPY_MODEL_4de362ac37224f968d2f6a4f2336b327",
+            "value": "Extraction completed...: 100%"
+          }
+        },
+        "c96bafeb19454961aa95b819bec135df": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_879fa9af242b4c698c15747315a91ba9",
+              "IPY_MODEL_cda35a02449442e48f3b6131547e1628",
+              "IPY_MODEL_9b65c0c17d3d49589feed767a4ced04c"
+            ],
+            "layout": "IPY_MODEL_09d1f164a0e941438f06a3feb58bdd30"
+          }
+        },
+        "cb1357c15a4c4a28a6198a61af82848b": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_e70cb8396ab341c1bb0d60d2ebdd708a",
+            "placeholder": "​",
+            "style": "IPY_MODEL_d5020ccf7aff400988f0f0e6d1ff3557",
+            "value": "Generating validation examples...:  73%"
+          }
+        },
+        "cbd21b3e1e5f4e318940805ac4d142ec": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "cd67f31d575f44d9944a5019b3e38c19": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HBoxModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_cb1357c15a4c4a28a6198a61af82848b",
+              "IPY_MODEL_0a9d89c682624b208b399ea9d8c26b0b",
+              "IPY_MODEL_130007f46b794a4c9d283aa5fb333d04"
+            ],
+            "layout": "IPY_MODEL_a7fb5284dab64287b108a848c7203dd5"
+          }
+        },
+        "cda35a02449442e48f3b6131547e1628": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "FloatProgressModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_8fe22f392c204fdeb8c8b34b9c4bdc12",
+            "max": 3,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_97717222d177460bad6c626004f3dd9d",
+            "value": 3
+          }
+        },
+        "d147c549823e407eb843750f0ec45cbc": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": "hidden",
+            "width": null
+          }
+        },
+        "d410bad163c24c1aa5e00d54090b05db": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "d5020ccf7aff400988f0f0e6d1ff3557": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "d929b9d20b1a4b00bb5902c0883b4034": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_e9113eab130140f9b9367a61b1e4b5c0",
+            "placeholder": "​",
+            "style": "IPY_MODEL_454dfc6558ae4a1b83621ea1352fc051",
+            "value": " 67015/67349 [00:08&lt;00:00, 8511.23 examples/s]"
+          }
+        },
+        "d9a622959c0d4c45b0b1932d7c85893d": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_758811c0d8284c76984fc2ac974d835a",
+            "placeholder": "​",
+            "style": "IPY_MODEL_4e06b19ebe63441ebb6a9b1cd57e2d21",
+            "value": "Shuffling ~/tensorflow_datasets/glue/sst2/2.0.0.incompleteQH5CX4/glue-train.tfrecord*...:  77%"
+          }
+        },
+        "d9d4d67205d340a8acb919b7c80a842a": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "DescriptionStyleModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "e5d9c734fb3e412aabc03e5790b12382": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "e70cb8396ab341c1bb0d60d2ebdd708a": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "e9113eab130140f9b9367a61b1e4b5c0": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "e9a881d7fcb84bafb6a063009bf4dd54": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_77c0fa586d4d4dcd982513c827799adb",
+            "placeholder": "​",
+            "style": "IPY_MODEL_a70a3e92a37247eeaad91ef1e140aa54",
+            "value": " 1/1 [00:01&lt;00:00,  1.11s/ url]"
+          }
+        },
+        "efdf1618f9224961b840bf58fb763876": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "fac49204e83e49c9920c29a61382d9b7": {
+          "model_module": "@jupyter-widgets/base",
+          "model_module_version": "1.2.0",
+          "model_name": "LayoutModel",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": "20px"
+          }
+        }
+      }
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}


Test patch:
diff --git a/keras_nlp/models/bert/bert_models_test.py b/keras_nlp/models/bert/bert_models_test.py
index 57f30b08..8a39cb17 100644
--- a/keras_nlp/models/bert/bert_models_test.py
+++ b/keras_nlp/models/bert/bert_models_test.py
@@ -69,59 +69,6 @@ class BertTest(tf.test.TestCase, parameterized.TestCase):
             }
             self.model(input_data)
 
-    def test_valid_call_presets(self):
-        # Test preset loading without weights
-        for preset in Bert.presets:
-            model = Bert.from_preset(preset, load_weights=False, name="encoder")
-            input_data = {
-                "token_ids": tf.ones(
-                    (self.batch_size, self.model.max_sequence_length),
-                    dtype="int32",
-                ),
-                "segment_ids": tf.ones(
-                    (self.batch_size, self.model.max_sequence_length),
-                    dtype="int32",
-                ),
-                "padding_mask": tf.ones(
-                    (self.batch_size, self.model.max_sequence_length),
-                    dtype="int32",
-                ),
-            }
-            model(input_data)
-
-    def test_unknown_preset_error(self):
-        # Not a preset name
-        with self.assertRaises(ValueError):
-            Bert.from_preset(
-                "bert_base_uncased_clowntown",
-                load_weights=False,
-            )
-
-    def test_preset_mutability(self):
-        preset = "bert_base_uncased_en"
-        # Cannot overwrite the presents attribute in an object
-        with self.assertRaises(AttributeError):
-            self.model.presets = {"my_model": "clowntown"}
-        # Cannot mutate presents in an object
-        config = self.model.presets[preset]["config"]
-        config["max_sequence_length"] = 1
-        self.assertEqual(config["max_sequence_length"], 1)
-        self.assertEqual(
-            self.model.presets[preset]["config"]["max_sequence_length"], 512
-        )
-        # Cannot mutate presets in the class
-        config = Bert.presets[preset]["config"]
-        config["max_sequence_length"] = 1
-        self.assertEqual(config["max_sequence_length"], 1)
-        self.assertEqual(
-            Bert.presets[preset]["config"]["max_sequence_length"], 512
-        )
-
-    def test_preset_docstring(self):
-        """Check we did our docstring formatting correctly."""
-        for name in Bert.presets:
-            self.assertRegex(Bert.from_preset.__doc__, name)
-
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
diff --git a/keras_nlp/models/bert/bert_preprocessing_test.py b/keras_nlp/models/bert/bert_preprocessing_test.py
index bf6ce9df..4f7ea449 100644
--- a/keras_nlp/models/bert/bert_preprocessing_test.py
+++ b/keras_nlp/models/bert/bert_preprocessing_test.py
@@ -14,7 +14,6 @@
 """Tests for BERT preprocessing layers."""
 
 import os
-import unittest
 
 import tensorflow as tf
 from absl.testing import parameterized
@@ -58,26 +57,6 @@ class BertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
         tokenizer = BertTokenizer(vocabulary=self.vocab)
         self.assertEqual(tokenizer.vocabulary_size(), 13)
 
-    def test_unknown_preset_error(self):
-        # Not a preset name
-        with self.assertRaises(ValueError):
-            BertPreprocessor.from_preset("bert_base_uncased_clowntown")
-
-    def test_preset_docstring(self):
-        """Check we did our docstring formatting correctly."""
-        for name in BertPreprocessor.presets:
-            self.assertRegex(BertPreprocessor.from_preset.__doc__, name)
-
-    @unittest.mock.patch("tensorflow.keras.utils.get_file")
-    def test_valid_call_presets(self, get_file_mock):
-        """Ensure presets have necessary structure, but no RPCs."""
-        input_data = ["THE QUICK BROWN FOX."]
-        get_file_mock.return_value = self.vocab
-        for preset in BertTokenizer.presets:
-            tokenizer = BertTokenizer.from_preset(preset)
-            tokenizer(input_data)
-        self.assertEqual(get_file_mock.call_count, len(BertTokenizer.presets))
-
     @parameterized.named_parameters(
         ("save_format_tf", "tf"), ("save_format_h5", "h5")
     )
@@ -114,14 +93,7 @@ class BertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(output["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 0])
 
     def test_tokenize_batch(self):
-        input_data = tf.constant(
-            [
-                "THE QUICK BROWN FOX.",
-                "THE QUICK BROWN FOX.",
-                "THE QUICK BROWN FOX.",
-                "THE QUICK BROWN FOX.",
-            ]
-        )
+        input_data = tf.constant(["THE QUICK BROWN FOX."] * 4)
         preprocessor = BertPreprocessor(
             BertTokenizer(vocabulary=self.vocab),
             sequence_length=8,
@@ -135,6 +107,46 @@ class BertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
             output["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 0]] * 4
         )
 
+    def test_tokenize_labeled_batch(self):
+        x = tf.constant(["THE QUICK BROWN FOX."] * 4)
+        y = tf.constant([1] * 4)
+        sw = tf.constant([1.0] * 4)
+        preprocessor = BertPreprocessor(
+            BertTokenizer(vocabulary=self.vocab),
+            sequence_length=8,
+        )
+        x_out, y_out, sw_out = preprocessor(x, y, sw)
+        self.assertAllEqual(x_out["token_ids"], [[2, 5, 6, 7, 8, 1, 3, 0]] * 4)
+        self.assertAllEqual(
+            x_out["segment_ids"], [[0, 0, 0, 0, 0, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(
+            x_out["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 0]] * 4
+        )
+        self.assertAllEqual(y_out, y)
+        self.assertAllEqual(sw_out, sw)
+
+    def test_tokenize_labeled_dataset(self):
+        x = tf.constant(["THE QUICK BROWN FOX."] * 4)
+        y = tf.constant([1] * 4)
+        sw = tf.constant([1.0] * 4)
+        ds = tf.data.Dataset.from_tensor_slices((x, y, sw))
+        preprocessor = BertPreprocessor(
+            BertTokenizer(vocabulary=self.vocab),
+            sequence_length=8,
+        )
+        ds = ds.map(preprocessor)
+        x_out, y_out, sw_out = ds.batch(4).take(1).get_single_element()
+        self.assertAllEqual(x_out["token_ids"], [[2, 5, 6, 7, 8, 1, 3, 0]] * 4)
+        self.assertAllEqual(
+            x_out["segment_ids"], [[0, 0, 0, 0, 0, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(
+            x_out["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 0]] * 4
+        )
+        self.assertAllEqual(y_out, y)
+        self.assertAllEqual(sw_out, sw)
+
     def test_tokenize_multiple_sentences(self):
         sentence_one = "THE QUICK"
         sentence_two = "BROWN FOX."
@@ -150,22 +162,8 @@ class BertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(output["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 1])
 
     def test_tokenize_multiple_batched_sentences(self):
-        sentence_one = tf.constant(
-            [
-                "THE QUICK",
-                "THE QUICK",
-                "THE QUICK",
-                "THE QUICK",
-            ]
-        )
-        sentence_two = tf.constant(
-            [
-                "BROWN FOX.",
-                "BROWN FOX.",
-                "BROWN FOX.",
-                "BROWN FOX.",
-            ]
-        )
+        sentence_one = tf.constant(["THE QUICK"] * 4)
+        sentence_two = tf.constant(["BROWN FOX."] * 4)
         preprocessor = BertPreprocessor(
             BertTokenizer(vocabulary=self.vocab),
             sequence_length=8,
@@ -181,50 +179,6 @@ class BertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
             output["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 1]] * 4
         )
 
-    def test_unknown_preset_error(self):
-        # Not a preset name
-        with self.assertRaises(ValueError):
-            BertPreprocessor.from_preset("bert_base_uncased_clowntown")
-
-    def test_preset_docstring(self):
-        """Check we did our docstring formatting correctly."""
-        for name in BertPreprocessor.presets:
-            self.assertRegex(BertPreprocessor.from_preset.__doc__, name)
-
-    @unittest.mock.patch("tensorflow.keras.utils.get_file")
-    def test_valid_call_presets(self, get_file_mock):
-        """Ensure presets have necessary structure, but no RPCs."""
-        input_data = ["THE QUICK BROWN FOX."]
-        get_file_mock.return_value = self.vocab
-        for preset in BertPreprocessor.presets:
-            preprocessor = BertPreprocessor.from_preset(preset)
-            preprocessor(input_data)
-        self.assertEqual(
-            get_file_mock.call_count, len(BertPreprocessor.presets)
-        )
-
-    @unittest.mock.patch("tensorflow.keras.utils.get_file")
-    def test_override_preprocessor_sequence_length(self, get_file_mock):
-        get_file_mock.return_value = self.vocab
-        preprocessor = BertPreprocessor.from_preset(
-            "bert_base_uncased_en",
-            sequence_length=64,
-        )
-        self.assertEqual(preprocessor.get_config()["sequence_length"], 64)
-        preprocessor("The quick brown fox.")
-        get_file_mock.assert_called_once()
-
-    @unittest.mock.patch("tensorflow.keras.utils.get_file")
-    def test_override_preprocessor_sequence_length_gt_max(self, get_file_mock):
-        """Override sequence length longer than model's maximum."""
-        get_file_mock.return_value = self.vocab
-        with self.assertRaises(ValueError):
-            BertPreprocessor.from_preset(
-                "bert_base_uncased_en",
-                sequence_length=1024,
-            )
-        get_file_mock.assert_called_once()
-
     @parameterized.named_parameters(
         ("save_format_tf", "tf"), ("save_format_h5", "h5")
     )
diff --git a/keras_nlp/models/bert/bert_presets_test.py b/keras_nlp/models/bert/bert_presets_test.py
index efedd0f3..167d7222 100644
--- a/keras_nlp/models/bert/bert_presets_test.py
+++ b/keras_nlp/models/bert/bert_presets_test.py
@@ -15,6 +15,7 @@
 
 import pytest
 import tensorflow as tf
+from absl.testing import parameterized
 
 from keras_nlp.models.bert.bert_models import Bert
 from keras_nlp.models.bert.bert_preprocessing import BertPreprocessor
@@ -23,7 +24,7 @@ from keras_nlp.models.bert.bert_tasks import BertClassifier
 
 
 @pytest.mark.large
-class BertPresetSmokeTest(tf.test.TestCase):
+class BertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
     """
     A smoke test for BERT presets we run continuously.
 
@@ -48,49 +49,124 @@ class BertPresetSmokeTest(tf.test.TestCase):
         expected_outputs = [101, 1996, 4248, 102]
         self.assertAllEqual(outputs, expected_outputs)
 
-    def test_backbone_output(self):
+    @parameterized.named_parameters(
+        ("load_weights", True), ("no_load_weights", False)
+    )
+    def test_backbone_output(self, load_weights):
         input_data = {
             "token_ids": tf.constant([[101, 1996, 4248, 102]]),
             "segment_ids": tf.constant([[0, 0, 0, 0]]),
             "padding_mask": tf.constant([[1, 1, 1, 1]]),
         }
         model = Bert.from_preset(
-            "bert_tiny_uncased_en",
+            "bert_tiny_uncased_en", load_weights=load_weights
         )
-        outputs = model(input_data)["sequence_output"][0, 0, :5]
-        # Outputs from our preset checkpoints should be stable!
-        # We should only update these numbers if we are updating a weights file,
-        # or have found a bug where output did not match the upstream source.
-        expected_outputs = [-1.38173, 0.16598, -2.92788, -2.66958, -0.61556]
-        # Keep a high tolerance, so we are robust to different hardware.
-        self.assertAllClose(outputs, expected_outputs, atol=0.01, rtol=0.01)
-
-    def test_classifier_output(self):
+        outputs = model(input_data)["sequence_output"]
+        if load_weights:
+            # The forward pass from a preset should be stable!
+            # This test should catch cases where we unintentionally change our
+            # network code in a way that would invalidate our preset weights.
+            # We should only update these numbers if we are updating a weights
+            # file, or have found a discrepancy with the upstream source.
+            outputs = outputs[0, 0, :5]
+            expected = [-1.38173, 0.16598, -2.92788, -2.66958, -0.61556]
+            # Keep a high tolerance, so we are robust to different hardware.
+            self.assertAllClose(outputs, expected, atol=0.01, rtol=0.01)
+
+    @parameterized.named_parameters(
+        ("load_weights", True), ("no_load_weights", False)
+    )
+    def test_classifier_output(self, load_weights):
         input_data = {
             "token_ids": tf.constant([[101, 1996, 4248, 102]]),
             "segment_ids": tf.constant([[0, 0, 0, 0]]),
             "padding_mask": tf.constant([[1, 1, 1, 1]]),
         }
         model = BertClassifier.from_preset(
-            "bert_tiny_uncased_en",
+            "bert_tiny_uncased_en", load_weights=load_weights
         )
-        # We don't assert output values, as the head weights are random.
+        # Never assert output values, as the head weights are random.
         model(input_data)
 
+    @parameterized.named_parameters(
+        ("bert_tokenizer", BertTokenizer),
+        ("bert_preprocessor", BertPreprocessor),
+        ("bert", Bert),
+        ("bert_classifier", BertClassifier),
+    )
+    def test_preset_mutability(self, cls):
+        preset = "bert_tiny_uncased_en"
+        obj = cls.from_preset(preset)
+        # Cannot overwrite the presents attribute in an object
+        with self.assertRaises(AttributeError):
+            obj.presets = {"my_model": "clowntown"}
+        # Cannot mutate presents in an object
+        config = obj.presets[preset]["config"]
+        config["num_layers"] = 1
+        self.assertEqual(config["num_layers"], 1)
+        self.assertEqual(obj.presets[preset]["config"]["num_layers"], 2)
+        # Cannot mutate presets in the class
+        config = Bert.presets[preset]["config"]
+        config["num_layers"] = 1
+        self.assertEqual(config["num_layers"], 1)
+        self.assertEqual(Bert.presets[preset]["config"]["num_layers"], 2)
+
+    @parameterized.named_parameters(
+        ("bert_tokenizer", BertTokenizer),
+        ("bert_preprocessor", BertPreprocessor),
+        ("bert", Bert),
+        ("bert_classifier", BertClassifier),
+    )
+    def test_preset_docstring(self, cls):
+        """Check we did our docstring formatting correctly."""
+        for name in cls.presets:
+            self.assertRegex(cls.from_preset.__doc__, name)
+
+    @parameterized.named_parameters(
+        ("bert_tokenizer", BertTokenizer),
+        ("bert_preprocessor", BertPreprocessor),
+        ("bert", Bert),
+        ("bert_classifier", BertClassifier),
+    )
+    def test_unknown_preset_error(self, cls):
+        # Not a preset name
+        with self.assertRaises(ValueError):
+            cls.from_preset("bert_base_uncased_clowntown")
+
+    def test_override_preprocessor_sequence_length(self):
+        """Override sequence length longer than model's maximum."""
+        preprocessor = BertPreprocessor.from_preset(
+            "bert_base_uncased_en",
+            sequence_length=64,
+        )
+        self.assertEqual(preprocessor.get_config()["sequence_length"], 64)
+        preprocessor("The quick brown fox.")
+
+    def test_override_preprocessor_sequence_length_gt_max(self):
+        """Override sequence length longer than model's maximum."""
+        with self.assertRaises(ValueError):
+            BertPreprocessor.from_preset(
+                "bert_base_uncased_en",
+                sequence_length=1024,
+            )
+
 
 @pytest.mark.extra_large
-class BertPresetTest(tf.test.TestCase):
+class BertPresetFullTest(tf.test.TestCase, parameterized.TestCase):
     """
     Test the full enumeration of our preset.
 
-    This only tests all enumeration of our presets and is only run manually.
+    This every presets for BERT and is only run manually.
     Run with:
     `pytest keras_nlp/models/bert_presets_test.py --run_extra_large`
     """
 
-    def test_load_bert(self):
+    @parameterized.named_parameters(
+        ("load_weights", True), ("no_load_weights", False)
+    )
+    def test_load_bert(self, load_weights):
         for preset in Bert.presets:
-            model = Bert.from_preset(preset, load_weights=True)
+            model = Bert.from_preset(preset, load_weights=load_weights)
             input_data = {
                 "token_ids": tf.random.uniform(
                     shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
@@ -102,10 +178,13 @@ class BertPresetTest(tf.test.TestCase):
             }
             model(input_data)
 
-    def test_load_bert_classifier(self):
+    @parameterized.named_parameters(
+        ("load_weights", True), ("no_load_weights", False)
+    )
+    def test_load_bert_classifier(self, load_weights):
         for preset in BertClassifier.presets:
             classifier = BertClassifier.from_preset(
-                preset, num_classes=4, name="classifier"
+                preset, load_weights=load_weights
             )
             input_data = {
                 "token_ids": tf.random.uniform(
diff --git a/keras_nlp/models/bert/bert_tasks_test.py b/keras_nlp/models/bert/bert_tasks_test.py
index e5a4f3f8..13dcfc06 100644
--- a/keras_nlp/models/bert/bert_tasks_test.py
+++ b/keras_nlp/models/bert/bert_tasks_test.py
@@ -58,39 +58,6 @@ class BertClassifierTest(tf.test.TestCase, parameterized.TestCase):
     def test_valid_call_classifier(self):
         self.classifier(self.input_batch)
 
-    def test_valid_call_presets(self):
-        # Test preset loading without weights
-        for preset in BertClassifier.presets:
-            classifier = BertClassifier.from_preset(preset, load_weights=False)
-            input_data = {
-                "token_ids": tf.ones(
-                    (self.batch_size, self.backbone.max_sequence_length),
-                    dtype="int32",
-                ),
-                "segment_ids": tf.ones(
-                    (self.batch_size, self.backbone.max_sequence_length),
-                    dtype="int32",
-                ),
-                "padding_mask": tf.ones(
-                    (self.batch_size, self.backbone.max_sequence_length),
-                    dtype="int32",
-                ),
-            }
-            classifier(input_data)
-
-    def test_unknown_preset_error(self):
-        # Not a preset name
-        with self.assertRaises(ValueError):
-            BertClassifier.from_preset(
-                "bert_base_uncased_clowntown",
-                load_weights=False,
-            )
-
-    def test_preset_docstring(self):
-        """Check we did our docstring formatting correctly."""
-        for name in BertClassifier.presets:
-            self.assertRegex(BertClassifier.from_preset.__doc__, name)
-
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
diff --git a/keras_nlp/models/gpt2/gpt2_preprocessing_test.py b/keras_nlp/models/gpt2/gpt2_preprocessing_test.py
index 7dbf79c0..73f758ee 100644
--- a/keras_nlp/models/gpt2/gpt2_preprocessing_test.py
+++ b/keras_nlp/models/gpt2/gpt2_preprocessing_test.py
@@ -25,7 +25,7 @@ from keras_nlp.models.gpt2.gpt2_preprocessing import GPT2Tokenizer
 
 class GPT2TokenizerTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        vocab = {
+        self.vocab = {
             "
": 0,
             "Ġair": 1,
             "plane": 2,
@@ -42,8 +42,11 @@ class GPT2TokenizerTest(tf.test.TestCase, parameterized.TestCase):
         merges += ["Ġa t", "p o", "r t", "o h", "l i", "Ġi s", "Ġb e", "s t"]
         merges += ["Ġt h", "Ġai r", "pl a", "Ġk oh", "Ġth e", "Ġbe st", "po rt"]
         merges += ["pla ne"]
+        self.merges = merges
 
-        self.tokenizer = GPT2Tokenizer(vocabulary=vocab, merges=merges)
+        self.tokenizer = GPT2Tokenizer(
+            vocabulary=self.vocab, merges=self.merges
+        )
 
     def test_tokenize(self):
         input_data = " airplane at airport"
diff --git a/keras_nlp/models/gpt2/gpt2_presets_test.py b/keras_nlp/models/gpt2/gpt2_presets_test.py
new file mode 100644
index 00000000..428cb52b
--- /dev/null
+++ b/keras_nlp/models/gpt2/gpt2_presets_test.py
@@ -0,0 +1,107 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for loading pretrained model presets."""
+
+import pytest
+import tensorflow as tf
+from absl.testing import parameterized
+
+from keras_nlp.models.gpt2.gpt2_models import GPT2
+from keras_nlp.models.gpt2.gpt2_preprocessing import GPT2Tokenizer
+
+
+@pytest.mark.large
+class GPT2PresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
+    """
+    A smoke test for GPT-2 presets we run continuously.
+
+    This only tests the smallest weights we have available. Run with:
+    `pytest keras_nlp/models/gpt2/gpt2_presets_test.py --run_large`
+    """
+
+    def test_tokenizer_output(self):
+        tokenizer = GPT2Tokenizer.from_preset("gpt2_base")
+        outputs = tokenizer("The quick brown fox.")
+        expected_outputs = [464, 2068, 7586, 21831, 13]
+        self.assertAllEqual(outputs, expected_outputs)
+
+    @parameterized.named_parameters(
+        ("preset_weights", True), ("random_weights", False)
+    )
+    def test_backbone_output(self, load_weights):
+        input_data = {
+            "token_ids": tf.constant([[1169, 2068, 7586, 21831, 13]]),
+            "padding_mask": tf.constant([[1, 1, 1, 1, 1]]),
+        }
+        model = GPT2.from_preset("gpt2_base", load_weights=load_weights)
+        outputs = model(input_data)[0, 0, :5]
+        if load_weights:
+            # The forward pass from a preset should be stable!
+            # This test should catch cases where we unintentionally change our
+            # network code in a way that would invalidate our preset weights.
+            # We should only update these numbers if we are updating a weights
+            # file, or have found a discrepancy with the upstream source.
+            expected_outputs = [-0.1116, -0.0375, -0.2624, 0.00891, -0.0061]
+            # Keep a high tolerance, so we are robust to different hardware.
+            self.assertAllClose(outputs, expected_outputs, atol=0.01, rtol=0.01)
+
+    @parameterized.named_parameters(
+        ("gpt2_tokenizer", GPT2Tokenizer),
+        ("gpt2", GPT2),
+    )
+    def test_preset_docstring(self, cls):
+        """Check we did our docstring formatting correctly."""
+        for name in cls.presets:
+            self.assertRegex(cls.from_preset.__doc__, name)
+
+    @parameterized.named_parameters(
+        ("gpt2_tokenizer", GPT2Tokenizer),
+        ("gpt2", GPT2),
+    )
+    def test_unknown_preset_error(self, cls):
+        # Not a preset name
+        with self.assertRaises(ValueError):
+            cls.from_preset("gpt2_base_clowntown")
+
+
+@pytest.mark.extra_large
+class GPT2PresetFullTest(tf.test.TestCase, parameterized.TestCase):
+    """
+    Test the full enumeration of our preset.
+
+    This tests every GPT-2 preset and is only run manually.
+    Run with:
+    `pytest keras_nlp/models/gpt2/gpt2_presets_test.py --run_extra_large`
+    """
+
+    @parameterized.named_parameters(
+        ("preset_weights", True), ("random_weights", False)
+    )
+    def test_load_gpt2(self, load_weights):
+        for preset in GPT2.presets:
+            model = GPT2.from_preset(preset, load_weights=load_weights)
+            input_data = {
+                "token_ids": tf.random.uniform(
+                    shape=(1, 1024),
+                    dtype=tf.int64,
+                    maxval=model.vocabulary_size,
+                ),
+                "padding_mask": tf.constant([1] * 1024, shape=(1, 1024)),
+            }
+            model(input_data)
+
+    def test_load_tokenizers(self):
+        for preset in GPT2Tokenizer.presets:
+            tokenizer = GPT2Tokenizer.from_preset(preset)
+            tokenizer("The quick brown fox.")
diff --git a/keras_nlp/utils/keras_utils_test.py b/keras_nlp/utils/keras_utils_test.py
index 498d7d5e..2e4f0a18 100644
--- a/keras_nlp/utils/keras_utils_test.py
+++ b/keras_nlp/utils/keras_utils_test.py
@@ -16,6 +16,7 @@ import tensorflow as tf
 from tensorflow import keras
 
 from keras_nlp.utils.keras_utils import clone_initializer
+from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
 
 
 class CloneInitializerTest(tf.test.TestCase):
@@ -39,3 +40,28 @@ class CloneInitializerTest(tf.test.TestCase):
         initializer = "glorot_uniform"
         clone = clone_initializer(initializer)
         self.assertAllEqual(initializer, clone)
+
+
+class PackTest(tf.test.TestCase):
+    def test_pack_dict(self):
+        tensor_dict = {"foo": tf.constant([1, 2])}
+        data = pack_x_y_sample_weight(tensor_dict)
+        self.assertAllEqual(data, tensor_dict)
+
+    def test_pack_tuple(self):
+        tensor_tuple = (tf.constant([1, 2]),)
+        data = pack_x_y_sample_weight(tensor_tuple)
+        self.assertAllEqual(data, (tensor_tuple,))
+
+    def test_pack_pair(self):
+        x = tf.constant([1, 2])
+        y = tf.constant([3, 4])
+        data = pack_x_y_sample_weight(x, y)
+        self.assertAllEqual(data, (x, y))
+
+    def test_pack_triplet(self):
+        x = tf.constant([1, 2])
+        y = tf.constant([3, 4])
+        sw = tf.constant([5, 6])
+        data = pack_x_y_sample_weight(x, y, sw)
+        self.assertAllEqual(data, (x, y, sw))
diff --git a/keras_nlp/utils/pipeline_model_test.py b/keras_nlp/utils/pipeline_model_test.py
new file mode 100644
index 00000000..5b52239d
--- /dev/null
+++ b/keras_nlp/utils/pipeline_model_test.py
@@ -0,0 +1,559 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.utils.pipeline_model import PipelineModel
+
+
+class NoopPipeline(PipelineModel):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.dense = keras.layers.Dense(1)
+
+    def call(self, inputs):
+        return self.dense(inputs)
+
+
+class FeaturePipeline(PipelineModel):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.dense = keras.layers.Dense(1)
+
+    def preprocess_features(self, x):
+        return tf.strings.to_number(x)
+
+    def call(self, inputs):
+        return self.dense(inputs)
+
+
+class LabelPipeline(PipelineModel):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.dense = keras.layers.Dense(1)
+
+    def preprocess_labels(self, y):
+        return tf.strings.to_number(y)
+
+    def call(self, inputs):
+        return self.dense(inputs)
+
+
+class DataPipeline(PipelineModel):
+    """This model generates labels straight from the input data."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.dense = keras.layers.Dense(1)
+
+    def preprocess_samples(self, x, y=None, sample_weight=None):
+        return tf.strings.to_number(x), tf.strings.to_number(x), sample_weight
+
+    def call(self, inputs):
+        return self.dense(inputs)
+
+
+class FunctionalPipeline(PipelineModel):
+    def __init__(self, **kwargs):
+        inputs = keras.Input(shape=(5,))
+        outputs = keras.layers.Dense(1)(inputs)
+        super().__init__(inputs, outputs, **kwargs)
+
+    def preprocess_features(self, x):
+        return tf.strings.to_number(x)
+
+    def get_config(self):
+        return {}
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
+
+
+class TestNoopPipelineModel(tf.test.TestCase, parameterized.TestCase):
+    def test_fit(self):
+        x = tf.random.uniform((8, 5))
+        y = tf.random.uniform((8, 1))
+        sw = tf.random.uniform((8, 1))
+        model = NoopPipeline()
+        model.compile(loss="mse")
+        # With sample weight.
+        model.fit(x=x, y=y, sample_weight=sw, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
+        # Without sample weight.
+        model.fit(x=x, y=y, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    def test_evaluate(self):
+        x = tf.random.uniform((8, 5))
+        y = tf.random.uniform((8, 1))
+        sw = tf.random.uniform((8, 1))
+        model = NoopPipeline()
+        model.compile(loss="mse")
+        # With sample weight.
+        model.evaluate(x=x, y=y, sample_weight=sw, batch_size=8)
+        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
+        # Without sample weight.
+        model.evaluate(x=x, y=y, batch_size=8)
+        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    def test_predict(self):
+        x = tf.random.uniform((8, 5))
+        model = NoopPipeline()
+        model.compile(loss="mse")
+        model.predict(x=x, batch_size=8)
+        model.predict(tf.data.Dataset.from_tensor_slices(x).batch(8))
+
+    def test_on_batch(self):
+        x = tf.random.uniform((8, 5))
+        y = tf.random.uniform((8, 1))
+        sw = tf.random.uniform((8, 1))
+        model = NoopPipeline()
+        model.compile(loss="mse")
+        # With sample weight.
+        model.train_on_batch(x=x, y=y, sample_weight=sw)
+        model.test_on_batch(x=x, y=y, sample_weight=sw)
+        # Without sample weight.
+        model.train_on_batch(x=x, y=y)
+        model.test_on_batch(x=x, y=y)
+        model.predict_on_batch(x=x)
+
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
+        model = NoopPipeline()
+        x = tf.random.uniform((8, 5))
+        model_output = model.predict(x)
+        save_path = os.path.join(self.get_temp_dir(), "model")
+        model.save(save_path, save_format)
+        restored_model = keras.models.load_model(
+            save_path, custom_objects={"NoopPipeline": NoopPipeline}
+        )
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, NoopPipeline)
+        # Check that output matches.
+        restored_output = restored_model.predict(x)
+        self.assertAllClose(model_output, restored_output)
+
+
+class TestFeaturePreprocessingModel(tf.test.TestCase, parameterized.TestCase):
+    def test_fit_with_preprocessing(self):
+        x = tf.strings.as_string(tf.random.uniform((100, 5)))
+        y = tf.random.uniform((100, 1))
+        sw = tf.random.uniform((100, 1))
+        model = FeaturePipeline()
+        model.compile(loss="mse")
+        # With sample weight.
+        model.fit(x=x, y=y, sample_weight=sw, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
+        # Without sample weight.
+        model.fit(x=x, y=y, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    def test_fit_no_preprocessing(self):
+        x = tf.random.uniform((100, 5))
+        y = tf.random.uniform((100, 1))
+        sw = tf.random.uniform((100, 1))
+        model = FeaturePipeline(include_preprocessing=False)
+        model.compile(loss="mse")
+        # With sample weight.
+        model.fit(x=x, y=y, sample_weight=sw, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
+        # Without sample weight.
+        model.fit(x=x, y=y, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    def test_evaluate_with_preprocessing(self):
+        x = tf.strings.as_string(tf.random.uniform((100, 5)))
+        y = tf.random.uniform((100, 1))
+        sw = tf.random.uniform((100, 1))
+        model = FeaturePipeline()
+        model.compile(loss="mse")
+        # With sample weight.
+        model.evaluate(x=x, y=y, sample_weight=sw, batch_size=8)
+        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
+        # Without sample weight.
+        model.evaluate(x=x, y=y, batch_size=8)
+        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    def test_evaluate_no_preprocessing(self):
+        x = tf.random.uniform((100, 5))
+        y = tf.random.uniform((100, 1))
+        sw = tf.random.uniform((100, 1))
+        model = FeaturePipeline(include_preprocessing=False)
+        model.compile(loss="mse")
+        # With sample weight.
+        model.evaluate(x=x, y=y, sample_weight=sw, batch_size=8)
+        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
+        # Without sample weight.
+        model.evaluate(x=x, y=y, batch_size=8)
+        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    def test_predict_with_preprocessing(self):
+        x = tf.strings.as_string(tf.random.uniform((100, 5)))
+        model = FeaturePipeline()
+        model.compile(loss="mse")
+        model.predict(x=x, batch_size=8)
+        model.predict(tf.data.Dataset.from_tensor_slices(x).batch(8))
+
+    def test_predict_no_preprocessing(self):
+        x = tf.random.uniform((100, 5))
+        model = FeaturePipeline(include_preprocessing=False)
+        model.compile(loss="mse")
+        model.predict(x=x, batch_size=8)
+        model.predict(tf.data.Dataset.from_tensor_slices(x).batch(8))
+
+    def test_on_batch(self):
+        x = tf.strings.as_string(tf.random.uniform((8, 5)))
+        y = tf.random.uniform((8, 1))
+        sw = tf.random.uniform((8, 1))
+        model = FeaturePipeline()
+        model.compile(loss="mse")
+        # With sample weight.
+        model.train_on_batch(x=x, y=y, sample_weight=sw)
+        model.test_on_batch(x=x, y=y, sample_weight=sw)
+        # Without sample weight.
+        model.train_on_batch(x=x, y=y)
+        model.test_on_batch(x=x, y=y)
+        model.predict_on_batch(x=x)
+
+    def test_on_batch_no_preprocessing(self):
+        x = tf.random.uniform((8, 5))
+        y = tf.random.uniform((8, 1))
+        sw = tf.random.uniform((8, 1))
+        model = FeaturePipeline(include_preprocessing=False)
+        model.compile(loss="mse")
+        # With sample weight.
+        model.train_on_batch(x=x, y=y, sample_weight=sw)
+        model.test_on_batch(x=x, y=y, sample_weight=sw)
+        # Without sample weight.
+        model.train_on_batch(x=x, y=y)
+        model.test_on_batch(x=x, y=y)
+        model.predict_on_batch(x=x)
+
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
+        model = FeaturePipeline()
+        x = tf.strings.as_string(tf.random.uniform((8, 5)))
+        model_output = model.predict(x)
+        save_path = os.path.join(self.get_temp_dir(), "model")
+        model.save(save_path, save_format)
+        restored_model = keras.models.load_model(
+            save_path, custom_objects={"FeaturePipeline": FeaturePipeline}
+        )
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, FeaturePipeline)
+        # Check that output matches.
+        restored_output = restored_model.predict(x)
+        self.assertAllClose(model_output, restored_output)
+
+
+class TestLabelPreprocessingModel(tf.test.TestCase, parameterized.TestCase):
+    def test_fit_with_preprocessing(self):
+        x = tf.random.uniform((100, 5))
+        y = tf.strings.as_string(tf.random.uniform((100, 1)))
+        sw = tf.random.uniform((100, 1))
+        model = LabelPipeline()
+        model.compile(loss="mse")
+        # With sample weight.
+        model.fit(x=x, y=y, sample_weight=sw, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
+        # Without sample weight.
+        model.fit(x=x, y=y, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    def test_fit_no_preprocessing(self):
+        x = tf.random.uniform((100, 5))
+        y = tf.random.uniform((100, 1))
+        sw = tf.random.uniform((100, 1))
+        model = LabelPipeline(include_preprocessing=False)
+        model.compile(loss="mse")
+        # With sample weight.
+        model.fit(x=x, y=y, sample_weight=sw, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
+        # Without sample weight.
+        model.fit(x=x, y=y, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    def test_evaluate_with_preprocessing(self):
+        x = tf.random.uniform((100, 5))
+        y = tf.strings.as_string(tf.random.uniform((100, 1)))
+        sw = tf.random.uniform((100, 1))
+        model = LabelPipeline()
+        model.compile(loss="mse")
+        # With sample weight.
+        model.evaluate(x=x, y=y, sample_weight=sw, batch_size=8)
+        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
+        # Without sample weight.
+        model.evaluate(x=x, y=y, batch_size=8)
+        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    def test_evaluate_no_preprocessing(self):
+        x = tf.random.uniform((100, 5))
+        y = tf.random.uniform((100, 1))
+        sw = tf.random.uniform((100, 1))
+        model = LabelPipeline(include_preprocessing=False)
+        model.compile(loss="mse")
+        # With sample weight.
+        model.evaluate(x=x, y=y, sample_weight=sw, batch_size=8)
+        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
+        # Without sample weight.
+        model.evaluate(x=x, y=y, batch_size=8)
+        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    def test_predict_with_preprocessing(self):
+        x = tf.random.uniform((100, 5))
+        model = LabelPipeline()
+        model.compile(loss="mse")
+        model.predict(x=x, batch_size=8)
+        model.predict(tf.data.Dataset.from_tensor_slices(x).batch(8))
+
+    def test_on_batch(self):
+        x = tf.random.uniform((8, 5))
+        y = tf.strings.as_string(tf.random.uniform((8, 1)))
+        sw = tf.random.uniform((8, 1))
+        model = LabelPipeline()
+        model.compile(loss="mse")
+        # With sample weight.
+        model.train_on_batch(x=x, y=y, sample_weight=sw)
+        model.test_on_batch(x=x, y=y, sample_weight=sw)
+        # Without sample weight.
+        model.train_on_batch(x=x, y=y)
+        model.test_on_batch(x=x, y=y)
+        model.predict_on_batch(x=x)
+
+    def test_on_batch_no_preprocessing(self):
+        x = tf.random.uniform((8, 5))
+        y = tf.random.uniform((8, 1))
+        sw = tf.random.uniform((8, 1))
+        model = LabelPipeline(include_preprocessing=False)
+        model.compile(loss="mse")
+        # With sample weight.
+        model.train_on_batch(x=x, y=y, sample_weight=sw)
+        model.test_on_batch(x=x, y=y, sample_weight=sw)
+        # Without sample weight.
+        model.train_on_batch(x=x, y=y)
+        model.test_on_batch(x=x, y=y)
+        model.predict_on_batch(x=x)
+
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
+        model = LabelPipeline()
+        x = tf.random.uniform((8, 5))
+        model_output = model.predict(x)
+        save_path = os.path.join(self.get_temp_dir(), "model")
+        model.save(save_path, save_format)
+        restored_model = keras.models.load_model(
+            save_path, custom_objects={"LabelPipeline": LabelPipeline}
+        )
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, LabelPipeline)
+        # Check that output matches.
+        restored_output = restored_model.predict(x)
+        self.assertAllClose(model_output, restored_output)
+
+
+class TestDataPreprocessingModel(tf.test.TestCase, parameterized.TestCase):
+    def test_fit_with_preprocessing(self):
+        data = tf.strings.as_string(tf.random.uniform((100, 1)))
+        model = DataPipeline()
+        model.compile(loss="mse")
+        model.fit(x=data, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices(data).batch(8))
+
+    def test_fit_no_preprocessing(self):
+        x = tf.random.uniform((100, 1))
+        y = tf.random.uniform((100, 1))
+        model = DataPipeline(include_preprocessing=False)
+        model.compile(loss="mse")
+        model.fit(x=x, y=y, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    def test_evaluate_with_preprocessing(self):
+        data = tf.strings.as_string(tf.random.uniform((100, 1)))
+        model = DataPipeline()
+        model.compile(loss="mse")
+        model.evaluate(x=data, batch_size=8)
+        model.evaluate(tf.data.Dataset.from_tensor_slices(data).batch(8))
+
+    def test_evaluate_no_preprocessing(self):
+        x = tf.random.uniform((100, 1))
+        y = tf.random.uniform((100, 1))
+        model = DataPipeline(include_preprocessing=False)
+        model.compile(loss="mse")
+        model.evaluate(x=x, y=y, batch_size=8)
+        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    def test_predict_with_preprocessing(self):
+        x = tf.strings.as_string(tf.random.uniform((100, 1)))
+        model = DataPipeline()
+        model.compile(loss="mse")
+        model.predict(x=x, batch_size=8)
+        model.predict(tf.data.Dataset.from_tensor_slices(x).batch(8))
+
+    def test_predict_no_preprocessing(self):
+        x = tf.random.uniform((100, 1))
+        model = DataPipeline(include_preprocessing=False)
+        model.compile(loss="mse")
+        model.predict(x=x, batch_size=8)
+        model.predict(tf.data.Dataset.from_tensor_slices(x).batch(8))
+
+    def test_on_batch(self):
+        data = tf.strings.as_string(tf.random.uniform((8, 1)))
+        model = DataPipeline()
+        model.compile(loss="mse")
+        # With sample weight.
+        model.train_on_batch(x=data)
+        model.test_on_batch(x=data)
+        # Without sample weight.
+        model.train_on_batch(x=data)
+        model.test_on_batch(x=data)
+        model.predict_on_batch(x=data)
+
+    def test_on_batch_no_preprocessing(self):
+        x = tf.random.uniform((8, 1))
+        y = tf.random.uniform((8, 1))
+        sw = tf.random.uniform((8, 1))
+        model = DataPipeline(include_preprocessing=False)
+        model.compile(loss="mse")
+        # With sample weight.
+        model.train_on_batch(x=x, y=y, sample_weight=sw)
+        model.test_on_batch(x=x, y=y, sample_weight=sw)
+        # Without sample weight.
+        model.train_on_batch(x=x, y=y)
+        model.test_on_batch(x=x, y=y)
+        model.predict_on_batch(x=x)
+
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
+        model = DataPipeline()
+        data = tf.strings.as_string(tf.random.uniform((8, 1)))
+        model_output = model.predict(data)
+        save_path = os.path.join(self.get_temp_dir(), "model")
+        model.save(save_path, save_format)
+        restored_model = keras.models.load_model(
+            save_path, custom_objects={"DataPipeline": DataPipeline}
+        )
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, DataPipeline)
+        # Check that output matches.
+        restored_output = restored_model.predict(data)
+        self.assertAllClose(model_output, restored_output)
+
+
+class TestFunctional(tf.test.TestCase, parameterized.TestCase):
+    def test_fit(self):
+        x = tf.strings.as_string(tf.random.uniform((100, 5)))
+        y = tf.random.uniform((100, 1))
+        sw = tf.random.uniform((100, 1))
+
+        model = FunctionalPipeline()
+        model.compile(loss="mse")
+        # With sample weight.
+        model.fit(x=x, y=y, sample_weight=sw, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
+        # Without sample weight.
+        model.fit(x=x, y=y, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    def test_fit_no_preprocessing(self):
+        x = tf.random.uniform((100, 5))
+        y = tf.random.uniform((100, 1))
+        sw = tf.random.uniform((100, 1))
+        model = FunctionalPipeline(include_preprocessing=False)
+        model.compile(loss="mse")
+        # With sample weight.
+        model.fit(x=x, y=y, sample_weight=sw, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
+        # Without sample weight.
+        model.fit(x=x, y=y, batch_size=8)
+        model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
+
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
+        model = FunctionalPipeline()
+        x = tf.strings.as_string(tf.random.uniform((8, 5)))
+        model_output = model.predict(x)
+        save_path = os.path.join(self.get_temp_dir(), "model")
+        model.save(save_path, save_format)
+        restored_model = keras.models.load_model(
+            save_path, custom_objects={"FunctionalPipeline": FunctionalPipeline}
+        )
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, FunctionalPipeline)
+        # Check that output matches.
+        restored_output = restored_model.predict(x)
+        self.assertAllClose(model_output, restored_output)
+
+
+class TestFitArguments(tf.test.TestCase):
+    def test_validation_data(self):
+        x = tf.strings.as_string(tf.random.uniform((80, 5)))
+        y = tf.random.uniform((80, 1))
+        val_x = tf.strings.as_string(tf.random.uniform((20, 5)))
+        val_y = tf.random.uniform((20, 1))
+
+        model = FeaturePipeline()
+        model.compile(loss="mse")
+
+        model.fit(x=x, y=y, validation_data=(val_x, val_y), batch_size=8)
+        model.fit(
+            x=tf.data.Dataset.from_tensor_slices((x, y)).batch(8),
+            validation_data=tf.data.Dataset.from_tensor_slices(
+                (val_x, val_y)
+            ).batch(8),
+        )
+
+    def test_validation_split(self):
+        x = tf.strings.as_string(tf.random.uniform((100, 5)))
+        y = tf.random.uniform((100, 1))
+
+        model = FeaturePipeline()
+        model.compile(loss="mse")
+
+        model.fit(x=x, y=y, validation_split=0.2, batch_size=8)
+
+    def test_error_dataset_and_invalid_arguments(self):
+        x = tf.strings.as_string(tf.random.uniform((100, 5)))
+        y = tf.random.uniform((100, 1))
+        sw = tf.random.uniform((100, 1))
+        ds = tf.data.Dataset.from_tensor_slices((x, y))
+
+        model = FeaturePipeline()
+        model.compile(loss="mse")
+        with self.assertRaises(ValueError):
+            model.fit(ds, validation_split=0.2)
+        with self.assertRaises(ValueError):
+            model.fit(ds, batch_size=0.2)
+        with self.assertRaises(ValueError):
+            model.fit(ds, y=y)
+        with self.assertRaises(ValueError):
+            model.fit(ds, sample_weight=sw)