Repository: keras-team/keras-nlp

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


Discussion:

Resolving patch:
diff --git a/keras_nlp/models/roberta/roberta_classifier.py b/keras_nlp/models/roberta/roberta_classifier.py
--- a/keras_nlp/models/roberta/roberta_classifier.py
+++ b/keras_nlp/models/roberta/roberta_classifier.py
@@ -19,14 +19,26 @@
 
 from keras_nlp.models.roberta.roberta_backbone import RobertaBackbone
 from keras_nlp.models.roberta.roberta_backbone import roberta_kernel_initializer
+from keras_nlp.models.roberta.roberta_preprocessor import RobertaPreprocessor
 from keras_nlp.models.roberta.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 RobertaClassifier(keras.Model):
-    """RoBERTa encoder model with a classification head.
+class RobertaClassifier(PipelineModel):
+    """An end-to-end RoBERTa model for classification tasks.
+
+    This model attaches a classification head to a
+    `keras_nlp.model.RobertaBackbone`, 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
@@ -34,20 +46,25 @@ class RobertaClassifier(keras.Model):
     [here](https://github.com/facebookresearch/fairseq).
 
     Args:
-        backbone: A `keras_nlp.models.Roberta` instance.
+        backbone: A `keras_nlp.models.RobertaBackbone` instance.
         num_classes: int. Number of classes to predict.
         hidden_dim: int. The size of the pooler layer.
+        preprocessor: A `keras_nlp.models.BertPreprocessor` or `None`. If
+            `None`, this model will not apply preprocessing, and inputs should
+            be preprocessed before calling the model.
 
     Example usage:
     ```python
-    input_data = {
-        "token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
+    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], shape=(1, 12)),
+            [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]] * 2, shape=(2, 12)
+        ),
     }
+    labels = [0, 3]
 
-    # Randomly initialized RoBERTa encoder
-    model = keras_nlp.models.RobertaBackbone(
+    # Randomly initialize a RoBERTa encoder
+    backbone = keras_nlp.models.RobertaBackbone(
         vocabulary_size=50265,
         num_layers=12,
         num_heads=12,
@@ -55,8 +72,16 @@ class RobertaClassifier(keras.Model):
         intermediate_dim=3072,
         max_sequence_length=12
     )
-    classifier = keras_nlp.models.RobertaClassifier(model, 4)
-    logits = classifier(input_data)
+    # Create a RoBERTa classifier and fit the data.
+    classifier = keras_nlp.models.RobertaClassifier(
+        backbone,
+        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
@@ -69,6 +94,7 @@ def __init__(
         num_classes=2,
         hidden_dim=None,
         dropout=0.0,
+        preprocessor=None,
         **kwargs,
     ):
         inputs = backbone.input
@@ -91,6 +117,7 @@ def __init__(
         super().__init__(
             inputs=inputs,
             outputs=outputs,
+            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
         # All references to `self` below this line
@@ -98,15 +125,25 @@ def __init__(
         self.num_classes = num_classes
         self.hidden_dim = hidden_dim
         self.dropout = dropout
+        self._preprocessor = preprocessor
+
+    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.Roberta` instance providing the encoder submodel."""
+        """The associated `keras_nlp.models.RobertaBackbone`."""
         return self._backbone
 
+    @property
+    def preprocessor(self):
+        """A `keras_nlp.models.RobertaPreprocessor` for preprocessing inputs."""
+        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,
@@ -116,8 +153,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
@@ -134,43 +177,92 @@ 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]
+
+        # Create a RobertClassifier and fit your data.
+        classifier = keras_nlp.models.RobertClassifier.from_preset(
+            "roberta_base",
+            num_classes=4,
+        )
+        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 backbone architecture and weights from preset
-        classifier = keras_nlp.models.RobertaClassifier.from_preset(
+        # Use a shorter sequence length.
+        preprocessor = keras_nlp.models.RobertPreprocessor.from_preset(
+            "roberta_base",
+            sequence_length=128,
+        )
+
+        # Create a RobertClassifier and fit your data.
+        classifier = keras_nlp.models.RobertClassifier.from_preset(
             "roberta_base",
             num_classes=4,
+            preprocessor=preprocessor,
+        )
+        classifier.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
         )
-        output = classifier(input_data)
+        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]
 
-        # Load randomly initialized model from preset architecture
-        classifier = keras_nlp.models.RobertaClassifier.from_preset(
+        # Create a RoBERTa classifier and fit your data.
+        classifier = keras_nlp.models.RobertClassifier.from_preset(
             "roberta_base",
-            load_weights=False,
             num_classes=4,
+            preprocessor=None,
         )
-        output = classifier(input_data)
+        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"] = RobertaPreprocessor.from_preset(preset)
+
         # Check if preset is backbone-only model
         if preset in RobertaBackbone.presets:
             backbone = RobertaBackbone.from_preset(preset, load_weights)
             return cls(backbone, **kwargs)
 
-        # Otherwise must be one of class presets
-        # Currently no classifier-level presets, so must throw.
+        # Otherwise must be one of class presets.
+        # Currently no classifier-level presets, so we raise ValueError.
         if preset not in cls.presets:
             raise ValueError(
                 "`preset` must be one of "
diff --git a/keras_nlp/models/roberta/roberta_preprocessor.py b/keras_nlp/models/roberta/roberta_preprocessor.py
--- a/keras_nlp/models/roberta/roberta_preprocessor.py
+++ b/keras_nlp/models/roberta/roberta_preprocessor.py
@@ -192,7 +192,7 @@ def get_config(self):
 
     @classmethod
     def from_config(cls, config):
-        if "tokenizer" in config:
+        if "tokenizer" in config and isinstance(config["tokenizer"], dict):
             config["tokenizer"] = keras.layers.deserialize(config["tokenizer"])
         return cls(**config)
 


Test patch:
diff --git a/keras_nlp/models/roberta/roberta_classifier_test.py b/keras_nlp/models/roberta/roberta_classifier_test.py
--- a/keras_nlp/models/roberta/roberta_classifier_test.py
+++ b/keras_nlp/models/roberta/roberta_classifier_test.py
@@ -21,11 +21,13 @@
 
 from keras_nlp.models.roberta.roberta_backbone import RobertaBackbone
 from keras_nlp.models.roberta.roberta_classifier import RobertaClassifier
+from keras_nlp.models.roberta.roberta_preprocessor import RobertaPreprocessor
+from keras_nlp.models.roberta.roberta_tokenizer import RobertaTokenizer
 
 
 class RobertaClassifierTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = RobertaBackbone(
+        self.backbone = RobertaBackbone(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
@@ -33,54 +35,105 @@ 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.vocab = {
+            "
": 0,
+            "<pad>": 1,
+            "
": 2,
+            "Ġair": 3,
+            "plane": 4,
+            "Ġat": 5,
+            "port": 6,
+            "Ġkoh": 7,
+            "li": 8,
+            "Ġis": 9,
+            "Ġthe": 10,
+            "Ġbest": 11,
         }
 
-        self.input_dataset = tf.data.Dataset.from_tensor_slices(
-            self.input_batch
+        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.merges = merges
+        self.preprocessor = RobertaPreprocessor(
+            RobertaTokenizer(vocabulary=self.vocab, merges=self.merges),
+            sequence_length=8,
+        )
+        self.classifier = RobertaClassifier(
+            self.backbone,
+            4,
+            preprocessor=self.preprocessor,
+        )
+        self.classifier_no_preprocessing = RobertaClassifier(
+            self.backbone,
+            4,
+            preprocessor=None,
+        )
+
+        self.raw_batch = tf.constant(
+            [
+                " airplane at airport",
+                " the airplane is the best",
+                " the best airport",
+                " kohli is the best",
+            ]
+        )
+        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 = RobertaClassifier(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_roberta_classifier_compile(self, jit_compile):
-        model = RobertaClassifier(self.model, 4, 128, name="classifier")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_batch)
+    def test_roberta_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_roberta_classifier_compile_batched_ds(self, jit_compile):
-        model = RobertaClassifier(self.model, 4, 128, name="classifier")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_dataset)
+    def test_roberta_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_roberta_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_roberta_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):
-        model = RobertaClassifier(self.model, 4, 128, name="classifier")
-        model_output = model(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)
-        model.save(save_path, save_format=save_format)
+        self.classifier.save(save_path, save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, RobertaClassifier)
 
         # Check that output matches.
-        restored_output = restored_model(self.input_batch)
+        restored_output = restored_model.predict(self.raw_batch)
         self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/models/roberta/roberta_presets_test.py b/keras_nlp/models/roberta/roberta_presets_test.py
--- a/keras_nlp/models/roberta/roberta_presets_test.py
+++ b/keras_nlp/models/roberta/roberta_presets_test.py
@@ -70,15 +70,28 @@ def test_backbone_output(self, load_weights):
         ("preset_weights", True), ("random_weights", False)
     )
     def test_classifier_output(self, load_weights):
+        input_data = ["Let's rock!"]
+        model = RobertaClassifier.from_preset(
+            "roberta_base", load_weights=load_weights
+        )
+        # Never assert output values, as the head weights are random.
+        model.predict(input_data)
+
+    @parameterized.named_parameters(
+        ("load_weights", True), ("no_load_weights", False)
+    )
+    def test_classifier_output_without_preprocessing(self, load_weights):
         input_data = {
-            "token_ids": tf.constant([[0, 133, 2119, 2]]),
+            "token_ids": tf.constant([[101, 1996, 4248, 102]]),
             "padding_mask": tf.constant([[1, 1, 1, 1]]),
         }
         model = RobertaClassifier.from_preset(
-            "roberta_base", load_weights=load_weights
+            "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(
         ("roberta_tokenizer", RobertaTokenizer),
@@ -137,6 +150,19 @@ def test_load_roberta_classifier(self, load_weights):
             classifier = RobertaClassifier.from_preset(
                 preset, num_classes=4, load_weights=load_weights
             )
+            input_data = ["The quick brown fox."]
+            classifier(input_data)
+
+    @parameterized.named_parameters(
+        ("load_weights", True), ("no_load_weights", False)
+    )
+    def test_load_roberta_classifier_without_preprocessing(self, load_weights):
+        for preset in RobertaClassifier.presets:
+            classifier = RobertaClassifier.from_preset(
+                preset,
+                preprocessor=None,
+                load_weights=load_weights,
+            )
             input_data = {
                 "token_ids": tf.random.uniform(
                     shape=(1, 512),
@@ -145,7 +171,7 @@ def test_load_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 RobertaTokenizer.presets: