Repository: keras-team/keras-nlp

Issue:
Allow XMLRobertaClasssifier to operate directly on raw strings
We should mirror the change in https://github.com/keras-team/keras-nlp/pull/485 for XLMRobertaClassifier, and allow fit/predict/evaluate to take in raw strings.

Discussion:

Patch:
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
index 2586ae81..967a56d6 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
@@ -19,14 +19,28 @@ from tensorflow import keras
 
 from keras_nlp.models.roberta.roberta_backbone import roberta_kernel_initializer
 from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRobertaBackbone
+from keras_nlp.models.xlm_roberta.xlm_roberta_preprocessor import (
+    XLMRobertaPreprocessor,
+)
 from keras_nlp.models.xlm_roberta.xlm_roberta_presets import backbone_presets
+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
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class XLMRobertaClassifier(keras.Model):
-    """XLM-RoBERTa encoder model with a classification head.
+class XLMRobertaClassifier(PipelineModel):
+    """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.
+
+    This model can optionally be configured with a `preprocessor` layer, in
+    which case it will automatically apply preprocessing to raw inputs during
+    `fit()`, `predict()`, and `evaluate()`. This is done by default when
+    creating the model with `from_preset()`.
 
     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
@@ -37,18 +51,23 @@ class XLMRobertaClassifier(keras.Model):
         backbone: A `keras_nlp.models.XLMRoberta` instance.
         num_classes: int. Number of classes to predict.
         hidden_dim: int. The size of the pooler layer.
+        preprocessor: A `keras_nlp.models.XLMRobertaPreprocessor` or `None`. If
+            `None`, this model will not apply preprocessing, and inputs should
+            be preprocessed before calling the model.
 
     Example usage:
     ```python
-    # Call classifier on the inputs.
-    input_data = {
-        "token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
+    preprocessed_features = {
+        "token_ids": tf.ones(shape=(2, 12), dtype=tf.int64),
+        "token_ids": tf.random.uniform(
+            shape=(2, 12), dtype=tf.int64, maxval=vocabulary_size),
         "padding_mask": tf.constant(
-            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+            [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]] * 2, shape=(1, 12)),
     }
+    labels = [0, 3]
 
     # Randomly initialized XLM-RoBERTa encoder
-    model = keras_nlp.models.XLMRobertaBackbone(
+    backbone = keras_nlp.models.XLMRobertaBackbone(
         vocabulary_size=250002,
         num_layers=12,
         num_heads=12,
@@ -56,11 +75,14 @@ class XLMRobertaClassifier(keras.Model):
         intermediate_dim=3072,
         max_sequence_length=12
     )
+
+    # Create a XLM-RoBERTa classifier and fit your data.
     classifier = keras_nlp.models.XLMRobertaClassifier(
-        backbone=model,
+        backbone,
         num_classes=4,
+        preprocessor=None,
     )
-    logits = classifier(input_data)
+    classifier.fit(x=preprocessed_features, y=labels, batch_size=2)
 
     # Access backbone programatically (e.g., to change `trainable`)
     classifier.backbone.trainable = False
@@ -73,6 +95,7 @@ class XLMRobertaClassifier(keras.Model):
         num_classes=2,
         hidden_dim=None,
         dropout=0.0,
+        preprocessor=None,
         **kwargs,
     ):
         inputs = backbone.input
@@ -95,24 +118,33 @@ class XLMRobertaClassifier(keras.Model):
         super().__init__(
             inputs=inputs,
             outputs=outputs,
+            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
         # All references to `self` below this line
         self._backbone = backbone
+        self._preprocessor = preprocessor
         self.num_classes = num_classes
         self.hidden_dim = hidden_dim
         self.dropout = dropout
 
+    def preprocess_samples(self, x, y=None, sample_weight=None):
+        return self.preprocessor(x, y=y, sample_weight=sample_weight)
+
     @property
     def backbone(self):
-        """A `keras_nlp.models.XLMRoberta` instance providing the encoder
-        submodel.
-        """
+        """A `keras_nlp.models.XLMRobertaBackbone` submodel."""
         return self._backbone
 
+    @property
+    def preprocessor(self):
+        """A `keras_nlp.models.XLMRobertaPreprocessor` preprocessing layer."""
+        return self._preprocessor
+
     def get_config(self):
         return {
             "backbone": keras.layers.serialize(self.backbone),
+            "preprocessor": keras.layers.serialize(self.preprocessor),
             "num_classes": self.num_classes,
             "hidden_dim": self.hidden_dim,
             "dropout": self.dropout,
@@ -122,8 +154,14 @@ class XLMRobertaClassifier(keras.Model):
 
     @classmethod
     def from_config(cls, config):
-        if "backbone" in config:
+        if "backbone" in config and isinstance(config["backbone"], dict):
             config["backbone"] = keras.layers.deserialize(config["backbone"])
+        if "preprocessor" in config and isinstance(
+            config["preprocessor"], dict
+        ):
+            config["preprocessor"] = keras.layers.deserialize(
+                config["preprocessor"]
+            )
         return cls(**config)
 
     @classproperty
@@ -140,36 +178,85 @@ class XLMRobertaClassifier(keras.Model):
     ):
         """Create a classification model from a preset architecture and weights.
 
+        By default, this method will automatically create a `preprocessor`
+        layer to preprocess raw inputs during `fit()`, `predict()`, and
+        `evaluate()`. If you would like to disable this behavior, pass
+        `preprocessor=None`.
+
         Args:
             preset: string. Must be one of {{names}}.
             load_weights: Whether to load pre-trained weights into model.
                 Defaults to `True`.
 
         Examples:
+
+        Raw string inputs.
         ```python
