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:
Hey [mattdangerw](https://github.com/mattdangerw) , can I take this issue ?
/assign
Sure thing! You will basically want to mirror the changes here https://github.com/keras-team/keras-nlp/pull/485 for the XLMRobertaClassifier class. Should be mostly 1-1, though you can ignore any part that is about `classifier_presets` (we don't have any weights that include the classification head for `XMLRobertaClasssifier`).

Note that this will be somewhat time sensitive, we will probably try to land these issues on the first half of next week. But if that timing is not good for you, will be plenty more contributor issues coming up! So no pressure to take this on!

Thanks!
> Sure thing! You will basically want to mirror the changes here #485 for the XLMRobertaClassifier class. Should be mostly 1-1, though you can ignore any part that is about `classifier_presets` (we don't have any weights that include the classification head for `XMLRobertaClasssifier`).
> 
> Note that this will be somewhat time sensitive, we will probably try to land these issues on the first half of next week. But if that timing is not good for you, will be plenty more contributor issues coming up! So no pressure to take this on!
> 
> Thanks!

thanks for the information, So the new issues are on the first half of next week right? 
Can I make the changes right now regarding this issue because you mention it is  time sensitive?
> Can I make the changes right now

No need to rush too much!

This is just time sensitive because we are lining up our 0.4 release. If you think you have time to take this on between now and this Wednesday (the 14th), please go ahead! If not, this issue might be best left for someone else to to grab.

Thanks!
Going to take this on! We will open up more contributions welcome issue soon!

Resolving patch:
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
--- 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 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 @@ def __init__(
         num_classes=2,
         hidden_dim=None,
         dropout=0.0,
+        preprocessor=None,
         **kwargs,
     ):
         inputs = backbone.input
@@ -95,24 +118,33 @@ def __init__(
         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 @@ def get_config(self):
 
     @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 @@ def from_preset(
     ):
         """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
--- 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_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 @@ def setUp(self):
             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
--- 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 @@ def test_backbone_output(self, load_weights):
         ("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 @@ def test_load_xlm_roberta(self, load_weights):
     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 @@ def test_load_xlm_roberta_classifier(self, load_weights):
                 ),
                 "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: