Repository: keras-team/keras-nlp
Pull request #901 (2023-03-22 05:43:56)

Possible fix for task.summary()
See https://github.com/keras-team/keras-nlp/issues/898.

[demo colab](https://colab.research.google.com/gist/mattdangerw/0b24ef4548e86c2872e9c21bd4d106ea/task-summaries.ipynb)

mattdangerw: Talk with @fchollet, think we should go for something like this (almost any fix will be a positive delta).

Added some tests and marking this ready for review!

Patch:
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_tokenizer.py b/keras_nlp/models/deberta_v3/deberta_v3_tokenizer.py
index d1e94576..c5400e7a 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_tokenizer.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_tokenizer.py
@@ -16,6 +16,8 @@
 
 import copy
 
+import tensorflow as tf
+
 from keras_nlp.api_export import keras_nlp_export
 from keras_nlp.models.deberta_v3.deberta_v3_presets import backbone_presets
 from keras_nlp.tokenizers.sentence_piece_tokenizer import SentencePieceTokenizer
@@ -120,6 +122,10 @@ class DebertaV3Tokenizer(SentencePieceTokenizer):
             return self.mask_token_id
         return super().token_to_id(token)
 
+    def detokenize(self, ids):
+        ids = tf.ragged.boolean_mask(ids, tf.not_equal(ids, self.mask_token_id))
+        return super().detokenize(ids)
+
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
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/models/whisper/whisper_backbone.py b/keras_nlp/models/whisper/whisper_backbone.py
index fcb35302..5f5e3d3a 100644
--- a/keras_nlp/models/whisper/whisper_backbone.py
+++ b/keras_nlp/models/whisper/whisper_backbone.py
@@ -38,7 +38,7 @@ def whisper_kernel_initializer(stddev=0.02):
 
 @keras_nlp_export("keras_nlp.models.WhisperBackbone")
 class WhisperBackbone(Backbone):
-    """Whisper encoder-decoder network for speech.
+    """A Whisper encoder-decoder network for speech.
 
     This class implements a Transformer-based encoder-decoder model as
     described in
@@ -48,7 +48,7 @@ class WhisperBackbone(Backbone):
 
     The default constructor gives a fully customizable, randomly initialized Whisper
     model with any number of layers, heads, and embedding dimensions. To load
-    preset architectures and weights, use the `from_preset` constructor.
+    preset architectures and weights, use the `from_preset()` constructor.
 
     Disclaimer: Pre-trained models are provided on an "as is" basis, without
     warranties or conditions of any kind. The underlying model is provided by a
@@ -83,17 +83,17 @@ class WhisperBackbone(Backbone):
         ),
     }
 
-    # Randomly initialized Whisper encoder-decoder model with a custom config
+    # Randomly initialized Whisper encoder-decoder model with a custom config.
     model = keras_nlp.models.WhisperBackbone(
         vocabulary_size=51864,
-        num_layers=6,
-        num_heads=8,
-        hidden_dim=512,
-        intermediate_dim=2048,
+        num_layers=4,
+        num_heads=4,
+        hidden_dim=256,
+        intermediate_dim=512,
         max_encoder_sequence_length=128,
-        max_decoder_sequence_length=64,
+        max_decoder_sequence_length=128,
     )
-    output = model(input_data)
+    model(input_data)
     ```
     """
 
diff --git a/keras_nlp/models/whisper/whisper_decoder.py b/keras_nlp/models/whisper/whisper_decoder.py
index d7d7d286..48551dbe 100644
--- a/keras_nlp/models/whisper/whisper_decoder.py
+++ b/keras_nlp/models/whisper/whisper_decoder.py
@@ -20,7 +20,7 @@ from keras_nlp.layers.transformer_decoder import TransformerDecoder
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
 class WhisperDecoder(TransformerDecoder):
-    """Whisper decoder.
+    """A Whisper decoder.
 
     Inherits from `keras_nlp.layers.TransformerDecoder`, and overrides the
     `_build` method so as to remove the bias term from the key projection layer.
diff --git a/keras_nlp/models/whisper/whisper_encoder.py b/keras_nlp/models/whisper/whisper_encoder.py
index ec104089..f779da26 100644
--- a/keras_nlp/models/whisper/whisper_encoder.py
+++ b/keras_nlp/models/whisper/whisper_encoder.py
@@ -20,7 +20,7 @@ from keras_nlp.layers.transformer_encoder import TransformerEncoder
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
 class WhisperEncoder(TransformerEncoder):
-    """Whisper encoder.
+    """A Whisper encoder.
 
     Inherits from `keras_nlp.layers.TransformerEncoder`, and overrides the
     `_build` method so as to remove the bias term from the key projection layer.
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py
index 6bef8e2c..db5967cb 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py
@@ -24,17 +24,16 @@ from keras_nlp.utils.python_utils import classproperty
 
 @keras_nlp_export("keras_nlp.models.XLMRobertaBackbone")
 class XLMRobertaBackbone(roberta_backbone.RobertaBackbone):
-    """XLM-RoBERTa encoder.
+    """An XLM-RoBERTa encoder network.
 
-    This network implements a bi-directional Transformer-based encoder as
-    described in
-    ["Unsupervised Cross-lingual Representation Learning at Scale"](https://arxiv.org/abs/1911.02116).
-    It includes the embedding lookups and transformer layers, but does not
+    This class implements a bi-directional Transformer-based encoder as
+    described in ["Unsupervised Cross-lingual Representation Learning at Scale"](https://arxiv.org/abs/1911.02116).
+    It includes the embedding lookups and transformer layers, but it does not
     include the masked language modeling head used during pretraining.
 
     The default constructor gives a fully customizable, randomly initialized
-    RoBERTa encoder with any number of layers, heads, and embedding
-    dimensions. To load preset architectures and weights, use the `from_preset`
+    RoBERTa encoder with any number of layers, heads, and embedding dimensions.
+    To load preset architectures and weights, use the `from_preset()`
     constructor.
 
     Disclaimer: Pre-trained models are provided on an "as is" basis, without
@@ -53,9 +52,10 @@ class XLMRobertaBackbone(roberta_backbone.RobertaBackbone):
         dropout: float. Dropout probability for the Transformer encoder.
         max_sequence_length: int. The maximum sequence length this encoder can
             consume. The sequence length of the input must be less than
-            `max_sequence_length`.
+            `max_sequence_length` default value. This determines the variable
+            shape for positional embeddings.
 
-    Example usage:
+    Examples:
     ```python
     input_data = {
         "token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
@@ -63,24 +63,22 @@ class XLMRobertaBackbone(roberta_backbone.RobertaBackbone):
             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
     }
 
-    # Pretrained XLM-R encoder
+    # Pretrained XLM-R encoder.
     model = keras_nlp.models.XLMRobertaBackbone.from_preset(
         "xlm_roberta_base_multi",
     )
-    output = model(input_data)
+    model(input_data)
 
-    # Randomly initialized XLM-R model with custom config
+    # Randomly initialized XLM-R model with custom config.
     model = keras_nlp.models.XLMRobertaBackbone(
         vocabulary_size=250002,
-        num_layers=12,
-        num_heads=12,
-        hidden_dim=768,
-        intermediate_dim=3072,
-        max_sequence_length=12
+        num_layers=4,
+        num_heads=4,
+        hidden_dim=256,
+        intermediate_dim=512,
+        max_sequence_length=128
     )
-
-    # Call the model on the input data.
-    output = model(input_data)
+    model(input_data)
     ```
     """
 
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
index 3ff68b00..8096bb74 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
@@ -34,9 +34,9 @@ class XLMRobertaClassifier(Task):
     """An end-to-end XLM-RoBERTa model for classification tasks.
 
     This model attaches a classification head to a
-    `keras_nlp.model.XLMRobertaBackbone` model, mapping from the backbone
-    outputs to logit output suitable for a classification task. For usage of
-    this model with pre-trained weights, see the `from_preset()` method.
+    `keras_nlp.model.XLMRobertaBackbone` instance, mapping from the backbone
+    outputs to logits suitable for a classification task. For usage of
+    this model with pre-trained weights, see the `from_preset()` constructor.
 
     This model can optionally be configured with a `preprocessor` layer, in
     which case it will automatically apply preprocessing to raw inputs during
@@ -49,7 +49,7 @@ class XLMRobertaClassifier(Task):
     [here](https://github.com/facebookresearch/fairseq).
 
     Args:
-        backbone: A `keras_nlp.models.XLMRoberta` instance.
+        backbone: A `keras_nlp.models.XLMRobertaBackbone` instance.
         num_classes: int. Number of classes to predict.
         hidden_dim: int. The size of the pooler layer.
         dropout: float. The dropout probability value, applied to the pooled
@@ -60,85 +60,34 @@ class XLMRobertaClassifier(Task):
 
     Examples:
 
-    Example usage.
+    Raw string data.
     ```python