-        input_data = {
-            "token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
-            "padding_mask": tf.constant(
-                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-        }
+        # 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]
 
-        # Load backbone architecture and weights from preset
+        # Create a XLMRobertaClassifier and fit your data.
         classifier = keras_nlp.models.XLMRobertaClassifier.from_preset(
             "xlm_roberta_base",
             num_classes=4,
         )
-        output = classifier(input_data)
+        classifier.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        )
+        classifier.fit(x=features, y=labels, 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]
 
-        # Load randomly initialized model from preset architecture
+        # Use a shorter sequence length.
+        preprocessor = keras_nlp.models.XLMRobertaPreprocessor.from_preset(
+            "xlm_roberta_base",
+            sequence_length=128,
+        )
+
+        # Create a XLMRobertaClassifier and fit your data.
         classifier = keras_nlp.models.XLMRobertaClassifier.from_preset(
             "xlm_roberta_base",
-            load_weights=False,
             num_classes=4,
+            preprocessor=preprocessor,
         )
-        output = classifier(input_data)
+        classifier.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        )
+        classifier.fit(x=features, y=labels, batch_size=2)
+        ```
+
+        Preprocessed inputs.
+        ```python
+        # Create a dataset with preprocessed features in an `(x, y)` format.
+        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]
+
+        # Create a XLMRobertaClassifier and fit your data.
+        classifier = keras_nlp.models.XLMRobertaClassifier.from_preset(
+            "xlm_roberta_base",
+            num_classes=4,
+            preprocessor=None,
+        )
+        classifier.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        )
+        classifier.fit(x=preprocessed_features, y=labels, batch_size=2)
         ```
         """
+        if "preprocessor" not in kwargs:
+            kwargs["preprocessor"] = XLMRobertaPreprocessor.from_preset(preset)
+
         # Check if preset is backbone-only model
         if preset in XLMRobertaBackbone.presets:
             backbone = XLMRobertaBackbone.from_preset(preset, load_weights)


Test patch:
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
index f3fabd56..a444c96c 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
@@ -13,8 +13,10 @@
 # limitations under the License.
 """Tests for XLM-RoBERTa task specific models and heads."""
 
+import io
 import os
 
+import sentencepiece
 import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
@@ -23,11 +25,34 @@ from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRobertaBackbone
 from keras_nlp.models.xlm_roberta.xlm_roberta_classifier import (
     XLMRobertaClassifier,
 )
+from keras_nlp.models.xlm_roberta.xlm_roberta_preprocessor import (
+    XLMRobertaPreprocessor,
+)
+from keras_nlp.models.xlm_roberta.xlm_roberta_preprocessor import (
+    XLMRobertaTokenizer,
+)
 
 
 class XLMRobertaClassifierTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = XLMRobertaBackbone(
+        bytes_io = io.BytesIO()
+        vocab_data = tf.data.Dataset.from_tensor_slices(
+            ["the quick brown fox", "the earth is round"]
+        )
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=vocab_data.as_numpy_iterator(),
+            model_writer=bytes_io,
+            vocab_size=10,
+            model_type="WORD",
+            unk_id=0,
+            bos_id=1,
+            eos_id=2,
+        )
+        self.preprocessor = XLMRobertaPreprocessor(
+            tokenizer=XLMRobertaTokenizer(proto=bytes_io.getvalue()),
+            sequence_length=12,
+        )
+        self.backbone = XLMRobertaBackbone(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
@@ -35,54 +60,81 @@ class XLMRobertaClassifierTest(tf.test.TestCase, parameterized.TestCase):
             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.classifier = XLMRobertaClassifier(
+            self.backbone,
+            4,
+            preprocessor=self.preprocessor,
+        )
+        self.classifier_no_preprocessing = XLMRobertaClassifier(
+            self.backbone,
+            4,
+            preprocessor=None,
+        )
 
-        self.input_dataset = tf.data.Dataset.from_tensor_slices(
-            self.input_batch
+        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,)))
         ).batch(2)
+        self.preprocessed_dataset = self.raw_dataset.map(self.preprocessor)
 
     def test_valid_call_classifier(self):
-        classifier = XLMRobertaClassifier(self.model, 4, 128, name="classifier")
-        classifier(self.input_batch)
+        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)
+        self.classifier.predict(self.raw_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)
+    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)
 
     @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)
+    def test_classifier_fit(self, jit_compile):
+        self.classifier.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            jit_compile=jit_compile,
+        )
+        self.classifier.fit(self.raw_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,
+        )
+        self.classifier_no_preprocessing.fit(self.preprocessed_dataset)
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saved_model(self, save_format, filename):
-        classifier = XLMRobertaClassifier(self.model, 4, 128, name="classifier")
-        classifier_output = classifier(self.input_batch)
+    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)
-        classifier.save(save_path, save_format=save_format)
-        restored_classifier = keras.models.load_model(save_path)
+        self.classifier.save(save_path, save_format=save_format)
+        restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
-        self.assertIsInstance(restored_classifier, XLMRobertaClassifier)
+        self.assertIsInstance(restored_model, XLMRobertaClassifier)
 
         # Check that output matches.
-        restored_output = restored_classifier(self.input_batch)
-        self.assertAllClose(classifier_output, restored_output)
+        restored_output = restored_model.predict(self.raw_batch)
+        self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_presets_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_presets_test.py
index fa578673..6ebe3f7e 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_presets_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_presets_test.py
@@ -76,15 +76,28 @@ class XLMRobertaPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
         ("preset_weights", True), ("random_weights", False)
     )
     def test_classifier_output(self, load_weights):
+        input_data = tf.constant(["The quick brown fox."])
+        model = XLMRobertaClassifier.from_preset(
+            "xlm_roberta_base", load_weights=load_weights
+        )
+        # Never assert output values, as the head weights are random.
+        model.predict(input_data)
+
+    @parameterized.named_parameters(
+        ("preset_weights", True), ("random_weights", False)
+    )
+    def test_classifier_output_without_preprocessing(self, load_weights):
         input_data = {
             "token_ids": tf.constant([[0, 581, 63773, 2]]),
             "padding_mask": tf.constant([[1, 1, 1, 1]]),
         }
         model = XLMRobertaClassifier.from_preset(
-            "xlm_roberta_base", load_weights=load_weights
+            "xlm_roberta_base",
+            load_weights=load_weights,
+            preprocessor=None,
         )
         # Never assert output values, as the head weights are random.
-        model(input_data)
+        model.predict(input_data)
 
     @parameterized.named_parameters(
         ("xlm_roberta_tokenizer", XLMRobertaTokenizer),
@@ -141,7 +154,25 @@ class XLMRobertaPresetFullTest(tf.test.TestCase, parameterized.TestCase):
     def test_load_xlm_roberta_classifier(self, load_weights):
         for preset in XLMRobertaClassifier.presets:
             classifier = XLMRobertaClassifier.from_preset(
-                preset, num_classes=4, load_weights=load_weights
+                preset,
+                num_classes=4,
+                load_weights=load_weights,
+            )
+            input_data = tf.constant(["This quick brown fox"])
+            classifier.predict(input_data)
+
+    @parameterized.named_parameters(
+        ("preset_weights", True), ("random_weights", False)
+    )
+    def test_load_xlm_roberta_classifier_without_preprocessing(
+        self, load_weights
+    ):
+        for preset in XLMRobertaClassifier.presets:
+            classifier = XLMRobertaClassifier.from_preset(
+                preset,
+                num_classes=4,
+                load_weights=load_weights,
+                preprocessor=None,
             )
             input_data = {
                 "token_ids": tf.random.uniform(
@@ -151,7 +182,7 @@ class XLMRobertaPresetFullTest(tf.test.TestCase, parameterized.TestCase):
                 ),
                 "padding_mask": tf.constant([1] * 512, shape=(1, 512)),
             }
-            classifier(input_data)
+            classifier.predict(input_data)
 
     def test_load_tokenizers(self):
         for preset in XLMRobertaTokenizer.presets: