Repository: keras-team/keras-nlp

Issue:
Return All Beams from Beam Search Utility
Follow-up to https://github.com/keras-team/keras-nlp/issues/750#issuecomment-1439963933

- Return all beams (tensor of shape `(batch_size, num_beams, sequence_length)`).
- Return the beam scores as well (tensor of shape `(batch_size, num_beams)`).

The changes are to be made to this utility function: https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/utils/text_generation.py#L213.

@mattdangerw, do we want to sort the outputs based on log prob and return? Or is that unnecessary overhead?

Some basic tinkering done because a user wanted a quick fix: https://github.com/abheesht17/keras-nlp/blob/beam_search_return_all_seqs/keras_nlp/utils/text_generation.py#L400-L414

Discussion:
TheAthleticCoder: @mattdangerw 
Done with the requested changes. I had to change the names of a few variables in the `BeamSampler` function. Do let me know if I need to rename them to something else 🤔 

Also added a unit test function. The tests were quite simple in which I just checked the dimensionality of all the returned `prompts` and the `log_probs`. Are there any more tests you would like me to add? 
Thanks!

Patch:
diff --git a/CONTRIBUTING_MODELS.md b/CONTRIBUTING_MODELS.md
new file mode 100644
index 00000000..de5e71af
--- /dev/null
+++ b/CONTRIBUTING_MODELS.md
@@ -0,0 +1,240 @@
+# Model Contribution Guide
+
+KerasNLP has a plethora of pre-trained large language models
+ranging from BERT to OPT. We are always looking for more models and are always
+open to contributions!
+
+In this guide, we will walk you through the steps one needs to take in order to
+contribute a new pre-trained model to KerasNLP. For illustration purposes, let's
+assume that you want to contribute the DistilBERT model. Before we dive in, we encourage you to go through
+[our getting started guide](https://keras.io/guides/keras_nlp/getting_started/)
+for an introduction to the library, and our
+[contribution guide](https://github.com/keras-team/keras-nlp/blob/master/CONTRIBUTING.md).
+
+## Checklist
+
+This to-do list is a brief outline of how a model can be contributed.
+Keep this checklist handy!
+
+### Step 1: Open an issue/find an issue
+
+- [ ] Open an issue or find an issue to contribute a backbone model.
+
+### Step 2: PR #1 - Add XXBackbone
+
+- [ ] An `xx/xx_backbone.py` file which has the model graph \[[Example](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_backbone.py)\].
+- [ ] An `xx/xx_backbone_test.py` file which has unit tests for the backbone \[[Example](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_backbone_test.py)\].
+- [ ] A Colab notebook link in the PR description which matches the outputs of the implemented backbone model with the original source \[[Example](https://colab.research.google.com/drive/1SeZWJorKWmwWJax8ORSdxKrxE25BfhHa?usp=sharing)\].
+
+### Step 3: PR #2 - Add XXTokenizer
+
+- [ ] An `xx/xx_tokenizer.py` file which has the tokenizer for the model \[[Example](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_tokenizer.py)\].
+- [ ] An `xx/xx_tokenizer_test.py` file which has unit tests for the model tokenizer \[[Example](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py)\].
+- [ ] A Colab notebook link in the PR description, demonstrating that the output of the tokenizer matches the original tokenizer \[[Example](https://colab.research.google.com/drive/1MH_rpuFB1Nz_NkKIAvVtVae2HFLjXZDA?usp=sharing)].
+
+### Step 4: PR #3 - Add XX Presets
+
+- [ ] An `xx/xx_presets.py` file with links to weights uploaded to a personal GCP bucket/Google Drive \[[Example](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_presets.py)\].
+- [ ] An `xx/xx_presets_test.py` file with runnable tests for each preset \[[Example](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_presets_test.py)\].
+- [ ] A `tools/checkpoint_conversion/convert_xx_checkpoints.py` which is reusable script for converting checkpoints \[[Example](https://github.com/keras-team/keras-nlp/blob/master/tools/checkpoint_conversion/convert_distilbert_checkpoints.py)\].
+- [ ] A Colab notebook link in the PR description, showing an end-to-end task such as text classification, etc. The task model can be built using the backbone model, with the task head on top \[[Example](https://gist.github.com/mattdangerw/bf0ca07fb66b6738150c8b56ee5bab4e)\].
+
+### Step 5: PR #4 and Beyond - Add XX Tasks and Preprocessors
+
+This PR is optional.
+
+- [ ] An `xx/xx_<task>.py` file for adding a task model like classifier, masked LM, etc. \[[Example](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_classifier.py)\]
+- [ ] An `xx/xx_<task>_preprocessor.py` file which has the preprocessor and can be used to get inputs suitable for the task model \[[Example](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_preprocessor.py)\].
+- [ ] `xx/xx_<task>_test.py` file and `xx/xx_<task>_preprocessor_test.py` files which have unit tests for the above two modules \[[Example 1](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_classifier_test.py) and [Example 2](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_preprocessor_test.py)\].
+- [ ] A Colab notebook link in the PR description, demonstrating that the output of the preprocessor matches the output of the original preprocessor \[[Example](https://colab.research.google.com/drive/1GFFC7Y1I_2PtYlWDToqKvzYhHWv1b3nC?usp=sharing)].
+
+## Detailed Instructions
+
+This section discusses, in details, every necessary step.
+
+### Step 1: Open an issue/Find an open issue
+
+Before getting started with the code, it's important to check if there are any
+[open issues](https://github.com/keras-team/keras-nlp/issues?q=is%3Aissue+is%3Aopen+label%3Amodel-contribution)
+related to the model you wish to contribute. If there is an open issue, you can
+claim it by commenting on the issue and letting us know that you're interested
+in working on it. This helps us keep track of who is working on what and avoid
+duplicated effort.
+
+If there aren't any open issues, you can create one by clicking the "New Issue"
+button on our repository page.
+
+Note that you need not have all the answers or complete knowledge of the inner
+workings of the model at the time of opening the issue. But it is appreciated if
+you can furnish as much detail as possible to enable us to help you with the
+contribution! 🙂
+
+### Step 2: PR #1 - Add XXBackbone
+
+#### Add the backbone class
+
+Once you are done identifying all the required layers, you should implement the
+model backbone class.
+
+To keep the code simple and readable, we follow
+[Keras' functional style model](https://keras.io/guides/functional_api/) wrapped
+around by a class to implement our models.
+
+A model is typically split into three/four sections. We would recommend you to
+compare this side-by-side with the
+[`keras_nlp.layers.DistilBertBackbone` source code](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_backbone.py)!
+
+**Inputs to the model**
+
+Generally, the standard inputs to any text model are:
+  - `token_ids`: tokenised inputs (An integer representation of the text sequence).
+  - `padding_mask`: Masks the padding tokens.
+
+**Embedding layer(s)**
+
+Standard layers used: `keras.layers.Embedding`,
+`keras_nlp.layers.PositionEmbedding`, `keras_nlp.layers.TokenAndPositionEmbedding`.
+
+**Encoder layers**
+
+Standard layers used: `keras_nlp.layers.TransformerEncoder`, `keras_nlp.layers.FNetEncoder`.
+
+**Decoder layers (possibly)**
+
+Standard layers used: `keras_nlp.layers.TransformerDecoder`.
+
+**Other layers which might be used**
+
+`keras.layers.LayerNorm`, `keras.layers.Dropout`, `keras.layers.Conv1D`, etc.
+
+<br/>
+
+The standard layers provided in Keras and KerasNLP are generally enough for
+most of the usecases and it is recommended to do a thorough search
+[here](https://keras.io/api/layers/) and [here](https://keras.io/api/keras_nlp/layers/).
+However, sometimes, models have small tweaks/paradigm changes in their architecture.
+This is when things might slightly get complicated.
+
+If the model introduces a paradigm shift, such as using relative attention instead
+of vanilla attention, the contributor will have to implement complete custom layers. A case
+in point is `keras_nlp.models.DebertaV3Backbone` where we had to [implement layers
+from scratch](https://github.com/keras-team/keras-nlp/tree/master/keras_nlp/models/deberta_v3).
+
+On the other hand, if the model has a small tweak, something simpler can be done.
+For instance, in the Whisper model, the self-attention and cross-attention mechanism
+is exactly the same as vanilla attention, with the exception that the key projection
+layer does not have a bias term. In this case, we can inherit the custom layer
+from one of the standard layers and make minor modifications. See [this PR](https://github.com/keras-team/keras-nlp/pull/801/files#diff-8533ae3a7755c0dbe95ccbb71f85c677297f687bf3884fadefc64f1d0fdce51aR22) for
+more details.
+
+Since the first PR is only to add the model backbone class, you should omit the
+`from_presets()` function; this will be added at a later stage when you open a PR
+for adding presets.
+
+#### Convert weights from the original source and check output!
+
+Before you open a PR for adding the model backbone class, it is essential to check
+whether the model has been implemented exactly as the source implementation. This
+also helps in adding model "presets" at a later stage.
+
+The preferred way of doing this is to add a Colab link in the PR description, which
+1) converts the original preset weights to our format, and
+2) checks whether the outputs of the original model and your implemented model are close enough.
+
+It is okay if you demonstrate it for one preset at this stage; you can do the conversion
+for the other presets when you officially add presets to the library at a later stage.
+
+#### Add Unit Tests
+
+It is essential to add units tests. These unit tests are basic and mostly check
+whether the forward pass goes through successfully, whether the model can be saved
+and loaded correctly, etc.
+
+### Step 3: PR #2 - Add XXTokenizer
+
+#### Tokenizer
+
+Most text models nowadays use subword tokenizers such as WordPiece, SentencePiece
+and BPE Tokenizer. Since KerasNLP has implementations of most of the popular
+subword tokenizers, the model tokenizer layer typically inherits from a base
+tokenizer class.
+
+For example, DistilBERT uses the WordPiece tokenizer. So, we can introduce a new
+class, `DistilBertTokenizer`, which inherits from `keras_nlp.tokenizers.WordPieceTokenizer`.
+All the underlying actual tokenization will be taken care of by the superclass.
+
+The important thing here is adding "special tokens". Most models have
+special tokens such as beginning-of-sequence token, end-of-sequence token,
+mask token, pad token, etc. These have to be
+[added as member attributes](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_tokenizer.py#L91-L105)
+to the tokenizer class. These member attributes are then accessed by the
+preprocessor layers.
+
+For a full list of the tokenizers KerasNLP offers, please visit
+[this link](https://keras.io/api/keras_nlp/tokenizers/) and make use of the
+tokenizer your model uses!
+
+#### Unit Tests
+
+The last step here is to add unit tests for the tokenizer. A dummy vocabulary is
+created, and the output of both these layers is verified including tokenization,
+detokenization, etc.
+
+### Step 4: PR #3 - Add XX Presets
+
+Once the backbone and tokenizer PRs have been merged, you can open a PR for
+adding presets. For every model, we have a separate file where we mention our
+preset configurations. This preset configuration has model-specific arguments
+such as number of layers, number of attention heads; preprocessor-specific
+arguments such as whether we want to lowercase the input text; checkpoint and
+vocabulary file URLs, etc. In the PR description, you can add
+Google Drive/personal GCP bucket links to the checkpoint and the vocabulary
+files. These files will then be uploaded to GCP by us!
+
+After wrapping up the preset configuration file, you need to
+add the `from_preset` function to all three classes, i.e., `DistilBertBackbone`,
+and `DistilBertTokenizer`. Here is an
+[example](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_backbone.py#L187-L189).
+
+The testing for presets is divided into two: "large" and "extra large".
+For "large" tests, we pick the smallest preset (in terms of number of parameters)
+and verify whether the output is correct. For "extra large tests", we loop over
+all the presets and just check whether the backbone and the tokenizer can
+be called without any error.
+
+Additionally, a checkpoint conversion script should be added. This script
+demonstrates that the outputs of our backbone model and outputs of the source
+model match. This should be done for all presets.
+
+### Step 5: PR #4 and Beyond: Add XXTasks and XXPreprocessors
+
+Once you are finished with Steps 1-4, you can add "task" models and
+preprocessors.
+
+### Task model
+
+Task models are essentially models which have "task heads" on top of the backbone
+models. For instance, for the text classification task, you can have a
+feedforward layer on top of a backbone model like DistilBERT. Task models are
+very essential since pretrained models are used extensively for downstream tasks
+like text classification, token classification, text summarization, neural
+machine translation, etc.
+
+#### Preprocessor
+
+The preprocessor class is responsible for making the inputs suitable for
+consumption by the model - it packs multiple inputs together, i.e., given
+multiple input texts, it will add appropriate special tokens, pad the inputs
+and return the dictionary in the form expected by the model.
+
+The preprocessor class might have a few intricacies depending on the model. For example,
+the DeBERTaV3 tokenizer does not have the `[MASK]` in the provided sentencepiece
+proto file, and we had to make some modifications [here](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/deberta_v3/deberta_v3_preprocessor.py). Secondly, we have
+a separate preprocessor class for every task. This is because different tasks
+might require different input formats. For instance, we have a [separate preprocessor](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor.py)
+for masked language modeling (MLM) for DistilBERT.
+
+## Conclusion
+
+Once all three PRs (and optionally, the fourth PR) have been merged, you have
+successfully contributed a model to KerasNLP. Congratulations! 🔥
diff --git a/keras_nlp/benchmarks/README.md b/benchmarks/README.md
similarity index 100%
rename from keras_nlp/benchmarks/README.md
rename to benchmarks/README.md
diff --git a/benchmarks/glue.py b/benchmarks/glue.py
new file mode 100644
index 00000000..43caf9a4
--- /dev/null
+++ b/benchmarks/glue.py
@@ -0,0 +1,191 @@
+# Copyright 2023 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.
+
+"""GLUE benchmark script to test model performance.
+
+To run the script, use this command:
+```
+python3 glue.py --model BertClassifier \
+                --preset bert_base_en \
+                --epochs 5 \
+                --batch_size 16 \
+                --learning_rate 0.001 \
+                --mixed_precision_policy mixed_float16
+```
+
+Disclaimer: This script only supports GLUE/mrpc (for now).
+"""
+
+import inspect
+import time
+
+import tensorflow as tf
+import tensorflow_datasets as tfds
+from absl import app
+from absl import flags
+from absl import logging
+from tensorflow import keras
+
+import keras_nlp
+
+seed = 42
+tf.random.set_seed(seed)
+
+
+flags.DEFINE_string(
+    "task",
+    "mrpc",
+    "The name of the GLUE task to finetune on.",
+)
+flags.DEFINE_string(
+    "model", None, "The name of the classifier such as BertClassifier."
+)
+flags.DEFINE_string(
+    "preset",
+    None,
+    "The model preset, e.g., 'bert_base_en_uncased' for `BertClassifier`",
+)
+flags.DEFINE_float(
+    "learning_rate", 0.005, "The learning_rate for the optimizer."
+)
+flags.DEFINE_string(
+    "mixed_precision_policy",
+    "mixed_float16",
+    "The global precision policy to use, e.g., 'mixed_float16' or 'float32'.",
+)
+flags.DEFINE_integer("epochs", 2, "The number of epochs.")
+flags.DEFINE_integer("batch_size", 8, "Batch Size.")
+
+
+FLAGS = flags.FLAGS
+
+
+def load_data():
+    """Load data.
+
+    Load GLUE/MRPC dataset, and convert the dictionary format to
+    (features, label), where `features` is a tuple of all input sentences.
+    """
+    feature_names = ("sentence1", "sentence2")
+
+    def split_features(x):
+        # GLUE comes with 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)
+
+    train_ds, test_ds, validation_ds = tfds.load(
+        "glue/mrpc",
+        split=["train", "test", "validation"],
+    )
+
+    train_ds = train_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)
+    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 load_model(model, preset, num_classes):
+    for name, symbol in keras_nlp.models.__dict__.items():
+        if inspect.isclass(symbol) and issubclass(symbol, keras.Model):
+            if model and name != model:
+                continue
+            if not hasattr(symbol, "from_preset"):
+                continue
+            for preset in symbol.presets:
+                if preset and preset != preset:
+                    continue
+                model = symbol.from_preset(preset, num_classes=num_classes)
+                logging.info(f"\nUsing model {name} with preset {preset}\n")
+                return model
+
+    raise ValueError(f"Model {model} or preset {preset} not found.")
+
+
+def main(_):
+    keras.mixed_precision.set_global_policy(FLAGS.mixed_precision_policy)
+
+    # Check task is supported.
+    # TODO(chenmoneygithub): Add support for other glue tasks.
+    if FLAGS.task != "mrpc":
+        raise ValueError(
+            f"For now only mrpc is supported, but received {FLAGS.task}."
+        )
+
+    logging.info(
+        "Benchmarking configs...\n"
+        "=========================\n"
+        f"MODEL: {FLAGS.model}\n"
+        f"PRESET: {FLAGS.preset}\n"
+        f"TASK: glue/{FLAGS.task}\n"
+        f"BATCH_SIZE: {FLAGS.batch_size}\n"
+        f"EPOCHS: {FLAGS.epochs}\n"
+        "=========================\n"
+    )
+
+    # Load datasets.
+    train_ds, test_ds, validation_ds = load_data()
+    train_ds = train_ds.batch(FLAGS.batch_size).prefetch(tf.data.AUTOTUNE)
+    test_ds = test_ds.batch(FLAGS.batch_size).prefetch(tf.data.AUTOTUNE)
+    validation_ds = validation_ds.batch(FLAGS.batch_size).prefetch(
+        tf.data.AUTOTUNE
+    )
+
+    # Load the model.
+    model = load_model(model=FLAGS.model, preset=FLAGS.preset, num_classes=2)
+    # Set loss and metrics.
+    loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+    metrics = [keras.metrics.SparseCategoricalAccuracy()]
+    # Configure optimizer.
+    lr = tf.keras.optimizers.schedules.PolynomialDecay(
+        FLAGS.learning_rate,
+        decay_steps=train_ds.cardinality() * FLAGS.epochs,
+        end_learning_rate=0.0,
+    )
+    optimizer = tf.keras.optimizers.experimental.AdamW(lr, weight_decay=0.01)
+    optimizer.exclude_from_weight_decay(
+        var_names=["LayerNorm", "layer_norm", "bias"]
+    )
+    model.compile(optimizer=optimizer, loss=loss, metrics=metrics)
+
+    # Start training.
+    logging.info("Starting Training...")
+
+    st = time.time()
+    history = model.fit(
+        train_ds,
+        validation_data=validation_ds,
+        epochs=FLAGS.epochs,
+    )
+
+    wall_time = time.time() - st
+    validation_accuracy = history.history["val_sparse_categorical_accuracy"][-1]
+    examples_per_second = (
+        FLAGS.epochs * FLAGS.batch_size * (len(train_ds) + len(validation_ds))
+    ) / wall_time
+
+    logging.info("Training Finished!")
+    logging.info(f"Wall Time: {wall_time:.4f} seconds.")
+    logging.info(f"Validation Accuracy: {validation_accuracy:.4f}")
+    logging.info(f"examples_per_second: {examples_per_second:.4f}")
+
+
+if __name__ == "__main__":
+    flags.mark_flag_as_required("model")
+    flags.mark_flag_as_required("preset")
+    app.run(main)
diff --git a/keras_nlp/benchmarks/sentiment_analysis.py b/benchmarks/sentiment_analysis.py
similarity index 100%
rename from keras_nlp/benchmarks/sentiment_analysis.py
rename to benchmarks/sentiment_analysis.py
diff --git a/keras_nlp/benchmarks/text_generation.py b/benchmarks/text_generation.py
similarity index 100%
rename from keras_nlp/benchmarks/text_generation.py
rename to benchmarks/text_generation.py
diff --git a/keras_nlp/benchmarks/__init__.py b/keras_nlp/benchmarks/__init__.py
deleted file mode 100644
index ba0c2545..00000000
--- a/keras_nlp/benchmarks/__init__.py
+++ /dev/null
@@ -1,13 +0,0 @@
-# Copyright 2023 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.
diff --git a/keras_nlp/layers/masked_lm_mask_generator.py b/keras_nlp/layers/masked_lm_mask_generator.py
index f235430d..cfa854f4 100644
--- a/keras_nlp/layers/masked_lm_mask_generator.py
+++ b/keras_nlp/layers/masked_lm_mask_generator.py
@@ -163,7 +163,7 @@ class MaskedLMMaskGenerator(keras.layers.Layer):
             token_ids = token_ids.to_tensor()
 
         mask_weights = tf.ones_like(mask_positions, self.compute_dtype)
-        # If mask_selection_length is set, covert to raggeds to dense.
+        # If `mask_selection_length` is set, convert to dense.
         if self.mask_selection_length:
             target_shape = tf.cast([-1, self.mask_selection_length], tf.int64)
             mask_positions = mask_positions.to_tensor(shape=target_shape)
diff --git a/keras_nlp/models/task.py b/keras_nlp/models/task.py
index f1497f06..8b710c44 100644
--- a/keras_nlp/models/task.py
+++ b/keras_nlp/models/task.py
@@ -15,8 +15,11 @@
 
 import os
 
+import tensorflow as tf
 from tensorflow import keras
 
+from keras_nlp.utils.keras_utils import print_msg
+from keras_nlp.utils.keras_utils import print_row
 from keras_nlp.utils.pipeline_model import PipelineModel
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
@@ -49,6 +52,7 @@ class Task(PipelineModel):
         return self._preprocessor
 
     @preprocessor.setter
+    @tf.__internal__.tracking.no_automatic_dependency_tracking
     def preprocessor(self, value):
         self.include_preprocessing = value is not None
         self._preprocessor = value
@@ -173,3 +177,42 @@ class Task(PipelineModel):
                 example_preset_name=next(iter(cls.presets), ""),
                 preset_names='", "'.join(cls.presets),
             )(cls.from_preset.__func__)
+
+    def summary(
+        self,
+        line_length=None,
+        positions=None,
+        print_fn=None,
+        **kwargs,
+    ):
+        """Override `model.summary()` to show a preprocessor if set."""
+        # Defaults are copied from core Keras; we should try to stay in sync.
+        line_length = line_length or 98
+        positions = positions or [0.33, 0.55, 0.67, 1.0]
+        if positions[-1] <= 1:
+            positions = [int(line_length * p) for p in positions]
+        if print_fn is None:
+            print_fn = print_msg
+
+        if self.preprocessor:
+            column_names = ["Tokenizer (type)", "Vocab #"]
+            tokenizer = self.preprocessor.tokenizer
+            column_values = [
+                f"{tokenizer.name} ({tokenizer.__class__.__name__})",
+                f"{tokenizer.vocabulary_size()}",
+            ]
+
+            print_fn(f'Preprocessor: "{self.preprocessor.name}"')
+            print_fn("_" * line_length)
+            print_row(column_names, positions[1:3], print_fn)
+            print_fn("=" * line_length)
+            print_row(column_values, positions[1:3], print_fn)
+            print_fn("_" * line_length)
+            print_fn(" " * line_length)
+
+        super().summary(
+            line_length=line_length,
+            positions=positions,
+            print_fn=print_fn,
+            **kwargs,
+        )
diff --git a/keras_nlp/samplers/beam_sampler.py b/keras_nlp/samplers/beam_sampler.py
index 03b390a3..e140be6f 100644
--- a/keras_nlp/samplers/beam_sampler.py
+++ b/keras_nlp/samplers/beam_sampler.py
@@ -36,11 +36,14 @@ class BeamSampler(Sampler):
     Args:
         num_beams: int. The number of beams that should be kept at each
             time-step. `num_beams` should be strictly positive.
+        return_all_beams: bool. When set to `True`, the sampler will return the top prompt,
+            all prompts and their respective probabilities score.
 
     Call Args:
         {{call_args}}
 
     Examples:
+    Return only the beam with the highest accumulated probability.
     ```python
     # Use a simple alphabet of lowercase characters to [0, 26).
     int_lookup = {i: chr(i + ord('a')) for i in range(26)}
@@ -60,14 +63,41 @@ class BeamSampler(Sampler):
     print(["".join([int_lookup[i] for i in s]) for s in output.numpy()])
     # >>> "zzzzzaaaaaaa"
     ```
+    Return all beams and their probabilities.
+    ```python
+    # Use a simple alphabet of lowercase characters to [0, 26).
+    int_lookup = {i: chr(i + ord('a')) for i in range(26)}
+    char_lookup = {v: k for k, v in int_lookup.items()}
+    batch_size, length, vocab_size = 1, 12, len(int_lookup)
+
+    def next(prompt, state, index):
+        # A uniform distribution over our alphabet.
+        logits = tf.ones((batch_size, vocab_size))
+        return logits, state
+
+    output = keras_nlp.samplers.BeamSampler(return_all_beams=True)(
+        next=next,
+        prompt=tf.fill((batch_size, length,), char_lookup['z']),
+        index=5,
+    )
+
+    print(output[0].shape)
+    # >>> (1, 5, 12)
+    print(output[1].shape)
+    # >>> (1, 5)
+    print(["".join([int_lookup[i] for i in s]) for s in output[0][0].numpy()])
+    # >>> "zzzzzaaaaaaa"
+    ```
     """
 
     def __init__(
         self,
         num_beams=5,
+        return_all_beams=False,
     ):
         super().__init__()
         self.num_beams = num_beams
+        self.return_all_beams = return_all_beams
 
     def __call__(
         self,
@@ -132,6 +162,9 @@ class BeamSampler(Sampler):
             )
             beam_indices = indices // vocab_size
             next_token = flatten_beams(indices % vocab_size)
+            # Ensure shape is `[None]`, otherwise it causes issues after
+            # converting to TFLite.
+            next_token = tf.ensure_shape(next_token, [None])
             # We need `ensure_shape` as `top_k` will change the static shape.
             next_log_probs = flatten_beams(next_log_probs)
             log_probs = tf.ensure_shape(next_log_probs, log_probs.shape)
@@ -161,17 +194,32 @@ class BeamSampler(Sampler):
             maximum_iterations=(max_length - index),
         )
 
-        # Gather the top beam at each batch index.
-        prompt, log_probs = unflatten_beams(prompt), unflatten_beams(log_probs)
-        top_beams = tf.math.argmax(log_probs, axis=-1)[:, tf.newaxis]
-        prompt = tf.gather(prompt, top_beams, axis=1, batch_dims=1)
-        return tf.squeeze(prompt, axis=1)
+        all_prompts = unflatten_beams(prompt)
+        all_log_probs = unflatten_beams(log_probs)
+
+        if self.return_all_beams:
+            sorted_indices = tf.argsort(
+                all_log_probs, axis=-1, direction="DESCENDING"
+            )
+            sorted_log_probs = tf.gather(
+                all_log_probs, sorted_indices, axis=-1, batch_dims=1
+            )
+            sorted_prompts = tf.gather(
+                all_prompts, sorted_indices, axis=1, batch_dims=1
+            )
+            return sorted_prompts, sorted_log_probs
+        else:
+            # Gather the top beam at each batch index.
+            top_beams = tf.math.argmax(all_log_probs, axis=-1)[:, tf.newaxis]
+            prompt = tf.gather(all_prompts, top_beams, axis=1, batch_dims=1)
+            return tf.squeeze(prompt, axis=1)
 
     def get_config(self):
         config = super().get_config()
         config.update(
             {
                 "num_beams": self.num_beams,
+                "return_all_beams": self.return_all_beams,
             }
         )
         return config
diff --git a/keras_nlp/samplers/sampler.py b/keras_nlp/samplers/sampler.py
index dde0598f..39cb80b7 100644
--- a/keras_nlp/samplers/sampler.py
+++ b/keras_nlp/samplers/sampler.py
@@ -110,11 +110,13 @@ class Sampler:
             # Compute the softmax distribution for the next token.
             logits, state = next(prompt, state, index)
             probabilities = keras.activations.softmax(logits)
-
             # Compute the next token.
             next_token = self.get_next_token(probabilities)
             # Don't overwrite anywhere mask is True.
             next_token = tf.cast(next_token, prompt.dtype)
+            # Ensure shape is `[None]`, otherwise it causes issues after
+            # converting to TFLite.
+            next_token = tf.ensure_shape(next_token, [None])
             next_token = tf.where(mask[:, index], prompt[:, index], next_token)
             # Update the prompt with the next token.
             next_token = next_token[:, tf.newaxis]
diff --git a/keras_nlp/utils/keras_utils.py b/keras_nlp/utils/keras_utils.py
index 8d41681b..84feb765 100644
--- a/keras_nlp/utils/keras_utils.py
+++ b/keras_nlp/utils/keras_utils.py
@@ -14,8 +14,10 @@
 
 
 import platform
+import sys
 
 import tensorflow as tf
+from absl import logging
 from tensorflow import keras
 
 from keras_nlp.utils.tf_utils import is_tensor_type
@@ -111,3 +113,60 @@ def is_xla_compatible(model):
             tf.distribute.TPUStrategy,
         ),
     )
+
+
+def print_msg(message, line_break=True):
+    """Print the message to absl logging or stdout."""
+    # Copied from core Keras.
+    if keras.utils.is_interactive_logging_enabled():
+        if line_break:
+            sys.stdout.write(message + "\n")
+        else:
+            sys.stdout.write(message)
+        sys.stdout.flush()
+    else:
+        logging.info(message)
+
+
+def print_row(fields, positions, print_fn, nested_level=0):
+    """Print a row of a summary message."""
+    # Copied from core Keras.
+    left_to_print = [str(x) for x in fields]
+    while any(left_to_print):
+        line = ""
+        for col in range(len(left_to_print)):
+            if col > 0:
+                start_pos = positions[col - 1]
+            else:
+                start_pos = 0
+            end_pos = positions[col]
+            # Leave room for 2 spaces to delineate columns
+            # we don't need any if we are printing the last column
+            space = 2 if col != len(positions) - 1 else 0
+            cutoff = end_pos - start_pos - space
+            fit_into_line = left_to_print[col][:cutoff]
+            # For nicer formatting we line-break on seeing end of
+            # tuple/dict etc.
+            line_break_conditions = ("),", "},", "],", "',")
+            candidate_cutoffs = [
+                fit_into_line.find(x) + len(x)
+                for x in line_break_conditions
+                if fit_into_line.find(x) >= 0
+            ]
+            if candidate_cutoffs:
+                cutoff = min(candidate_cutoffs)
+                fit_into_line = fit_into_line[:cutoff]
+
+            if col == 0:
+                line += "|" * nested_level + " "
+            line += fit_into_line
+            line += " " * space if space else ""
+            left_to_print[col] = left_to_print[col][cutoff:]
+
+            # Pad out to the next position
+            if nested_level:
+                line += " " * (positions[col] - len(line) - nested_level)
+            else:
+                line += " " * (positions[col] - len(line))
+        line += "|" * nested_level
+        print_fn(line)
diff --git a/tools/checkpoint_conversion/checkpoint_conversion_utils.py b/tools/checkpoint_conversion/checkpoint_conversion_utils.py
index 2fea7157..99b133b1 100644
--- a/tools/checkpoint_conversion/checkpoint_conversion_utils.py
+++ b/tools/checkpoint_conversion/checkpoint_conversion_utils.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import hashlib
+import tarfile
+import zipfile
 
 
 def get_md5_checksum(file_path):
@@ -20,3 +22,12 @@ def get_md5_checksum(file_path):
         for byte_block in iter(lambda: f.read(4096), b""):
             md5_hash.update(byte_block)
     return md5_hash.hexdigest()
+
+
+def extract_files_from_archive(archive_file_path):
+    if archive_file_path.endswith(".tar.gz"):
+        with tarfile.open(archive_file_path, "r:gz") as tar:
+            return tar.extractall()
+    elif archive_file_path.endswith(".zip"):
+        with zipfile.ZipFile(archive_file_path, "r") as zip_ref:
+            return zip_ref.extractall()
diff --git a/tools/checkpoint_conversion/convert_albert_checkpoints.py b/tools/checkpoint_conversion/convert_albert_checkpoints.py
index 180478b0..a089f800 100644
--- a/tools/checkpoint_conversion/convert_albert_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_albert_checkpoints.py
@@ -19,11 +19,9 @@ 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
-from tools.checkpoint_conversion.checkpoint_conversion_utils import (
-    get_md5_checksum,
-)
 
 PRESET_MAP = {
     "albert_base_en_uncased": "albert-base-v2",
diff --git a/tools/checkpoint_conversion/convert_bart_checkpoints.py b/tools/checkpoint_conversion/convert_bart_checkpoints.py
index 54d54b96..f5f800f2 100644
--- a/tools/checkpoint_conversion/convert_bart_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_bart_checkpoints.py
@@ -19,11 +19,9 @@ 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
-from tools.checkpoint_conversion.checkpoint_conversion_utils import (
-    get_md5_checksum,
-)
 
 PRESET_MAP = {
     "bart_base_en": "facebook/bart-base",
diff --git a/tools/checkpoint_conversion/convert_deberta_v3_checkpoints.py b/tools/checkpoint_conversion/convert_deberta_v3_checkpoints.py
index 85e49194..82f35a8e 100644
--- a/tools/checkpoint_conversion/convert_deberta_v3_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_deberta_v3_checkpoints.py
@@ -20,15 +20,13 @@ import tensorflow as tf
 import transformers
 from absl import app
 from absl import flags
+from checkpoint_conversion_utils import get_md5_checksum
 
 from keras_nlp.models.deberta_v3.deberta_v3_backbone import DebertaV3Backbone
 from keras_nlp.models.deberta_v3.deberta_v3_preprocessor import (
     DebertaV3Preprocessor,
 )
 from keras_nlp.models.deberta_v3.deberta_v3_tokenizer import DebertaV3Tokenizer
-from tools.checkpoint_conversion.checkpoint_conversion_utils import (
-    get_md5_checksum,
-)
 
 PRESET_MAP = {
     "deberta_v3_extra_small_en": "microsoft/deberta-v3-xsmall",
@@ -46,10 +44,10 @@ flags.DEFINE_string(
 )
 
 
-def download_files(preset, hf_model_name):
+def download_files(hf_model_name):
     print("-> Download original vocabulary and config.")
 
-    extract_dir = EXTRACT_DIR.format(preset)
+    extract_dir = EXTRACT_DIR.format(FLAGS.preset)
     if not os.path.exists(extract_dir):
         os.makedirs(extract_dir)
 
@@ -70,9 +68,9 @@ def download_files(preset, hf_model_name):
     print(f"`{spm_path}`")
 
 
-def define_preprocessor(preset, hf_model_name):
+def define_preprocessor(hf_model_name):
     print("\n-> Define the tokenizers.")
-    extract_dir = EXTRACT_DIR.format(preset)
+    extract_dir = EXTRACT_DIR.format(FLAGS.preset)
     spm_path = os.path.join(extract_dir, "spm.model")
 
     keras_nlp_tokenizer = DebertaV3Tokenizer(proto=spm_path)
@@ -81,7 +79,7 @@ def define_preprocessor(preset, hf_model_name):
     # padding token may be vastly different from the representations computed in
     # the original model. See https://github.com/keras-team/keras/pull/16619#issuecomment-1156338394.
     sequence_length = 14
-    if preset == "deberta_v3_base_multi":
+    if FLAGS.preset == "deberta_v3_base_multi":
         sequence_length = 17
     keras_nlp_preprocessor = DebertaV3Preprocessor(
         keras_nlp_tokenizer, sequence_length=sequence_length
@@ -95,10 +93,10 @@ def define_preprocessor(preset, hf_model_name):
     return keras_nlp_preprocessor, hf_tokenizer
 
 
-def convert_checkpoints(preset, keras_nlp_model, hf_model):
+def convert_checkpoints(keras_nlp_model, hf_model):
     print("\n-> Convert original weights to KerasNLP format.")
 
-    extract_dir = EXTRACT_DIR.format(preset)
+    extract_dir = EXTRACT_DIR.format(FLAGS.preset)
     config_path = os.path.join(extract_dir, "config.json")
 
     # Build config.
@@ -256,14 +254,13 @@ def convert_checkpoints(preset, keras_nlp_model, hf_model):
         )
 
     # Save the model.
-    print(f"\n-> Save KerasNLP model weights to `{preset}.h5`.")
-    keras_nlp_model.save_weights(f"{preset}.h5")
+    print(f"\n-> Save KerasNLP model weights to `{FLAGS.preset}.h5`.")
+    keras_nlp_model.save_weights(f"{FLAGS.preset}.h5")
 
     return keras_nlp_model
 
 
 def check_output(
-    preset,
     keras_nlp_preprocessor,
     keras_nlp_model,
     hf_tokenizer,
@@ -287,17 +284,15 @@ def check_output(
     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"))
+    print("Model md5sum: ", get_md5_checksum(f"./{FLAGS.preset}.h5"))
 
 
 def main(_):
     hf_model_name = PRESET_MAP[FLAGS.preset]
 
-    download_files(FLAGS.preset, hf_model_name)
+    download_files(hf_model_name)
 
-    keras_nlp_preprocessor, hf_tokenizer = define_preprocessor(
-        FLAGS.preset, hf_model_name
-    )
+    keras_nlp_preprocessor, hf_tokenizer = define_preprocessor(hf_model_name)
 
     print("\n-> Load KerasNLP model.")
     keras_nlp_model = DebertaV3Backbone.from_preset(
@@ -308,12 +303,9 @@ def main(_):
     hf_model = transformers.AutoModel.from_pretrained(hf_model_name)
     hf_model.eval()
 
-    keras_nlp_model = convert_checkpoints(
-        FLAGS.preset, keras_nlp_model, hf_model
-    )
+    keras_nlp_model = convert_checkpoints(keras_nlp_model, hf_model)
 
     check_output(
-        FLAGS.preset,
         keras_nlp_preprocessor,
         keras_nlp_model,
         hf_tokenizer,
diff --git a/tools/checkpoint_conversion/convert_distilbert_checkpoints.py b/tools/checkpoint_conversion/convert_distilbert_checkpoints.py
index 4280eede..80af89eb 100644
--- a/tools/checkpoint_conversion/convert_distilbert_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_distilbert_checkpoints.py
@@ -40,10 +40,10 @@ flags.DEFINE_string(
 )
 
 
-def download_files(preset, hf_model_name):
+def download_files(hf_model_name):
     print("-> Download original vocab and config.")
 
-    extract_dir = EXTRACT_DIR.format(preset)
+    extract_dir = EXTRACT_DIR.format(FLAGS.preset)
     if not os.path.exists(extract_dir):
         os.makedirs(extract_dir)
 
@@ -64,9 +64,9 @@ def download_files(preset, hf_model_name):
     print(f"`{vocab_path}`")
 
 
-def define_preprocessor(preset, hf_model_name):
+def define_preprocessor(hf_model_name):
     print("\n-> Define the tokenizers.")
-    extract_dir = EXTRACT_DIR.format(preset)
+    extract_dir = EXTRACT_DIR.format(FLAGS.preset)
     vocab_path = os.path.join(extract_dir, "vocab.txt")
 
     keras_nlp_tokenizer = keras_nlp.models.DistilBertTokenizer(
@@ -84,10 +84,10 @@ def define_preprocessor(preset, hf_model_name):
     return keras_nlp_preprocessor, hf_tokenizer
 
 
-def convert_checkpoints(preset, keras_nlp_model, hf_model):
+def convert_checkpoints(keras_nlp_model, hf_model):
     print("\n-> Convert original weights to KerasNLP format.")
 
-    extract_dir = EXTRACT_DIR.format(preset)
+    extract_dir = EXTRACT_DIR.format(FLAGS.preset)
     config_path = os.path.join(extract_dir, "config.json")
 
     # Build config.
@@ -244,14 +244,13 @@ def convert_checkpoints(preset, keras_nlp_model, hf_model):
         )
 
     # Save the model.
-    print(f"\n-> Save KerasNLP model weights to `{preset}.h5`.")
-    keras_nlp_model.save_weights(f"{preset}.h5")
+    print(f"\n-> Save KerasNLP model weights to `{FLAGS.preset}.h5`.")
+    keras_nlp_model.save_weights(f"{FLAGS.preset}.h5")
 
     return keras_nlp_model
 
 
 def check_output(
-    preset,
     keras_nlp_preprocessor,
     keras_nlp_model,
     hf_tokenizer,
@@ -275,17 +274,15 @@ def check_output(
     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"))
+    print("Model md5sum: ", get_md5_checksum(f"./{FLAGS.preset}.h5"))
 
 
 def main(_):
     hf_model_name = PRESET_MAP[FLAGS.preset]
 
-    download_files(FLAGS.preset, hf_model_name)
+    download_files(hf_model_name)
 
-    keras_nlp_preprocessor, hf_tokenizer = define_preprocessor(
-        FLAGS.preset, hf_model_name
-    )
+    keras_nlp_preprocessor, hf_tokenizer = define_preprocessor(hf_model_name)
 
     print("\n-> Load KerasNLP model.")
     keras_nlp_model = keras_nlp.models.DistilBertBackbone.from_preset(
@@ -296,12 +293,9 @@ def main(_):
     hf_model = transformers.AutoModel.from_pretrained(hf_model_name)
     hf_model.eval()
 
-    keras_nlp_model = convert_checkpoints(
-        FLAGS.preset, keras_nlp_model, hf_model
-    )
+    keras_nlp_model = convert_checkpoints(keras_nlp_model, hf_model)
 
     check_output(
-        FLAGS.preset,
         keras_nlp_preprocessor,
         keras_nlp_model,
         hf_tokenizer,
diff --git a/tools/checkpoint_conversion/convert_f_net_checkpoints.py b/tools/checkpoint_conversion/convert_f_net_checkpoints.py
index 68e52057..a5800203 100644
--- a/tools/checkpoint_conversion/convert_f_net_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_f_net_checkpoints.py
@@ -19,11 +19,9 @@ 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
-from tools.checkpoint_conversion.checkpoint_conversion_utils import (
-    get_md5_checksum,
-)
 
 PRESET_MAP = {
     "f_net_base_en": "google/fnet-base",
diff --git a/tools/checkpoint_conversion/convert_gpt2_checkpoints.py b/tools/checkpoint_conversion/convert_gpt2_checkpoints.py
index 93df3c42..79e67eef 100644
--- a/tools/checkpoint_conversion/convert_gpt2_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_gpt2_checkpoints.py
@@ -20,19 +20,17 @@ import tensorflow as tf
 import transformers
 from absl import app
 from absl import flags
+from checkpoint_conversion_utils import get_md5_checksum
 
 # Temporarily directly import gpt2 until we expose it.
 from keras_nlp.models.gpt2.gpt2_backbone import GPT2Backbone
 from keras_nlp.models.gpt2.gpt2_tokenizer import GPT2Tokenizer
-from tools.checkpoint_conversion.checkpoint_conversion_utils import (
-    get_md5_checksum,
-)
 
 PRESET_MAP = {
-    "gpt2_base": ("124M", "gpt2"),
-    "gpt2_medium": ("355M", "gpt2-medium"),
-    "gpt2_large": ("774M", "gpt2-large"),
-    "gpt2_extra_large": ("1558M", "gpt2-xl"),
+    "gpt2_base_en": ("124M", "gpt2"),
+    "gpt2_medium_en": ("355M", "gpt2-medium"),
+    "gpt2_large_en": ("774M", "gpt2-large"),
+    "gpt2_extra_large_en": ("1558M", "gpt2-xl"),
 }
 
 DOWNLOAD_SCRIPT_URL = (
@@ -47,7 +45,7 @@ flags.DEFINE_string(
 )
 
 
-def download_model(preset, num_params):
+def download_model(num_params):
     print("-> Download original weights.")
     response = requests.get(DOWNLOAD_SCRIPT_URL)
     open("download_model.py", "wb").write(response.content)
@@ -55,7 +53,7 @@ def download_model(preset, num_params):
     os.system(f"python download_model.py {num_params}")
 
 
-def convert_checkpoints(preset, num_params):
+def convert_checkpoints(num_params):
     print("\n-> Convert original weights to KerasNLP format.")
     # GPT-2 paths.
     extract_dir = EXTRACT_DIR.format(num_params)
@@ -76,7 +74,7 @@ def convert_checkpoints(preset, num_params):
 
     # Temporary direct import, as we aren't exposing this quite yet.
     keras_nlp_model = GPT2Backbone.from_preset(
-        preset,
+        FLAGS.preset,
         load_weights=False,
     )
 
@@ -191,13 +189,13 @@ def convert_checkpoints(preset, num_params):
     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")
+    print(f"\n-> Save KerasNLP model weights to `{FLAGS.preset}.h5`.")
+    keras_nlp_model.save_weights(f"{FLAGS.preset}.h5")
 
     return keras_nlp_model
 
 
-def define_tokenizer(preset, num_params, hf_model_name):
+def define_tokenizer(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")
@@ -217,7 +215,6 @@ def define_tokenizer(preset, num_params, hf_model_name):
 
 
 def check_output(
-    preset,
     keras_nlp_model,
     keras_nlp_tokenizer,
     hf_model,
@@ -245,7 +242,7 @@ def check_output(
     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"))
+    print("Model md5sum: ", get_md5_checksum(f"./{FLAGS.preset}.h5"))
 
     return keras_nlp_output
 
@@ -257,20 +254,19 @@ def main(_):
     num_params = PRESET_MAP[FLAGS.preset][0]
     hf_model_name = PRESET_MAP[FLAGS.preset][1]
 
-    download_model(FLAGS.preset, num_params)
+    download_model(num_params)
 
-    keras_nlp_model = convert_checkpoints(FLAGS.preset, num_params)
+    keras_nlp_model = convert_checkpoints(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
+        num_params, hf_model_name
     )
 
     check_output(
-        FLAGS.preset,
         keras_nlp_model,
         keras_nlp_tokenizer,
         hf_model,
diff --git a/tools/checkpoint_conversion/convert_opt_checkpoints.py b/tools/checkpoint_conversion/convert_opt_checkpoints.py
index 7a5a867e..2f577bfb 100644
--- a/tools/checkpoint_conversion/convert_opt_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_opt_checkpoints.py
@@ -19,11 +19,9 @@ 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
-from tools.checkpoint_conversion.checkpoint_conversion_utils import (
-    get_md5_checksum,
-)
 
 PRESET_MAP = {
     "opt_125m_en": "facebook/opt-125m",
diff --git a/tools/checkpoint_conversion/convert_roberta_checkpoints.py b/tools/checkpoint_conversion/convert_roberta_checkpoints.py
index ae0900a6..092eb470 100644
--- a/tools/checkpoint_conversion/convert_roberta_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_roberta_checkpoints.py
@@ -20,16 +20,15 @@ import torch
 import transformers
 from absl import app
 from absl import flags
+from checkpoint_conversion_utils import extract_files_from_archive
+from checkpoint_conversion_utils import get_md5_checksum
 from tensorflow import keras
 
 import keras_nlp
-from tools.checkpoint_conversion.checkpoint_conversion_utils import (
-    get_md5_checksum,
-)
 
 PRESET_MAP = {
-    "roberta_base": ("roberta.base", "roberta-base"),
-    "roberta_large": ("roberta.large", "roberta-large"),
+    "roberta_base_en": ("roberta.base", "roberta-base"),
+    "roberta_large_en": ("roberta.large", "roberta-large"),
 }
 
 DOWNLOAD_SCRIPT_URL = "https://dl.fbaipublicfiles.com/fairseq/models/{}.tar.gz"
@@ -42,16 +41,16 @@ flags.DEFINE_string(
 )
 
 
-def download_model(preset, size, hf_model_name):
+def download_model(size, hf_model_name):
     print("-> Download original weights.")
     extract_dir = EXTRACT_DIR.format(size)
     archive_file_path = keras.utils.get_file(
         fname=None,
         origin=DOWNLOAD_SCRIPT_URL.format(size),
-        cache_subdir=os.path.join("checkpoint_conversion", preset),
+        cache_subdir=os.path.join("checkpoint_conversion", FLAGS.preset),
     )
 
-    os.system(f"tar -xvf {archive_file_path}")
+    extract_files_from_archive(archive_file_path)
 
     # The original `tar.gz` file does not have the vocab files. Let's fetch
     # them from HF.
@@ -67,7 +66,7 @@ def download_model(preset, size, hf_model_name):
     shutil.copy(merges_path, extract_dir)
 
 
-def convert_checkpoints(preset, size):
+def convert_checkpoints(size):
     print("\n-> Convert original weights to KerasNLP format.")
     # RoBERTa paths.
     extract_dir = EXTRACT_DIR.format(size)
@@ -94,7 +93,7 @@ def convert_checkpoints(preset, size):
     print("Config:", cfg)
 
     keras_nlp_model = keras_nlp.models.RobertaBackbone.from_preset(
-        preset, load_weights=False
+        FLAGS.preset, load_weights=False
     )
 
     # Embedding Layer.
@@ -265,13 +264,13 @@ def convert_checkpoints(preset, size):
         )
 
     # Save the model.
-    print(f"\n-> Save KerasNLP model weights to `{preset}.h5`.")
-    keras_nlp_model.save_weights(f"{preset}.h5")
+    print(f"\n-> Save KerasNLP model weights to `{FLAGS.preset}.h5`.")
+    keras_nlp_model.save_weights(f"{FLAGS.preset}.h5")
 
     return keras_nlp_model
 
 
-def define_preprocessor(preset, hf_model_name, size):
+def define_preprocessor(hf_model_name, size):
     print("\n-> Define the tokenizers.")
     extract_dir = EXTRACT_DIR.format(size)
     vocabulary_path = os.path.join(extract_dir, "vocab.json")
@@ -294,7 +293,6 @@ def define_preprocessor(preset, hf_model_name, size):
 
 
 def check_output(
-    preset,
     keras_nlp_model,
     keras_nlp_preprocessor,
     hf_model,
@@ -318,7 +316,7 @@ def check_output(
     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"))
+    print("Model md5sum: ", get_md5_checksum(f"./{FLAGS.preset}.h5"))
 
     return keras_nlp_output
 
@@ -330,20 +328,19 @@ def main(_):
     size = PRESET_MAP[FLAGS.preset][0]
     hf_model_name = PRESET_MAP[FLAGS.preset][1]
 
-    download_model(FLAGS.preset, size, hf_model_name)
+    download_model(size, hf_model_name)
 
-    keras_nlp_model = convert_checkpoints(FLAGS.preset, size)
+    keras_nlp_model = convert_checkpoints(size)
 
     print("\n-> Load HF model.")
     hf_model = transformers.AutoModel.from_pretrained(hf_model_name)
     hf_model.eval()
 
     keras_nlp_preprocessor, hf_tokenizer = define_preprocessor(
-        FLAGS.preset, hf_model_name, size
+        hf_model_name, size
     )
 
     check_output(
-        FLAGS.preset,
         keras_nlp_model,
         keras_nlp_preprocessor,
         hf_model,
diff --git a/tools/checkpoint_conversion/convert_xlm_roberta_checkpoints.py b/tools/checkpoint_conversion/convert_xlm_roberta_checkpoints.py
index 1705836f..40cf3862 100644
--- a/tools/checkpoint_conversion/convert_xlm_roberta_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_xlm_roberta_checkpoints.py
@@ -19,16 +19,15 @@ import torch
 import transformers
 from absl import app
 from absl import flags
+from checkpoint_conversion_utils import extract_files_from_archive
+from checkpoint_conversion_utils import get_md5_checksum
 from tensorflow import keras
 
 import keras_nlp
-from tools.checkpoint_conversion.checkpoint_conversion_utils import (
-    get_md5_checksum,
-)
 
 PRESET_MAP = {
-    "xlm_roberta_base": ("xlmr.base", "xlm-roberta-base"),
-    "xlm_roberta_large": ("xlmr.large", "xlm-roberta-large"),
+    "xlm_roberta_base_multi": ("xlmr.base", "xlm-roberta-base"),
+    "xlm_roberta_large_multi": ("xlmr.large", "xlm-roberta-large"),
 }
 
 DOWNLOAD_SCRIPT_URL = "https://dl.fbaipublicfiles.com/fairseq/models/{}.tar.gz"
@@ -41,18 +40,18 @@ flags.DEFINE_string(
 )
 
 
-def download_model(preset, size):
+def download_model(size):
     print("-> Download original weights.")
     archive_file_path = keras.utils.get_file(
         fname=None,
         origin=DOWNLOAD_SCRIPT_URL.format(size),
-        cache_subdir=os.path.join("checkpoint_conversion", preset),
+        cache_subdir=os.path.join("checkpoint_conversion", FLAGS.preset),
     )
 
-    os.system(f"tar -xvf {archive_file_path}")
+    extract_files_from_archive(archive_file_path)
 
 
-def convert_checkpoints(preset, size):
+def convert_checkpoints(size):
     print("\n-> Convert original weights to KerasNLP format.")
     # XLM-RoBERTa paths.
     extract_dir = EXTRACT_DIR.format(size)
@@ -79,7 +78,7 @@ def convert_checkpoints(preset, size):
     print("Config:", cfg)
 
     keras_nlp_model = keras_nlp.models.XLMRobertaBackbone.from_preset(
-        preset, load_weights=False
+        FLAGS.preset, load_weights=False
     )
 
     # Embedding Layer.
@@ -246,13 +245,13 @@ def convert_checkpoints(preset, size):
         )
 
     # Save the model.
-    print(f"\n-> Save KerasNLP model weights to `{preset}.h5`.")
-    keras_nlp_model.save_weights(f"{preset}.h5")
+    print(f"\n-> Save KerasNLP model weights to `{FLAGS.preset}.h5`.")
+    keras_nlp_model.save_weights(f"{FLAGS.preset}.h5")
 
     return keras_nlp_model
 
 
-def define_preprocessor(preset, hf_model_name, size):
+def define_preprocessor(hf_model_name, size):
     print("\n-> Define the tokenizers.")
     extract_dir = EXTRACT_DIR.format(size)
     spm_path = os.path.join(extract_dir, "sentencepiece.bpe.model")
@@ -273,7 +272,6 @@ def define_preprocessor(preset, hf_model_name, size):
 
 
 def check_output(
-    preset,
     keras_nlp_model,
     keras_nlp_preprocessor,
     hf_model,
@@ -297,7 +295,7 @@ def check_output(
     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"))
+    print("Model md5sum: ", get_md5_checksum(f"./{FLAGS.preset}.h5"))
 
     return keras_nlp_output
 
@@ -309,20 +307,19 @@ def main(_):
     size = PRESET_MAP[FLAGS.preset][0]
     hf_model_name = PRESET_MAP[FLAGS.preset][1]
 
-    download_model(FLAGS.preset, size)
+    download_model(size)
 
-    keras_nlp_model = convert_checkpoints(FLAGS.preset, size)
+    keras_nlp_model = convert_checkpoints(size)
 
     print("\n-> Load HF model.")
     hf_model = transformers.AutoModel.from_pretrained(hf_model_name)
     hf_model.eval()
 
     keras_nlp_preprocessor, hf_tokenizer = define_preprocessor(
-        FLAGS.preset, hf_model_name, size
+        hf_model_name, size
     )
 
     check_output(
-        FLAGS.preset,
         keras_nlp_model,
         keras_nlp_preprocessor,
         hf_model,


Test patch:
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_backbone_test.py b/keras_nlp/models/deberta_v3/deberta_v3_backbone_test.py
index 936773ce..931410a5 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_backbone_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_backbone_test.py
@@ -93,3 +93,30 @@ class DebertaV3BackboneTest(tf.test.TestCase, parameterized.TestCase):
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
         self.assertAllClose(model_output, restored_output)
+
+
+@pytest.mark.tpu
+@pytest.mark.usefixtures("tpu_test_class")
+class DebertaV3BackboneTPUTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        with self.tpu_strategy.scope():
+            self.backbone = DebertaV3Backbone(
+                vocabulary_size=10,
+                num_layers=2,
+                num_heads=2,
+                hidden_dim=2,
+                intermediate_dim=4,
+                max_sequence_length=5,
+                bucket_size=2,
+            )
+        self.input_batch = {
+            "token_ids": tf.ones((2, 5), dtype="int32"),
+            "padding_mask": tf.ones((2, 5), dtype="int32"),
+        }
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
+    def test_predict(self):
+        self.backbone.compile()
+        self.backbone.predict(self.input_dataset)
diff --git a/keras_nlp/models/f_net/f_net_backbone_test.py b/keras_nlp/models/f_net/f_net_backbone_test.py
index 08720afa..47b0ec9c 100644
--- a/keras_nlp/models/f_net/f_net_backbone_test.py
+++ b/keras_nlp/models/f_net/f_net_backbone_test.py
@@ -25,22 +25,17 @@ from keras_nlp.models.f_net.f_net_backbone import FNetBackbone
 
 class FNetBackboneTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = FNetBackbone(
-            vocabulary_size=1000,
+        self.backbone = FNetBackbone(
+            vocabulary_size=10,
             num_layers=2,
-            hidden_dim=16,
-            intermediate_dim=32,
-            max_sequence_length=128,
+            hidden_dim=2,
+            intermediate_dim=4,
+            max_sequence_length=5,
             num_segments=4,
         )
-        self.batch_size = 8
         self.input_batch = {
-            "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"
-            ),
+            "token_ids": tf.ones((2, 5), dtype="int32"),
+            "segment_ids": tf.ones((2, 5), dtype="int32"),
         }
 
         self.input_dataset = tf.data.Dataset.from_tensor_slices(
@@ -48,45 +43,38 @@ class FNetBackboneTest(tf.test.TestCase, parameterized.TestCase):
         ).batch(2)
 
     def test_valid_call_f_net(self):
-        self.model(self.input_batch)
+        self.backbone(self.input_batch)
 
         # Check default name passed through
-        self.assertRegexpMatches(self.model.name, "f_net_backbone")
+        self.assertRegexpMatches(self.backbone.name, "f_net_backbone")
 
     def test_variable_sequence_length_call_f_net(self):
-        for seq_length in (25, 50, 75):
+        for seq_length in (2, 3, 4):
             input_data = {
-                "token_ids": tf.ones(
-                    (self.batch_size, seq_length), dtype="int32"
-                ),
-                "segment_ids": tf.ones(
-                    (self.batch_size, seq_length), dtype="int32"
-                ),
+                "token_ids": tf.ones((2, seq_length), dtype="int32"),
+                "segment_ids": tf.ones((2, seq_length), dtype="int32"),
             }
-            self.model(input_data)
+            self.backbone(input_data)
 
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_compile(self, jit_compile):
-        self.model.compile(jit_compile=jit_compile)
-        self.model.predict(self.input_batch)
+    def test_predict(self):
+        self.backbone.predict(self.input_batch)
+        self.backbone.predict(self.input_dataset)
 
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_compile_batched_ds(self, jit_compile):
-        self.model.compile(jit_compile=jit_compile)
-        self.model.predict(self.input_dataset)
+    def test_serialization(self):
+        new_backbone = keras.utils.deserialize_keras_object(
+            keras.utils.serialize_keras_object(self.backbone)
+        )
+        self.assertEqual(new_backbone.get_config(), self.backbone.get_config())
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
     )
+    @pytest.mark.large
     def test_saved_model(self, save_format, filename):
-        model_output = self.model(self.input_batch)
+        model_output = self.backbone(self.input_batch)
         save_path = os.path.join(self.get_temp_dir(), filename)
-        self.model.save(save_path, save_format=save_format)
+        self.backbone.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
@@ -104,8 +92,8 @@ class FNetBackboneTest(tf.test.TestCase, parameterized.TestCase):
 class FNetBackboneTPUTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         with self.tpu_strategy.scope():
-            self.model = FNetBackbone(
-                vocabulary_size=1000,
+            self.backbone = FNetBackbone(
+                vocabulary_size=100,
                 num_layers=2,
                 hidden_dim=16,
                 intermediate_dim=32,
@@ -121,5 +109,5 @@ class FNetBackboneTPUTest(tf.test.TestCase, parameterized.TestCase):
         ).batch(2)
 
     def test_predict(self):
-        self.model.compile()
-        self.model.predict(self.input_dataset)
+        self.backbone.compile()
+        self.backbone.predict(self.input_dataset)
diff --git a/keras_nlp/models/f_net/f_net_classifier_test.py b/keras_nlp/models/f_net/f_net_classifier_test.py
index cb91535a..6f3e01d0 100644
--- a/keras_nlp/models/f_net/f_net_classifier_test.py
+++ b/keras_nlp/models/f_net/f_net_classifier_test.py
@@ -16,6 +16,7 @@
 import io
 import os
 
+import pytest
 import sentencepiece
 import tensorflow as tf
 from absl.testing import parameterized
@@ -29,15 +30,7 @@ from keras_nlp.models.f_net.f_net_tokenizer import FNetTokenizer
 
 class FNetClassifierTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.backbone = FNetBackbone(
-            vocabulary_size=1000,
-            num_layers=2,
-            hidden_dim=64,
-            intermediate_dim=128,
-            max_sequence_length=128,
-            name="encoder",
-        )
-
+        # Setup Model
         bytes_io = io.BytesIO()
         vocab_data = tf.data.Dataset.from_tensor_slices(
             ["the quick brown fox", "the earth is round"]
@@ -63,78 +56,68 @@ class FNetClassifierTest(tf.test.TestCase, parameterized.TestCase):
 
         self.preprocessor = FNetPreprocessor(
             tokenizer=FNetTokenizer(proto=self.proto),
-            sequence_length=12,
+            sequence_length=8,
+        )
+        self.backbone = FNetBackbone(
+            vocabulary_size=self.preprocessor.tokenizer.vocabulary_size(),
+            num_layers=2,
+            hidden_dim=2,
+            intermediate_dim=4,
+            max_sequence_length=self.preprocessor.packer.sequence_length,
         )
-
         self.classifier = FNetClassifier(
             self.backbone,
             4,
             preprocessor=self.preprocessor,
         )
-        self.classifier_no_preprocessing = FNetClassifier(
-            self.backbone,
-            4,
-            preprocessor=None,
-        )
 
+        # Setup data.
         self.raw_batch = tf.constant(
             [
                 "the quick brown fox.",
                 "the slow brown fox.",
-                "the smelly brown fox.",
-                "the old brown fox.",
             ]
         )
         self.preprocessed_batch = self.preprocessor(self.raw_batch)
         self.raw_dataset = tf.data.Dataset.from_tensor_slices(
-            (self.raw_batch, tf.ones((4,)))
+            (self.raw_batch, tf.ones((2,)))
         ).batch(2)
         self.preprocessed_dataset = self.raw_dataset.map(self.preprocessor)
 
     def test_valid_call_classifier(self):
         self.classifier(self.preprocessed_batch)
 
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_fnet_classifier_predict(self, jit_compile):
-        self.classifier.compile(jit_compile=jit_compile)
+    def test_classifier_predict(self):
         self.classifier.predict(self.raw_batch)
+        self.classifier.preprocessor = None
+        self.classifier.predict(self.preprocessed_batch)
 
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_fnet_classifier_predict_no_preprocessing(self, jit_compile):
-        self.classifier_no_preprocessing.compile(jit_compile=jit_compile)
-        self.classifier_no_preprocessing.predict(self.preprocessed_batch)
-
-    def test_fnet_classifier_fit_default_compile(self):
+    def test_fnet_classifier_fit(self):
         self.classifier.fit(self.raw_dataset)
+        self.classifier.preprocessor = None
+        self.classifier.fit(self.preprocessed_dataset)
 
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_fnet_classifier_fit(self, jit_compile):
+    def test_classifier_fit_no_xla(self):
+        self.classifier.preprocessor = None
         self.classifier.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
-            jit_compile=jit_compile,
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=False),
+            jit_compile=False,
         )
-        self.classifier.fit(self.raw_dataset)
-
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_fnet_classifier_fit_no_preprocessing(self, jit_compile):
-        self.classifier_no_preprocessing.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
-            jit_compile=jit_compile,
+        self.classifier.fit(self.preprocessed_dataset)
+
+    def test_serialization(self):
+        config = keras.utils.serialize_keras_object(self.classifier)
+        new_classifier = keras.utils.deserialize_keras_object(config)
+        self.assertEqual(
+            new_classifier.get_config(),
+            self.classifier.get_config(),
         )
-        self.classifier_no_preprocessing.fit(self.preprocessed_dataset)
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
     )
+    @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         model_output = self.classifier.predict(self.raw_batch)
         save_path = os.path.join(self.get_temp_dir(), filename)
diff --git a/keras_nlp/models/f_net/f_net_masked_lm_preprocessor_test.py b/keras_nlp/models/f_net/f_net_masked_lm_preprocessor_test.py
index b95b4a7a..600afc09 100644
--- a/keras_nlp/models/f_net/f_net_masked_lm_preprocessor_test.py
+++ b/keras_nlp/models/f_net/f_net_masked_lm_preprocessor_test.py
@@ -14,6 +14,7 @@
 import io
 import os
 
+import pytest
 import sentencepiece
 import tensorflow as tf
 from absl.testing import parameterized
@@ -120,10 +121,19 @@ class FNetMaskedLMPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(y, [0, 0, 0, 0])
         self.assertAllEqual(sw, [0.0, 0.0, 0.0, 0.0])
 
+    def test_serialization(self):
+        config = keras.utils.serialize_keras_object(self.preprocessor)
+        new_preprocessor = keras.utils.deserialize_keras_object(config)
+        self.assertEqual(
+            new_preprocessor.get_config(),
+            self.preprocessor.get_config(),
+        )
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
     )
+    @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         input_data = tf.constant(["the quick brown fox"])
 
diff --git a/keras_nlp/models/f_net/f_net_masked_lm_test.py b/keras_nlp/models/f_net/f_net_masked_lm_test.py
index 3e7e7c9c..508b3f3e 100644
--- a/keras_nlp/models/f_net/f_net_masked_lm_test.py
+++ b/keras_nlp/models/f_net/f_net_masked_lm_test.py
@@ -14,6 +14,7 @@
 import io
 import os
 
+import pytest
 import sentencepiece
 import tensorflow as tf
 from absl.testing import parameterized
@@ -29,21 +30,15 @@ from keras_nlp.models.f_net.f_net_tokenizer import FNetTokenizer
 
 class FNetMaskedLMTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.backbone = FNetBackbone(
-            vocabulary_size=1000,
-            num_layers=2,
-            hidden_dim=16,
-            intermediate_dim=32,
-            max_sequence_length=128,
-        )
+        # Setup Model.
         bytes_io = io.BytesIO()
         vocab_data = tf.data.Dataset.from_tensor_slices(
-            ["the quick brown fox", "the earth is round", "an eagle flew"]
+            ["the quick brown fox", "the slow brown fox"]
         )
         sentencepiece.SentencePieceTrainer.train(
             sentence_iterator=vocab_data.as_numpy_iterator(),
             model_writer=bytes_io,
-            vocab_size=15,
+            vocab_size=5,
             model_type="WORD",
             pad_id=0,
             bos_id=1,
@@ -58,24 +53,25 @@ class FNetMaskedLMTest(tf.test.TestCase, parameterized.TestCase):
         self.proto = bytes_io.getvalue()
         self.preprocessor = FNetMaskedLMPreprocessor(
             FNetTokenizer(proto=self.proto),
-            sequence_length=10,
-            mask_selection_length=4,
+            sequence_length=5,
+            mask_selection_length=2,
+        )
+        self.backbone = FNetBackbone(
+            vocabulary_size=self.preprocessor.tokenizer.vocabulary_size(),
+            num_layers=2,
+            hidden_dim=2,
+            intermediate_dim=4,
+            max_sequence_length=self.preprocessor.packer.sequence_length,
         )
         self.masked_lm = FNetMaskedLM(
             self.backbone,
             preprocessor=self.preprocessor,
         )
-        self.masked_lm_no_preprocessing = FNetMaskedLM(
-            self.backbone,
-            preprocessor=None,
-        )
 
         self.raw_batch = tf.constant(
             [
-                "quick brown fox",
-                "eagle flew over fox",
-                "the eagle flew quick",
-                "a brown eagle",
+                "the quick brown fox",
+                "the slow brown fox",
             ]
         )
         self.preprocessed_batch = self.preprocessor(self.raw_batch)[0]
@@ -84,51 +80,42 @@ class FNetMaskedLMTest(tf.test.TestCase, parameterized.TestCase):
         ).batch(2)
         self.preprocessed_dataset = self.raw_dataset.map(self.preprocessor)
 
-    def test_valid_call_masked_lm(self):
+    def test_valid_call_classifier(self):
         self.masked_lm(self.preprocessed_batch)
 
-    def test_fnet_masked_lm_fit_default_compile(self):
-        self.masked_lm.fit(self.raw_dataset)
+    def test_classifier_predict(self):
+        # self.masked_lm.predict(self.raw_batch)
+        self.masked_lm.preprocessor = None
+        self.masked_lm.predict(self.preprocessed_batch)
 
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_f_net_masked_lm_predict(self, jit_compile):
-        self.masked_lm.compile(jit_compile=jit_compile)
-        self.masked_lm.predict(self.raw_batch)
-
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_f_net_masked_lm_predict_no_preprocessing(self, jit_compile):
-        self.masked_lm_no_preprocessing.compile(jit_compile=jit_compile)
-        self.masked_lm_no_preprocessing.predict(self.preprocessed_batch)
+    def test_classifier_fit(self):
+        self.masked_lm.fit(self.raw_dataset)
+        self.masked_lm.preprocessor = None
+        self.masked_lm.fit(self.preprocessed_dataset)
 
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_f_net_masked_lm_fit(self, jit_compile):
+    def test_classifier_fit_no_xla(self):
+        self.masked_lm.preprocessor = None
         self.masked_lm.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
-            jit_compile=jit_compile,
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=False),
+            jit_compile=False,
         )
-        self.masked_lm.fit(self.raw_dataset)
+        self.masked_lm.fit(self.preprocessed_dataset)
 
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_f_net_masked_lm_fit_no_preprocessing(self, jit_compile):
-        self.masked_lm_no_preprocessing.compile(
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
-            jit_compile=jit_compile,
+    def test_serialization(self):
+        config = keras.utils.serialize_keras_object(self.masked_lm)
+        new_classifier = keras.utils.deserialize_keras_object(config)
+        self.assertEqual(
+            new_classifier.get_config(),
+            self.masked_lm.get_config(),
         )
-        self.masked_lm_no_preprocessing.fit(self.preprocessed_dataset)
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
     )
+    @pytest.mark.large
     def test_saved_model(self, save_format, filename):
+        model_output = self.masked_lm.predict(self.raw_batch)
         save_path = os.path.join(self.get_temp_dir(), filename)
         self.masked_lm.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(save_path)
@@ -136,7 +123,6 @@ class FNetMaskedLMTest(tf.test.TestCase, parameterized.TestCase):
         # Check we got the real object back.
         self.assertIsInstance(restored_model, FNetMaskedLM)
 
-        model_output = self.masked_lm(self.preprocessed_batch)
-        restored_output = restored_model(self.preprocessed_batch)
-
+        # Check that output matches.
+        restored_output = restored_model.predict(self.raw_batch)
         self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/models/f_net/f_net_preprocessor_test.py b/keras_nlp/models/f_net/f_net_preprocessor_test.py
index 8812ae3f..230247e8 100644
--- a/keras_nlp/models/f_net/f_net_preprocessor_test.py
+++ b/keras_nlp/models/f_net/f_net_preprocessor_test.py
@@ -17,6 +17,7 @@
 import io
 import os
 
+import pytest
 import sentencepiece
 import tensorflow as tf
 from absl.testing import parameterized
@@ -139,10 +140,19 @@ class FNetPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         with self.assertRaises(ValueError):
             self.preprocessor(ambiguous_input)
 
+    def test_serialization(self):
+        config = keras.utils.serialize_keras_object(self.preprocessor)
+        new_preprocessor = keras.utils.deserialize_keras_object(config)
+        self.assertEqual(
+            new_preprocessor.get_config(),
+            self.preprocessor.get_config(),
+        )
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
     )
+    @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         input_data = tf.constant(["the quick brown fox"])
         inputs = keras.Input(dtype="string", shape=())
diff --git a/keras_nlp/models/f_net/f_net_tokenizer_test.py b/keras_nlp/models/f_net/f_net_tokenizer_test.py
index 3abb0d06..15673a04 100644
--- a/keras_nlp/models/f_net/f_net_tokenizer_test.py
+++ b/keras_nlp/models/f_net/f_net_tokenizer_test.py
@@ -17,6 +17,7 @@
 import io
 import os
 
+import pytest
 import sentencepiece
 import tensorflow as tf
 from absl.testing import parameterized
@@ -82,10 +83,19 @@ class FNetTokenizerTest(tf.test.TestCase, parameterized.TestCase):
         with self.assertRaises(ValueError):
             FNetTokenizer(proto=bytes_io.getvalue())
 
+    def test_serialization(self):
+        config = keras.utils.serialize_keras_object(self.tokenizer)
+        new_tokenizer = keras.utils.deserialize_keras_object(config)
+        self.assertEqual(
+            new_tokenizer.get_config(),
+            self.tokenizer.get_config(),
+        )
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
     )
+    @pytest.mark.large
     def test_saved_model(self, save_format, filename):
         input_data = tf.constant(["the quick brown fox"])
 
diff --git a/keras_nlp/models/t5/t5_backbone_test.py b/keras_nlp/models/t5/t5_backbone_test.py
index 53c28395..813052b1 100644
--- a/keras_nlp/models/t5/t5_backbone_test.py
+++ b/keras_nlp/models/t5/t5_backbone_test.py
@@ -136,7 +136,7 @@ class T5BackboneTPUTest(tf.test.TestCase, parameterized.TestCase):
         ).batch(2)
 
     def test_predict(self):
-        self.model.compile()
+        self.backbone.compile()
         outputs = self.backbone.predict(self.input_dataset)
         self.assertIn("encoder_sequence_output", outputs)
         self.assertIn("decoder_sequence_output", outputs)
diff --git a/keras_nlp/models/task_test.py b/keras_nlp/models/task_test.py
new file mode 100644
index 00000000..72831a2c
--- /dev/null
+++ b/keras_nlp/models/task_test.py
@@ -0,0 +1,53 @@
+# Copyright 2023 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 tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.models.preprocessor import Preprocessor
+from keras_nlp.models.task import Task
+from keras_nlp.tokenizers.tokenizer import Tokenizer
+
+
+class SimpleTokenizer(Tokenizer):
+    def vocabulary_size(self):
+        return 10
+
+
+class SimplePreprocessor(Preprocessor):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.tokenizer = SimpleTokenizer()
+
+
+class SimpleTask(Task):
+    def __init__(self, preprocessor=None, **kwargs):
+        inputs = keras.Input(shape=(5,))
+        outputs = keras.layers.Dense(5)(inputs)
+        super().__init__(inputs, outputs, **kwargs)
+        self.preprocessor = preprocessor
+
+
+class TestTask(tf.test.TestCase):
+    def test_summary_with_preprocessor(self):
+        preprocessor = SimplePreprocessor()
+        model = SimpleTask(preprocessor)
+        summary = []
+        model.summary(print_fn=lambda x: summary.append(x))
+        self.assertRegex("\n".join(summary), "Preprocessor:")
+
+    def test_summary_without_preprocessor(self):
+        model = SimpleTask()
+        summary = []
+        model.summary(print_fn=lambda x: summary.append(x))
+        self.assertNotRegex("\n".join(summary), "Preprocessor:")
diff --git a/keras_nlp/models/whisper/whisper_backbone_test.py b/keras_nlp/models/whisper/whisper_backbone_test.py
index b71bbfda..456cb7b5 100644
--- a/keras_nlp/models/whisper/whisper_backbone_test.py
+++ b/keras_nlp/models/whisper/whisper_backbone_test.py
@@ -142,16 +142,16 @@ class WhisperBackboneTPUTest(tf.test.TestCase, parameterized.TestCase):
             "encoder_features": tf.ones(
                 (
                     8,
-                    self.model.max_encoder_sequence_length,
+                    self.backbone.max_encoder_sequence_length,
                     NUM_MELS,
                 ),
                 dtype="int32",
             ),
             "decoder_token_ids": tf.ones(
-                (8, self.model.max_decoder_sequence_length), dtype="int32"
+                (8, self.backbone.max_decoder_sequence_length), dtype="int32"
             ),
             "decoder_padding_mask": tf.ones(
-                (8, self.model.max_decoder_sequence_length), dtype="int32"
+                (8, self.backbone.max_decoder_sequence_length), dtype="int32"
             ),
         }
 
diff --git a/keras_nlp/samplers/beam_sampler_test.py b/keras_nlp/samplers/beam_sampler_test.py
index 95290b8f..e20c6603 100644
--- a/keras_nlp/samplers/beam_sampler_test.py
+++ b/keras_nlp/samplers/beam_sampler_test.py
@@ -37,6 +37,7 @@ class BeamSamplerTest(tf.test.TestCase, parameterized.TestCase):
 
         self.next = next
         self.sampler = BeamSampler(num_beams=5)
+        self.sampler_all_beams = BeamSampler(num_beams=5, return_all_beams=True)
 
     def join_as_string(self, x):
         return ["".join([self.int_lookup[i] for i in s]) for s in x.numpy()]
@@ -67,6 +68,27 @@ class BeamSamplerTest(tf.test.TestCase, parameterized.TestCase):
         )
         self.assertEqual(self.join_as_string(output), ["sequentially"])
 
+    def test_return_all_beams(self):
+        state_chars = list("sequentially")
+        state = tf.constant([[self.char_lookup[c] for c in state_chars]])
+        prompt = tf.fill((self.batch_size, self.length), self.char_lookup["z"])
+        sorted_prompts, sorted_log_probs = self.sampler_all_beams(
+            next=self.next,
+            prompt=prompt,
+            state=state,
+        )
+
+        self.assertEqual(
+            sorted_prompts.shape, (self.batch_size, 5, self.length)
+        )
+        self.assertEqual(sorted_log_probs.shape, (self.batch_size, 5))
+        self.assertTrue(
+            tf.reduce_all(sorted_log_probs[:, 1:] <= sorted_log_probs[:, :-1])
+        )
+        self.assertEqual(
+            self.join_as_string(sorted_prompts[:, 0, :]), ["sequentially"]
+        )
+
     def test_early_stopping(self):
         state_chars = list("sequentially")
         state = tf.constant([[self.char_lookup[c] for c in state_chars]])