-    preprocessed_features = {
-        "token_ids": tf.ones(shape=(2, 12), dtype=tf.int64),
-        "padding_mask": tf.constant(
-            [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]] * 2, shape=(2, 12)),
-    }
-    labels = [0, 3]
-
-    # Randomly initialized XLM-RoBERTa encoder
-    backbone = keras_nlp.models.XLMRobertaBackbone(
-        vocabulary_size=250002,
-        num_layers=12,
-        num_heads=12,
-        hidden_dim=768,
-        intermediate_dim=3072,
-        max_sequence_length=12
-    )
-
-    # Create a XLM-RoBERTa classifier and fit your data.
-    classifier = keras_nlp.models.XLMRobertaClassifier(
-        backbone,
-        num_classes=4,
-        preprocessor=None,
-    )
-    classifier.compile(
-        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
-    )
-    classifier.fit(x=preprocessed_features, y=labels, batch_size=2)
-
-    # Access backbone programatically (e.g., to change `trainable`)
-    classifier.backbone.trainable = False
-    ```
-
-    Raw string inputs.
-    ```python
-    # Create a dataset with raw string features in an `(x, y)` format.
-    features = ["The quick brown fox jumped.", "I forgot my homework."]
+    features = ["The quick brown fox jumped.", "نسيت الواجب"]
     labels = [0, 3]
 
-    # Create a XLMRobertaClassifier and fit your data.
+    # Pretrained classifier.
     classifier = keras_nlp.models.XLMRobertaClassifier.from_preset(
         "xlm_roberta_base_multi",
         num_classes=4,
     )
-    classifier.compile(
-        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
-    )
     classifier.fit(x=features, y=labels, batch_size=2)
-    ```
+    classifier.predict(x=features, batch_size=2)
 
-    Raw string inputs with customized preprocessing.
-    ```python
-    # Create a dataset with raw string features in an `(x, y)` format.
-    features = ["The quick brown fox jumped.", "I forgot my homework."]
-    labels = [0, 3]
-
-    # Use a shorter sequence length.
-    preprocessor = keras_nlp.models.XLMRobertaPreprocessor.from_preset(
-        "xlm_roberta_base_multi",
-        sequence_length=128,
-    )
-
-    # Create a XLMRobertaClassifier and fit your data.
-    classifier = keras_nlp.models.XLMRobertaClassifier.from_preset(
-        "xlm_roberta_base_multi",
-        num_classes=4,
-        preprocessor=preprocessor,
-    )
+    # Re-compile (e.g., with a new learning rate).
     classifier.compile(
         loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        optimizer=keras.optimizers.Adam(5e-5),
+        jit_compile=True,
     )
+    # Access backbone programmatically (e.g., to change `trainable`).
+    classifier.backbone.trainable = False
+    # Fit again.
     classifier.fit(x=features, y=labels, batch_size=2)
     ```
 
-    Preprocessed inputs.
+    Preprocessed integer data.
     ```python
-    # Create a dataset with preprocessed features in an `(x, y)` format.
-    preprocessed_features = {
+    features = {
         "token_ids": tf.ones(shape=(2, 12), dtype=tf.int64),
         "padding_mask": tf.constant(
             [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]] * 2, shape=(2, 12)
@@ -146,16 +95,57 @@ class XLMRobertaClassifier(Task):
     }
     labels = [0, 3]
 
-    # Create a XLMRobertaClassifier and fit your data.
+    # Pretrained classifier without preprocessing.
     classifier = keras_nlp.models.XLMRobertaClassifier.from_preset(
         "xlm_roberta_base_multi",
         num_classes=4,
         preprocessor=None,
     )
-    classifier.compile(
-        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+    classifier.fit(x=features, y=labels, batch_size=2)
+    ```
+
+    Custom backbone and vocabulary.
+    ```python
+    features = ["The quick brown fox jumped.", "نسيت الواجب"]
+    labels = [0, 3]
+
+    def train_sentencepiece(ds, vocab_size):
+        bytes_io = io.BytesIO()
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=ds.as_numpy_iterator(),
+            model_writer=bytes_io,
+            vocab_size=vocab_size,
+            model_type="WORD",
+            unk_id=0,
+            bos_id=1,
+            eos_id=2,
+        )
+        return bytes_io.getvalue()
+    ds = tf.data.Dataset.from_tensor_slices(
+        ["the quick brown fox", "the earth is round"]
+    )
+    proto = train_sentencepiece(ds, vocab_size=10)
+    tokenizer = keras_nlp.models.XLMRobertaTokenizer(
+        proto=proto
+    )
+    preprocessor = keras_nlp.models.XLMRobertaPreprocessor(
+        tokenizer,
+        sequence_length=128,
+    )
+    backbone = keras_nlp.models.XLMRobertaBackbone(
+        vocabulary_size=250002,
+        num_layers=4,
+        num_heads=4,
+        hidden_dim=256,
+        intermediate_dim=512,
+        max_sequence_length=128,
     )
-    classifier.fit(x=preprocessed_features, y=labels, batch_size=2)
+    classifier = keras_nlp.models.XLMRobertaClassifier(
+        backbone=backbone,
+        preprocessor=preprocessor,
+        num_classes=4,
+    )
+    classifier.fit(x=features, y=labels, batch_size=2)
     ```
     """
 
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor.py b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor.py
index fa38f1e3..6ecb5016 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor.py
@@ -34,35 +34,23 @@ from keras_nlp.utils.python_utils import classproperty
 
 @keras_nlp_export("keras_nlp.models.XLMRobertaPreprocessor")
 class XLMRobertaPreprocessor(Preprocessor):
-    """XLM-RoBERTa preprocessing layer.
+    """An XLM-RoBERTa preprocessing layer which tokenizes and packs inputs.
 
     This preprocessing layer will do three things:
 
-    - Tokenize any number of input segments using the `tokenizer`.
-    - Pack the inputs together with the appropriate `"
"`, `"
"` and
-      `"<pad>"` tokens, i.e., adding a single `"
"` at the start of the
-      entire sequence, `"

"` at the end of each segment, save the last
-      and a `"
"` at the end of the entire sequence.
-    - Construct a dictionary with keys `"token_ids"` and `"padding_mask"`,
-      that can be passed directly to a XLM-RoBERTa model.
+    1. Tokenize any number of input segments using the `tokenizer`.
+    2. Pack the inputs together using a `keras_nlp.layers.MultiSegmentPacker`.
+      with the appropriate `"
"`, `"
"` and `"<pad>"` tokens, i.e., adding
+      a single `"
"` at the start of the entire sequence, `"

"` at the
+      end of each segment, save the last and a `"
"` at the end of the
+      entire sequence.
+    3. Construct a dictionary with keys `"token_ids"` and `"padding_mask"`,
+      that can be passed directly to an XLM-RoBERTa model.
 
     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_weight`. `x` can be a python string or tensor representing a single
-    segment, a list of python strings representing a batch of single segments,
-    or a list of tensors representing multiple segments to be packed together.
-    `y` and `sample_weight` 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.Dataset.map` 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.tokenizers.XLMRobertaTokenizer` instance.
         sequence_length: The length of the packed inputs.
@@ -77,77 +65,84 @@ class XLMRobertaPreprocessor(Preprocessor):
                     left-to-right manner and fills up the buckets until we run
                     out of budget. It supports an arbitrary number of segments.
 
+    Call arguments:
+        x: A tensor of single string sequences, or a tuple of multiple
+            tensor sequences to be packed together. Inputs may be batched or
+            unbatched. For single sequences, raw python inputs will be converted
+            to tensors. For multiple sequences, pass tensors directly.
+        y: Any label data. Will be passed through unaltered.
+        sample_weight: Any label weight data. Will be passed through unaltered.
+
     Examples:
     ```python
-    # Load the preprocessor from a preset.
-    preprocessor = keras_nlp.models.XLMRobertaPreprocessor.from_preset("xlm_roberta_base_multi")
+    preprocessor = keras_nlp.models.XLMRobertaPreprocessor.from_preset(
+        "xlm_roberta_base_multi"
+    )
 
     # Tokenize and pack a single sentence.
-    sentence = tf.constant("The quick brown fox jumped.")
-    preprocessor(sentence)
-    # Same output.
     preprocessor("The quick brown fox jumped.")
 
-    # Tokenize and a batch of single sentences.
-    sentences = tf.constant(
-        ["The quick brown fox jumped.", "Call me Ishmael."]
+    # Tokenize a batch of single sentences.
+    preprocessor(["The quick brown fox jumped.", "اسمي اسماعيل"])
+
+    # Preprocess a batch of sentence pairs.
+    # When handling multiple sequences, always convert to tensors first!
+    first = tf.constant(["The quick brown fox jumped.", "اسمي اسماعيل"])
+    second = tf.constant(["The fox tripped.", "الأسد ملك الغابة"])
+    preprocessor((first, second))
+
+    # Custom vocabulary.
+    def train_sentencepiece(ds, vocab_size):
+        bytes_io = io.BytesIO()
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=ds.as_numpy_iterator(),
+            model_writer=bytes_io,
+            vocab_size=vocab_size,
+            model_type="WORD",
+            unk_id=0,
+            bos_id=1,
+            eos_id=2,
+        )
+        return bytes_io.getvalue()
+    ds = tf.data.Dataset.from_tensor_slices(
+        ["the quick brown fox", "the earth is round"]
     )
-    preprocessor(sentences)
-    # Same output.
-    preprocessor(
-        ["The quick brown fox jumped.", "Call me Ishmael."]
+    proto = train_sentencepiece(ds, vocab_size=10)
+    tokenizer = keras_nlp.models.XLMRobertaTokenizer(proto=proto)
+    preprocessor = keras_nlp.models.XLMRobertaPreprocessor(tokenizer)
+    preprocessor("The quick brown fox jumped.")
+    ```
+
+    Mapping with `tf.data.Dataset`.
+    ```python
+    preprocessor = keras_nlp.models.XLMRobertaPreprocessor.from_preset(
+        "xlm_roberta_base_multi"
     )
 
-    # Tokenize and pack a sentence pair.
-    first_sentence = tf.constant("The quick brown fox jumped.")
-    second_sentence = tf.constant("The fox tripped.")
-    preprocessor((first_sentence, second_sentence))
+    first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
+    second = tf.constant(["The fox tripped.", "Oh look, a whale."])
+    label = tf.constant([1, 1])
 
-    # Map a dataset to preprocess a single sentence.
-    features = tf.constant(
-        ["The quick brown fox jumped.", "Call me Ishmael."]
-    )
-    labels = tf.constant([0, 1])
-    ds = tf.data.Dataset.from_tensor_slices((features, labels))
+    # Map labeled single sentences.
+    ds = tf.data.Dataset.from_tensor_slices((first, label))
     ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
 
-    # Map a dataset to preprocess sentence pairs.
-    first_sentences = tf.constant(
-        ["The quick brown fox jumped.", "Call me Ishmael."]
-    )
-    second_sentences = tf.constant(
-        ["The fox tripped.", "Oh look, a whale."]
-    )
-    labels = tf.constant([1, 1])
-    ds = tf.data.Dataset.from_tensor_slices(
-        (
-            (first_sentences, second_sentences), labels
-        )
-    )
+    # Map unlabeled single sentences.
+    ds = tf.data.Dataset.from_tensor_slices(first)
     ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
 
-    # Map a dataset to preprocess unlabeled sentence pairs.
-    first_sentences = tf.constant(
-        ["The quick brown fox jumped.", "Call me Ishmael."]
-    )
-    second_sentences = tf.constant(
-        ["The fox tripped.", "Oh look, a whale."]
-    )
-    ds = tf.data.Dataset.from_tensor_slices((first_sentences, second_sentences))
+    # Map labeled sentence pairs.
+    ds = tf.data.Dataset.from_tensor_slices(((first, second), label))
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Map unlabeled sentence pairs.
+    ds = tf.data.Dataset.from_tensor_slices((first, second))
     # Watch out for tf.data's default unpacking of tuples here!
     # Best to invoke the `preprocessor` directly in this case.
     ds = ds.map(
-        lambda s1, s2: preprocessor(x=(s1, s2)),
+        lambda first, second: preprocessor(x=(first, second)),
         num_parallel_calls=tf.data.AUTOTUNE,
     )
-
-    # Alternatively, you can create a preprocessor from your own vocabulary.
-    # The usage is exactly the same as above.
-    tokenizer = keras_nlp.models.XLMRobertaTokenizer(proto="model.spm")
-    preprocessor = keras_nlp.models.XLMRobertaPreprocessor(
-        tokenizer=tokenizer,
-        sequence_length=10,
-    )
     ```
     """
 
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer.py b/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer.py
index e077a764..d679ab9b 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer.py
@@ -27,19 +27,18 @@ from keras_nlp.utils.tf_utils import tensor_to_string_list
 
 @keras_nlp_export("keras_nlp.models.XLMRobertaTokenizer")
 class XLMRobertaTokenizer(SentencePieceTokenizer):
-    """XLM-RoBERTa tokenizer layer based on SentencePiece.
+    """An XLM-RoBERTa tokenizer using SentencePiece subword segmentation.
 
     This tokenizer class will tokenize raw strings into integer sequences and
     is based on `keras_nlp.tokenizers.SentencePieceTokenizer`. Unlike the
     underlying tokenizer, it will check for all special tokens needed by
     XLM-RoBERTa models and provides a `from_preset()` method to automatically
-    download a matching vocabulary for a XLM-RoBERTa preset.
+    download a matching vocabulary for an XLM-RoBERTa preset.
 
-    The original fairseq implementation of XLM-RoBERTa modifies the indices of
-    the SentencePiece tokenizer output. To preserve compatibility, we make the
-    same changes, i.e., `"
"`, `"<pad>"`, `"
"` and `"<unk>"` are mapped to
-    0, 1, 2, 3, respectively, and non-special token indices are shifted right
-    by one.
+    Note: If you are providing your own custom SentencePiece model, the original
+    fairseq implementation of XLM-RoBERTa re-maps some token indices from the
+    underlying sentencepiece output. To preserve compatibility, we do the same
+    re-mapping here.
 
     If input is a batch of strings (rank > 0), the layer will output a
     `tf.RaggedTensor` where the last dimension of the output is ragged.
@@ -48,14 +47,27 @@ class XLMRobertaTokenizer(SentencePieceTokenizer):
     `tf.Tensor` with static shape `[None]`.
 
     Args:
-        proto: Either a `string` path to a SentencePiece proto file, or a
+        proto: Either a `string` path to a SentencePiece proto file or a
             `bytes` object with a serialized SentencePiece proto. See the
             [SentencePiece repository](https://github.com/google/sentencepiece)
             for more details on the format.
 
     Examples:
-
     ```python
+    tokenizer = keras_nlp.models.XLMRobertaTokenizer.from_preset(
+        "xlm_roberta_base_multi",
+    )
+
+    # Unbatched inputs.
+    tokenizer("the quick brown fox")
+
+    # Batched inputs.
+    tokenizer(["the quick brown fox", "الأرض كروية"])
+
+    # Detokenization.
+    tokenizer.detokenize(tokenizer("the quick brown fox"))
+
+    # Custom vocabulary
     def train_sentencepiece(ds, vocab_size):
         bytes_io = io.BytesIO()
         sentencepiece.SentencePieceTrainer.train(
@@ -72,18 +84,8 @@ class XLMRobertaTokenizer(SentencePieceTokenizer):
     ds = tf.data.Dataset.from_tensor_slices(
         ["the quick brown fox", "the earth is round"]
     )
-
     proto = train_sentencepiece(ds, vocab_size=10)
     tokenizer = keras_nlp.models.XLMRobertaTokenizer(proto=proto)
-
-    # Batched inputs.
-    tokenizer(["the quick brown fox", "the earth is round"])
-
-    # Unbatched inputs.
-    tokenizer("the quick brown fox")
-
-    # Detokenization.
-    tokenizer.detokenize(tf.constant([[0, 4, 9, 5, 7, 2]]))
     ```
     """
 
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)


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 5409364e..936773ce 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_backbone_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_backbone_test.py
@@ -15,6 +15,7 @@
 
 import os
 
+import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
@@ -24,23 +25,19 @@ from keras_nlp.models.deberta_v3.deberta_v3_backbone import DebertaV3Backbone
 
 class DebertaV3BackboneTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = DebertaV3Backbone(
-            vocabulary_size=1000,
+        self.backbone = DebertaV3Backbone(
+            vocabulary_size=10,
             num_layers=2,
             num_heads=2,
-            hidden_dim=64,
-            intermediate_dim=128,
-            max_sequence_length=128,
-            bucket_size=64,
+            hidden_dim=2,
+            intermediate_dim=4,
+            max_sequence_length=5,
+            bucket_size=2,
         )
         self.batch_size = 8
         self.input_batch = {
-            "token_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"
-            ),
+            "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(
@@ -48,45 +45,46 @@ class DebertaV3BackboneTest(tf.test.TestCase, parameterized.TestCase):
         ).batch(2)
 
     def test_valid_call_deberta(self):
-        self.model(self.input_batch)
+        self.backbone(self.input_batch)
 
-        # Check default name passed through
-        self.assertRegexpMatches(self.model.name, "deberta_v3_backbone")
+    def test_name(self):
+        self.assertRegexpMatches(self.backbone.name, "deberta_v3_backbone")
+
+    def test_token_embedding(self):
+        output = self.backbone.token_embedding(self.input_batch["token_ids"])
+        self.assertEqual(output.shape, (2, 5, 2))
 
     def test_variable_sequence_length_call_deberta(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"
-                ),
-                "padding_mask": tf.ones(
-                    (self.batch_size, seq_length), dtype="int32"
-                ),
+                "token_ids": tf.ones((2, seq_length), dtype="int32"),
+                "padding_mask": tf.ones((2, seq_length), dtype="int32"),
             }
-            self.model(input_data)
+            output = self.backbone(input_data)
+            self.assertAllEqual(
+                tf.shape(output),
+                [2, seq_length, self.backbone.hidden_dim],
+            )
 
-    @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.
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py b/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py
index 55e3a2f1..79cdb6c6 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py
@@ -16,6 +16,7 @@
 import io
 import os
 
+import pytest
 import sentencepiece
 import tensorflow as tf
 from absl.testing import parameterized
@@ -53,87 +54,70 @@ class DebertaV3ClassifierTest(tf.test.TestCase, parameterized.TestCase):
             user_defined_symbols="[MASK]",
         )
         self.preprocessor = DebertaV3Preprocessor(
-            tokenizer=DebertaV3Tokenizer(proto=bytes_io.getvalue()),
-            sequence_length=12,
+            DebertaV3Tokenizer(proto=bytes_io.getvalue()),
+            sequence_length=5,
         )
         self.backbone = DebertaV3Backbone(
-            vocabulary_size=1000,
+            vocabulary_size=self.preprocessor.tokenizer.vocabulary_size(),
             num_layers=2,
             num_heads=2,
-            hidden_dim=64,
-            intermediate_dim=128,
-            max_sequence_length=128,
-            bucket_size=64,
+            hidden_dim=2,
+            intermediate_dim=4,
+            max_sequence_length=self.preprocessor.packer.sequence_length,
+            bucket_size=2,
         )
         self.classifier = DebertaV3Classifier(
             self.backbone,
             4,
             preprocessor=self.preprocessor,
         )
-        self.classifier_no_preprocessing = DebertaV3Classifier(
-            self.backbone,
-            4,
-            preprocessor=None,
-        )
 
         self.raw_batch = tf.constant(
             [
                 "the quick brown fox.",
                 "the slow brown fox.",
-                "the earth is round",
-                "the earth is spherical",
             ]
         )
         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_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_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_debertav3_classifier_fit_default_compile(self):
+    def test_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_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)
+        self.classifier.fit(self.preprocessed_dataset)
 
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_classifier_fit_no_preprocessing(self, jit_compile):
-        self.classifier_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.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_saving_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/deberta_v3/deberta_v3_masked_lm_preprocessor_test.py b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_preprocessor_test.py
index def702de..f84c4037 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_preprocessor_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_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 DebertaV3PreprocessorTest(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/deberta_v3/deberta_v3_masked_lm_test.py b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_test.py
index 67909376..406fc616 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm_test.py
@@ -16,6 +16,7 @@
 import io
 import os
 
+import pytest
 import sentencepiece
 import tensorflow as tf
 from absl.testing import parameterized
@@ -31,14 +32,6 @@ from keras_nlp.models.deberta_v3.deberta_v3_tokenizer import DebertaV3Tokenizer
 
 class DebertaV3MaskedLMTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.backbone = DebertaV3Backbone(
-            vocabulary_size=1000,
-            num_layers=2,
-            num_heads=2,
-            hidden_dim=64,
-            intermediate_dim=128,
-            max_sequence_length=128,
-        )
         bytes_io = io.BytesIO()
         vocab_data = tf.data.Dataset.from_tensor_slices(
             ["the quick brown fox", "the earth is round", "an eagle flew"]
@@ -60,78 +53,71 @@ class DebertaV3MaskedLMTest(tf.test.TestCase, parameterized.TestCase):
         )
         proto = bytes_io.getvalue()
         self.preprocessor = DebertaV3MaskedLMPreprocessor(
-            tokenizer=DebertaV3Tokenizer(proto=proto),
-            mask_selection_length=4,
-            sequence_length=10,
+            DebertaV3Tokenizer(proto=proto),
+            mask_selection_length=2,
+            sequence_length=5,
+        )
+        self.backbone = DebertaV3Backbone(
+            vocabulary_size=self.preprocessor.tokenizer.vocabulary_size(),
+            num_layers=2,
+            num_heads=2,
+            hidden_dim=2,
+            intermediate_dim=4,
+            max_sequence_length=self.preprocessor.packer.sequence_length,
         )
         self.masked_lm = DebertaV3MaskedLM(
             self.backbone,
             preprocessor=self.preprocessor,
         )
-        self.masked_lm_no_preprocessing = DebertaV3MaskedLM(
-            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 eagle flew over fox.",
             ]
         )
-        self.preprocessed_batch = self.preprocessor(self.raw_batch)[0]
+        self.preprocessed_batch = self.preprocessor(self.raw_batch)
         self.raw_dataset = tf.data.Dataset.from_tensor_slices(
             self.raw_batch
         ).batch(2)
         self.preprocessed_dataset = self.raw_dataset.map(self.preprocessor)
 
-    def test_valid_call_masked_lm(self):
-        self.masked_lm(self.preprocessed_batch)
+    def test_valid_call_classifier(self):
+        self.masked_lm(self.preprocessed_batch[0])
 
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_deberta_v3_masked_lm_predict(self, jit_compile):
-        self.masked_lm.compile(jit_compile=jit_compile)
+    def test_classifier_predict(self):
         self.masked_lm.predict(self.raw_batch)
+        self.masked_lm.preprocessor = None
+        self.masked_lm.predict(self.preprocessed_batch[0])
 
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_deberta_v3_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_delberta_masked_lm_fit_default_compile(self):
-            self.masked_lm.fit(self.raw_dataset)
+    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_deberta_v3_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)
-
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_deberta_v3_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,
+        self.masked_lm.fit(self.preprocessed_dataset)
+
+    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)
@@ -139,7 +125,6 @@ class DebertaV3MaskedLMTest(tf.test.TestCase, parameterized.TestCase):
         # Check we got the real object back.
         self.assertIsInstance(restored_model, DebertaV3MaskedLM)
 
-        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/deberta_v3/deberta_v3_preprocessor_test.py b/keras_nlp/models/deberta_v3/deberta_v3_preprocessor_test.py
index 233d4330..ec02a6ee 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_preprocessor_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_preprocessor_test.py
@@ -17,6 +17,7 @@
 import io
 import os
 
+import pytest
 import sentencepiece
 import tensorflow as tf
 from absl.testing import parameterized
@@ -136,10 +137,19 @@ class DebertaV3PreprocessorTest(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/deberta_v3/deberta_v3_tokenizer_test.py b/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py
index 3b81bd81..6672068e 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py
@@ -17,6 +17,7 @@
 import io
 import os
 
+import pytest
 import sentencepiece
 import tensorflow as tf
 from absl.testing import parameterized
@@ -64,6 +65,11 @@ class DebertaV3TokenizerTest(tf.test.TestCase, parameterized.TestCase):
         output = self.tokenizer.detokenize(input_data)
         self.assertEqual(output, tf.constant(["the quick brown fox"]))
 
+    def test_detokenize_mask_token(self):
+        input_data = tf.constant([[4, 9, 5, 7, self.tokenizer.mask_token_id]])
+        output = self.tokenizer.detokenize(input_data)
+        self.assertEqual(output, tf.constant(["the quick brown fox"]))
+
     def test_vocabulary_size(self):
         self.assertEqual(self.tokenizer.vocabulary_size(), 11)
 
@@ -89,10 +95,19 @@ class DebertaV3TokenizerTest(tf.test.TestCase, parameterized.TestCase):
         with self.assertRaises(ValueError):
             DebertaV3Tokenizer(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/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:")