Repository: keras-team/keras-nlp

Issue:
Add a RoBERTa preprocessing layer
Now that we have landed the `BytePairTokenizer`, we should be able to land the preprocessing for the RoBERTa model, modeled after the `BertPreprocessor` layer.

Discussion:

Patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index da6035ec..9e9d27e4 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -19,8 +19,9 @@ from keras_nlp.models.distilbert.distilbert_models import DistilBert
 from keras_nlp.models.distilbert.distilbert_preprocessing import (
     DistilBertPreprocessor,
 )
-from keras_nlp.models.gpt2.gpt2_models import Gpt2
+from keras_nlp.models.gpt2.gpt2_models import GPT2
 from keras_nlp.models.roberta.roberta_models import Roberta
+from keras_nlp.models.roberta.roberta_preprocessing import RobertaPreprocessor
 from keras_nlp.models.roberta.roberta_tasks import RobertaClassifier
 from keras_nlp.models.xlm_roberta.xlm_roberta_models import XLMRoberta
 from keras_nlp.models.xlm_roberta.xlm_roberta_preprocessing import (
diff --git a/keras_nlp/models/gpt2/gpt2_models.py b/keras_nlp/models/gpt2/gpt2_models.py
index 14f0937f..93711333 100644
--- a/keras_nlp/models/gpt2/gpt2_models.py
+++ b/keras_nlp/models/gpt2/gpt2_models.py
@@ -26,7 +26,7 @@ def _gpt_2_kernel_initializer(stddev=0.02):
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class Gpt2(keras.Model):
+class GPT2(keras.Model):
     """GPT-2 core network with hyperparameters.
 
     This network implements a Transformer-based decoder network,
@@ -65,7 +65,7 @@ class Gpt2(keras.Model):
     }
 
     # Randomly initialized GPT-2 decoder
-    model = keras_nlp.models.Gpt2(
+    model = keras_nlp.models.GPT2(
         vocabulary_size=50257,
         num_layers=12,
         num_heads=12,
diff --git a/keras_nlp/models/roberta/roberta_preprocessing.py b/keras_nlp/models/roberta/roberta_preprocessing.py
index 9344b47d..1f78ebae 100644
--- a/keras_nlp/models/roberta/roberta_preprocessing.py
+++ b/keras_nlp/models/roberta/roberta_preprocessing.py
@@ -12,10 +12,187 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""RoBERTa preprocessing layers."""
+
 import tensorflow as tf
 import tensorflow_text as tf_text
 from tensorflow import keras
 
+from keras_nlp.tokenizers.byte_pair_tokenizer import BytePairTokenizer
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class RobertaPreprocessor(keras.layers.Layer):
+    """RoBERTa preprocessing layer.
+
+    This preprocessing layer will do three things:
+
+    - Tokenize any number of inputs using a
+      `keras_nlp.tokenizers.BytePairTokenizer`.
+    - 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 RoBERTa model.
+
+    This layer will accept either a tuple of (possibly batched) inputs, or a
+    single input tensor. If a single tensor is passed, it will be packed
+    equivalently to a tuple with a single element.
+
+    The `BytePairTokenizer` can be accessed via the `tokenizer` property on
+    this layer, and can be used directly for custom packing on inputs.
+
+    Args:
+        vocabulary: string or dict, maps token to integer ids. If it is a
+            string, it should be the file path to a json file.
+        merges: string or list, contains the merge rule. If it is a string,
+            it should be the file path to merge rules. The merge rule file
+            should have one merge rule per line. Every merge rule contains
+            merge entities separated by a space. Please refer to this example:
+            https://storage.googleapis.com/keras-nlp/models/roberta_base/merges.txt.
+        sequence_length: The length of the packed inputs.
+        truncate: string. The algorithm to truncate a list of batched segments
+            to fit within `sequence_length`. The value can be either
+            `round_robin` or `waterfall`:
+                - `"round_robin"`: Available space is assigned one token at a
+                    time in a round-robin fashion to the inputs that still need
+                    some, until the limit is reached.
+                - `"waterfall"`: The allocation of the budget is done using a
+                    "waterfall" algorithm that allocates quota in a
+                    left-to-right manner and fills up the buckets until we run
+                    out of budget. It supports an arbitrary number of segments.
+
+    Examples:
+    ```python
+    vocab = {
+        "
": 0,
+        "<pad>": 1,
+        "
": 2,
+        "reful": 3,
+        "gent": 4,
+        "Ġafter": 5,
+        "noon": 6,
+        "Ġsun": 7,
+        "Ġbright": 8,
+        "Ġnight": 9,
+        "Ġmoon": 10,
+    }
+    merges = ["Ġ a", "Ġ m", "Ġ s", "Ġ b", "Ġ n", "r e", "f u", "g e", "n t"]
+    merges += ["e r", "n o", "o n", "i g", "h t"]
+    merges += ["Ġs u", "Ġa f", "Ġm o", "Ġb r","ge nt", "no on", "re fu", "ig ht"]
+    merges += ["Ġn ight", "Ġsu n", "Ġaf t", "Ġmo on", "Ġbr ight", "refu l", "Ġaft er"]
+
+    preprocessor = keras_nlp.models.RobertaPreprocessor(
+        vocabulary=vocab,
+        merges=merges,
+        sequence_length=20,
+    )
+
+    # Tokenize and pack a single sentence directly.
+    preprocessor(" afternoon sun")
+
+    # Tokenize and pack multiples sentences directly.
+    preprocessor((" afternoon sun", "refulgent sun"))
+
+    # Map a dataset to preprocess a single sentence at a time.
+    features = [" afternoon sun", "refulgent sun"]
+    labels = [0, 1]
+    ds = tf.data.Dataset.from_tensor_slices((features, labels))
+    ds = ds.map(
+        lambda x, y: (preprocessor(x), y),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+
+    # Map a dataset to preprocess multiple sentences.
+    first_sentences = [" afternoon sun", "refulgent sun"]
+    second_sentences = [" night moon", " bright moon"]
+    labels = [1, 1]
+    ds = tf.data.Dataset.from_tensor_slices(
+        (
+            (first_sentences, second_sentences), labels
+        )
+    )
+    ds = ds.map(
+        lambda x, y: (preprocessor(x), y),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary,
+        merges,
+        sequence_length=512,
+        truncate="round_robin",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.tokenizer = BytePairTokenizer(
+            vocabulary=vocabulary,
+            merges=merges,
+        )
+
+        # Check for necessary special tokens.
+        start_token = "
"
+        pad_token = "<pad>"
+        end_token = "
"
+        for token in [start_token, pad_token, end_token]:
+            if token not in self.tokenizer.get_vocabulary():
+                raise ValueError(
+                    f"Cannot find token `'{token}'` in the provided "
+                    f"`vocabulary`. Please provide `'{token}'` in your "
+                    "`vocabulary` or use a pretrained `vocabulary` name."
+                )
+
+        self.pad_token_id = self.tokenizer.token_to_id(pad_token)
+        self.packer = RobertaMultiSegmentPacker(
+            start_value=self.tokenizer.token_to_id(start_token),
+            end_value=self.tokenizer.token_to_id(end_token),
+            pad_value=self.pad_token_id,
+            truncate=truncate,
+            sequence_length=sequence_length,
+        )
+
+    def vocabulary_size(self):
+        """Returns the vocabulary size of the tokenizer."""
+        return self.tokenizer.vocabulary_size()
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary": self.tokenizer.vocabulary,
+                "merges": self.tokenizer.merges,
+                "sequence_length": self.packer.sequence_length,
+                "truncate": self.packer.truncate,
+            }
+        )
+        return config
+
+    def call(self, inputs):
+        if not isinstance(inputs, (list, tuple)):
+            inputs = [inputs]
+
+        inputs = [self.tokenizer(x) for x in inputs]
+        token_ids = self.packer(inputs)
+        return {
+            "token_ids": token_ids,
+            "padding_mask": token_ids != self.pad_token_id,
+        }
+
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        sequence_length=None,
+        truncate="round_robin",
+        **kwargs,
+    ):
+        raise NotImplementedError
+
 
 # TODO: This is a temporary, unexported layer until we find a way to make the
 # `MultiSegmentPacker` layer more generic.
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_tasks.py b/keras_nlp/models/xlm_roberta/xlm_roberta_tasks.py
new file mode 100644
index 00000000..2fc8580e
--- /dev/null
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_tasks.py
@@ -0,0 +1,112 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""XLM-RoBERTa task specific models and heads."""
+
+from tensorflow import keras
+
+from keras_nlp.models.roberta.roberta_models import roberta_kernel_initializer
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class XLMRobertaClassifier(keras.Model):
+    """XLM-RoBERTa encoder model with a classification head.
+
+    Args:
+        backbone: A `keras_nlp.models.XLMRoberta` instance.
+        num_classes: int. Number of classes to predict.
+        hidden_dim: int. The size of the pooler layer.
+
+    Example usage:
+    ```python
+    # Randomly initialized XLM-RoBERTa encoder
+    model = keras_nlp.models.XLMRoberta(
+        vocabulary_size=50265,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=12
+    )
+
+    # Call classifier on the inputs.
+    input_data = {
+        "token_ids": tf.random.uniform(
+            shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size),
+        "padding_mask": tf.constant(
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+    }
+    classifier = keras_nlp.models.XLMRobertaClassifier(
+        backbone=model,
+        num_classes=4,
+    )
+    logits = classifier(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        backbone,
+        num_classes,
+        hidden_dim=None,
+        dropout=0.0,
+        **kwargs,
+    ):
+        inputs = backbone.input
+        if hidden_dim is None:
+            hidden_dim = backbone.hidden_dim
+
+        x = backbone(inputs)[:, backbone.start_token_index, :]
+        x = keras.layers.Dropout(dropout, name="pooled_dropout")(x)
+        x = keras.layers.Dense(
+            hidden_dim, activation="tanh", name="pooled_dense"
+        )(x)
+        x = keras.layers.Dropout(dropout, name="classifier_dropout")(x)
+        outputs = keras.layers.Dense(
+            num_classes,
+            kernel_initializer=roberta_kernel_initializer(),
+            name="logits",
+        )(x)
+
+        # Instantiate using Functional API Model constructor
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            **kwargs,
+        )
+        # All references to `self` below this line
+        self._backbone = backbone
+        self.num_classes = num_classes
+        self.hidden_dim = hidden_dim
+        self.dropout = dropout
+
+    @property
+    def backbone(self):
+        """A `keras_nlp.models.XLMRoberta` instance providing the encoder submodel."""
+        return self._backbone
+
+    def get_config(self):
+        return {
+            "backbone": keras.layers.serialize(self.backbone),
+            "num_classes": self.num_classes,
+            "hidden_dim": self.hidden_dim,
+            "dropout": self.dropout,
+            "name": self.name,
+            "trainable": self.trainable,
+        }
+
+    @classmethod
+    def from_config(cls, config):
+        if "backbone" in config:
+            config["backbone"] = keras.layers.deserialize(config["backbone"])
+        return cls(**config)
diff --git a/keras_nlp/tokenizers/byte_pair_tokenizer.py b/keras_nlp/tokenizers/byte_pair_tokenizer.py
index def254e8..0666cdb6 100644
--- a/keras_nlp/tokenizers/byte_pair_tokenizer.py
+++ b/keras_nlp/tokenizers/byte_pair_tokenizer.py
@@ -516,16 +516,13 @@ class BytePairTokenizer(tokenizer.Tokenizer):
             inputs = tf.expand_dims(inputs, 0)
 
         unicode_text = tf.strings.reduce_join(
-            self.id_to_token_map.lookup(inputs), axis=1
+            self.id_to_token_map.lookup(inputs), axis=-1
         )
         split_unicode_text = tf.strings.unicode_split(unicode_text, "UTF-8")
         byte_text = tf.strings.reduce_join(
-            self.unicode2byte.lookup(split_unicode_text)
+            self.unicode2byte.lookup(split_unicode_text), axis=-1
         )
 
-        if not scalar_input:
-            byte_text = tf.expand_dims(byte_text, 0)
-
         return byte_text
 
     def _transform_bytes(self, tokens):
diff --git a/keras_nlp/utils/text_generation.py b/keras_nlp/utils/text_generation.py
index 0e07806c..07e7ab0a 100644
--- a/keras_nlp/utils/text_generation.py
+++ b/keras_nlp/utils/text_generation.py
@@ -170,11 +170,11 @@ def greedy_search(
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
 
-    _validate_token_probability_fn(token_probability_fn, prompt)
-
     batch_size, length = _get_prompt_shape(prompt)
     prompt, mask = _pad_prompt(prompt, max_length)
 
+    _validate_token_probability_fn(token_probability_fn, prompt)
+
     def one_step(length, prompt):
         pred = token_probability_fn(prompt[:, :length])
         next_token = tf.cast(tf.argmax(pred, axis=-1), dtype=prompt.dtype)
@@ -484,11 +484,11 @@ def random_search(
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
 
-    _validate_token_probability_fn(token_probability_fn, prompt)
-
     batch_size, length = _get_prompt_shape(prompt)
     prompt, mask = _pad_prompt(prompt, max_length)
 
+    _validate_token_probability_fn(token_probability_fn, prompt)
+
     def one_step(length, prompt):
         pred = token_probability_fn(prompt[:, :length])
         if from_logits:
@@ -622,6 +622,9 @@ def top_k_search(
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
 
+    batch_size, length = _get_prompt_shape(prompt)
+    prompt, mask = _pad_prompt(prompt, max_length)
+
     _validate_token_probability_fn(token_probability_fn, prompt)
 
     # If k is greater than the vocabulary size, use the entire vocabulary.
@@ -633,9 +636,6 @@ def top_k_search(
         )
         k = pred.shape[1]
 
-    batch_size, length = _get_prompt_shape(prompt)
-    prompt, mask = _pad_prompt(prompt, max_length)
-
     def one_step(length, prompt):
         pred = token_probability_fn(prompt[:, :length])
         if from_logits:
@@ -777,11 +777,11 @@ def top_p_search(
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
 
-    _validate_token_probability_fn(token_probability_fn, prompt)
-
     batch_size, length = _get_prompt_shape(prompt)
     prompt, mask = _pad_prompt(prompt, max_length)
 
+    _validate_token_probability_fn(token_probability_fn, prompt)
+
     def one_step(length, prompt):
         pred = token_probability_fn(prompt[:, :length])
         if from_logits:


Test patch:
diff --git a/keras_nlp/models/gpt2/gpt2_models_test.py b/keras_nlp/models/gpt2/gpt2_models_test.py
index 3b0b3a34..f8c0765a 100644
--- a/keras_nlp/models/gpt2/gpt2_models_test.py
+++ b/keras_nlp/models/gpt2/gpt2_models_test.py
@@ -19,12 +19,12 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.gpt2.gpt2_models import Gpt2
+from keras_nlp.models.gpt2.gpt2_models import GPT2
 
 
-class Gpt2Test(tf.test.TestCase, parameterized.TestCase):
+class GPT2Test(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = Gpt2(
+        self.model = GPT2(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
@@ -88,7 +88,7 @@ class Gpt2Test(tf.test.TestCase, parameterized.TestCase):
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
-        self.assertIsInstance(restored_model, Gpt2)
+        self.assertIsInstance(restored_model, GPT2)
 
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
diff --git a/keras_nlp/models/roberta/roberta_preprocessing_test.py b/keras_nlp/models/roberta/roberta_preprocessing_test.py
new file mode 100644
index 00000000..d325a36b
--- /dev/null
+++ b/keras_nlp/models/roberta/roberta_preprocessing_test.py
@@ -0,0 +1,152 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for RoBERTa preprocessing layers."""
+
+import os
+
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.roberta.roberta_preprocessing import RobertaPreprocessor
+
+
+class RobertaPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        vocab = {
+            "
": 0,
+            "<pad>": 1,
+            "
": 2,
+            "Ġair": 3,
+            "plane": 4,
+            "Ġat": 5,
+            "port": 6,
+            "Ġkoh": 7,
+            "li": 8,
+            "Ġis": 9,
+            "Ġthe": 10,
+            "Ġbest": 11,
+        }
+
+        merges = ["Ġ a", "Ġ t", "Ġ k", "Ġ i", "Ġ b", "Ġa i", "p l", "n e"]
+        merges += ["Ġa t", "p o", "r t", "o h", "l i", "Ġi s", "Ġb e", "s t"]
+        merges += ["Ġt h", "Ġai r", "pl a", "Ġk oh", "Ġth e", "Ġbe st", "po rt"]
+        merges += ["pla ne"]
+
+        self.preprocessor = RobertaPreprocessor(
+            vocabulary=vocab,
+            merges=merges,
+            sequence_length=12,
+        )
+
+    def test_tokenize(self):
+        input_data = [" airplane at airport"]
+
+        output = self.preprocessor(input_data)
+        self.assertAllEqual(
+            output["token_ids"], [0, 3, 4, 5, 3, 6, 2, 1, 1, 1, 1, 1]
+        )
+        self.assertAllEqual(
+            output["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0]
+        )
+
+    def test_tokenize_batch(self):
+        input_data = tf.constant(
+            [
+                " airplane at airport",
+                " airplane at airport",
+                " airplane at airport",
+                " airplane at airport",
+            ]
+        )
+
+        output = self.preprocessor(input_data)
+        self.assertAllEqual(
+            output["token_ids"],
+            [[0, 3, 4, 5, 3, 6, 2, 1, 1, 1, 1, 1]] * 4,
+        )
+
+        self.assertAllEqual(
+            output["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0]] * 4
+        )
+
+    def test_tokenize_multiple_sentences(self):
+        sentence_one = " airplane at airport"
+        sentence_two = " kohli is the best"
+
+        output = self.preprocessor((sentence_one, sentence_two))
+        print(output)
+        self.assertAllEqual(
+            output["token_ids"], [0, 3, 4, 5, 3, 2, 2, 7, 8, 9, 10, 2]
+        )
+        self.assertAllEqual(
+            output["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
+        )
+
+    def test_tokenize_multiple_batched_sentences(self):
+        sentence_one = tf.constant(
+            [
+                " airplane at airport",
+                " airplane at airport",
+                " airplane at airport",
+                " airplane at airport",
+            ]
+        )
+        sentence_two = tf.constant(
+            [
+                " kohli is the best",
+                " kohli is the best",
+                " kohli is the best",
+                " kohli is the best",
+            ]
+        )
+
+        output = self.preprocessor((sentence_one, sentence_two))
+        self.assertAllEqual(
+            output["token_ids"],
+            [[0, 3, 4, 5, 3, 2, 2, 7, 8, 9, 10, 2]] * 4,
+        )
+        self.assertAllEqual(
+            output["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]] * 4
+        )
+
+    def test_detokenize(self):
+        input_tokens = [[0, 3, 4, 5, 3, 6, 2, 1, 1, 1, 1, 1]]
+        output = self.preprocessor.tokenizer.detokenize(input_tokens)
+        self.assertAllEqual(
+            output, ["
 airplane at airport
<pad><pad><pad><pad><pad>"]
+        )
+
+    def test_vocabulary_size(self):
+        self.assertEqual(self.preprocessor.vocabulary_size(), 12)
+
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
+        input_data = tf.constant([" airplane at airport"])
+
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = self.preprocessor(inputs)
+        model = keras.Model(inputs, outputs)
+
+        path = os.path.join(self.get_temp_dir(), "model")
+        model.save(path, save_format=save_format)
+
+        restored_model = keras.models.load_model(path)
+        self.assertAllEqual(
+            model(input_data)["token_ids"],
+            restored_model(input_data)["token_ids"],
+        )
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessing_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessing_test.py
index fefc00f7..9945f09e 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessing_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessing_test.py
@@ -11,7 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for BERT preprocessing layers."""
+
+"""Tests for XLM-RoBERTa preprocessing layers."""
 
 import io
 
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_tasks_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_tasks_test.py
new file mode 100644
index 00000000..393bcd4c
--- /dev/null
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_tasks_test.py
@@ -0,0 +1,85 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for XLM-RoBERTa task specific models and heads."""
+
+import os
+
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.xlm_roberta.xlm_roberta_models import XLMRoberta
+from keras_nlp.models.xlm_roberta.xlm_roberta_tasks import XLMRobertaClassifier
+
+
+class XLMRobertaClassifierTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        self.model = XLMRoberta(
+            vocabulary_size=1000,
+            num_layers=2,
+            num_heads=2,
+            hidden_dim=64,
+            intermediate_dim=128,
+            max_sequence_length=128,
+        )
+        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"
+            ),
+        }
+
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
+    def test_valid_call_classifier(self):
+        classifier = XLMRobertaClassifier(self.model, 4, 128, name="classifier")
+        classifier(self.input_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_xlm_roberta_classifier_compile(self, jit_compile):
+        classifier = XLMRobertaClassifier(self.model, 4, 128, name="classifier")
+        classifier.compile(jit_compile=jit_compile)
+        classifier.predict(self.input_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_xlm_roberta_classifier_compile_batched_ds(self, jit_compile):
+        classifier = XLMRobertaClassifier(self.model, 4, 128, name="classifier")
+        classifier.compile(jit_compile=jit_compile)
+        classifier.predict(self.input_dataset)
+
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
+        classifier = XLMRobertaClassifier(self.model, 4, 128, name="classifier")
+        classifier_output = classifier(self.input_batch)
+        save_path = os.path.join(self.get_temp_dir(), "model")
+        classifier.save(save_path, save_format)
+        restored_classifier = keras.models.load_model(save_path)
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_classifier, XLMRobertaClassifier)
+
+        # Check that output matches.
+        restored_output = restored_classifier(self.input_batch)
+        self.assertAllClose(classifier_output, restored_output)
diff --git a/keras_nlp/tokenizers/byte_pair_tokenizer_test.py b/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
index 243891e5..6c54dfca 100644
--- a/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
+++ b/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
@@ -57,8 +57,14 @@ class BytePairTokenizerTest(tf.test.TestCase, parameterized.TestCase):
         encoded = self.tokenizer.tokenize(input_data)
         self.assertAllEqual(encoded, [31876, 4])
 
-    def test_detokenize(self):
-        input_data = ["brown."]
+    def test_detokenize_scalar_input(self):
+        input_data = ["quick brown fox."]
+        encoded = self.tokenizer.tokenize(input_data)
+        decoded = self.tokenizer.detokenize(encoded)
+        self.assertAllEqual(input_data, decoded)
+
+    def test_detokenize_list_input(self):
+        input_data = ["quick brown fox.", "slow black bear."]
         encoded = self.tokenizer.tokenize(input_data)
         decoded = self.tokenizer.detokenize(encoded)
         self.assertAllEqual(input_data, decoded)
diff --git a/keras_nlp/utils/text_generation_test.py b/keras_nlp/utils/text_generation_test.py
index f1882df8..9cd9b405 100644
--- a/keras_nlp/utils/text_generation_test.py
+++ b/keras_nlp/utils/text_generation_test.py
@@ -68,6 +68,8 @@ class GreedySearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
     def test_generate_with_ragged_prompt(self):
         def token_probability_fn(inputs):
+            # Assert that user function is passed only dense tensors.
+            self.assertIsInstance(inputs, tf.Tensor)
             prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, 2, axis=0)
 
@@ -260,6 +262,8 @@ class BeamSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
     def test_generate_with_ragged_prompt(self):
         def token_probability_fn(inputs):
+            # Assert that user function is passed only dense tensors.
+            self.assertIsInstance(inputs, tf.Tensor)
             repeat_times = tf.shape(inputs)[0]
             prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, repeat_times, axis=0)
@@ -446,6 +450,8 @@ class RandomSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
     def test_generate_with_ragged_prompt(self):
         def token_probability_fn(inputs):
+            # Assert that user function is passed only dense tensors.
+            self.assertIsInstance(inputs, tf.Tensor)
             prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, 2, axis=0)
 
@@ -667,6 +673,8 @@ class TopKSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
     def test_generate_with_ragged_prompt(self):
         def token_probability_fn(inputs):
+            # Assert that user function is passed only dense tensors.
+            self.assertIsInstance(inputs, tf.Tensor)
             prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, 2, axis=0)
 
@@ -906,6 +914,8 @@ class TopPSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
     def test_generate_with_ragged_prompt(self):
         def token_probability_fn(inputs):
+            # Assert that user function is passed only dense tensors.
+            self.assertIsInstance(inputs, tf.Tensor)
             prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, 2, axis=0)