Repository: keras-team/keras-nlp

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

Discussion:

Patch:
diff --git a/.cloudbuild/cloudbuild.yaml b/.cloudbuild/cloudbuild.yaml
index 20f6593c..58b6052e 100644
--- a/.cloudbuild/cloudbuild.yaml
+++ b/.cloudbuild/cloudbuild.yaml
@@ -29,7 +29,7 @@ steps:
     'install',
     'github.com/google/go-jsonnet/cmd/jsonnet@latest',
   ]
-- name: 'google/cloud-sdk'
+- name: 'gcr.io/cloud-builders/gcloud'
   id: clone-templates
   waitFor: ['-']
   entrypoint: 'git'
@@ -52,7 +52,7 @@ steps:
     '--ext-str', 'gcs_bucket=$_GCS_BUCKET',
     '-o', 'output.yaml',
   ]
-- name: 'google/cloud-sdk'
+- name: 'gcr.io/cloud-builders/gcloud'
   id: create-job
   waitFor:
   - push-image
diff --git a/keras_nlp/models/distil_bert/distil_bert_classifier.py b/keras_nlp/models/distil_bert/distil_bert_classifier.py
index 76137ef9..4d5f54c2 100644
--- a/keras_nlp/models/distil_bert/distil_bert_classifier.py
+++ b/keras_nlp/models/distil_bert/distil_bert_classifier.py
@@ -21,14 +21,28 @@ from keras_nlp.models.distil_bert.distil_bert_backbone import DistilBertBackbone
 from keras_nlp.models.distil_bert.distil_bert_backbone import (
     distilbert_kernel_initializer,
 )
+from keras_nlp.models.distil_bert.distil_bert_preprocessor import (
+    DistilBertPreprocessor,
+)
 from keras_nlp.models.distil_bert.distil_bert_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 DistilBertClassifier(keras.Model):
-    """DistilBERT encoder model with a classification head.
+class DistilBertClassifier(PipelineModel):
+    """An end-to-end DistilBERT model for classification tasks.
+
+    This model attaches a classification head to a
+    `keras_nlp.model.DistilBertBackbone` 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
@@ -39,19 +53,23 @@ class DistilBertClassifier(keras.Model):
         backbone: A `keras_nlp.models.DistilBert` instance.
         num_classes: int. Number of classes to predict.
         hidden_dim: int. The size of the pooler layer.
+        preprocessor: A `keras_nlp.models.DistilBertPreprocessor` 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),
         "token_ids": tf.random.uniform(
-            shape=(1, 12), dtype=tf.int64, maxval=vocabulary_size),
+            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 DistilBERT encoder
-    model = keras_nlp.models.DistilBertBackbone(
+    backbone = keras_nlp.models.DistilBertBackbone(
         vocabulary_size=30552,
         num_layers=6,
         num_heads=12,
@@ -59,8 +77,14 @@ class DistilBertClassifier(keras.Model):
         intermediate_dim=3072,
         max_sequence_length=512
     )
-    classifier = keras_nlp.models.DistilBertClassifier(model, 4)
-    logits = classifier(input_data)
+
+    # Create a DistilBERT classifier and fit your data.
+    classifier = keras_nlp.models.DistilBertClassifier(
+        backbone,
+        num_classes=4,
+        preprocessor=None,
+    )
+    classifier.fit(x=preprocessed_features, y=labels, batch_size=2)
 
     # Access backbone programatically (e.g., to change `trainable`)
     classifier.backbone.trainable = False
@@ -73,6 +97,7 @@ class DistilBertClassifier(keras.Model):
         num_classes=2,
         hidden_dim=None,
         dropout=0.2,
+        preprocessor=None,
         **kwargs,
     ):
         inputs = backbone.input
@@ -97,22 +122,33 @@ class DistilBertClassifier(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.DistilBert` instance providing the encoder submodel."""
         return self._backbone
 
+    @property
+    def preprocessor(self):
+        """A `keras_nlp.models.BertPreprocessor` 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,
@@ -122,8 +158,14 @@ class DistilBertClassifier(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 +182,87 @@ class DistilBertClassifier(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 DistilBertClassifier and fit your data.
         classifier = keras_nlp.models.DistilBertClassifier.from_preset(
             "distil_bert_base_en_uncased",
             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)
+        ```
 
-        # Load randomly initialized model from preset architecture
+        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.BertPreprocessor.from_preset(
+            "bert_base_en_uncased",
+            sequence_length=128,
+        )
+        # Create a DistilBertClassifier and fit your data.
         classifier = keras_nlp.models.DistilBertClassifier.from_preset(
-            "distil_bert_base_en_uncased",
-            load_weights=False,
+            "bert_base_en_uncased",
             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),
+            "segment_ids": tf.constant(
+                [[0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0]] * 2, shape=(2, 12)
+            ),
+            "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 DistilBERT classifier and fit your data.
+        classifier = keras_nlp.models.DistilBertClassifier.from_preset(
+            "bert_base_en_uncased",
+            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"] = DistilBertPreprocessor.from_preset(preset)
+
         # Check if preset is backbone-only model
         if preset in DistilBertBackbone.presets:
             backbone = DistilBertBackbone.from_preset(preset, load_weights)


Test patch:
diff --git a/keras_nlp/models/distil_bert/distil_bert_classifier_test.py b/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
index 4af7fe0d..8f0d62be 100644
--- a/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
@@ -23,11 +23,17 @@ from keras_nlp.models.distil_bert.distil_bert_backbone import DistilBertBackbone
 from keras_nlp.models.distil_bert.distil_bert_classifier import (
     DistilBertClassifier,
 )
+from keras_nlp.models.distil_bert.distil_bert_preprocessor import (
+    DistilBertPreprocessor,
+)
+from keras_nlp.models.distil_bert.distil_bert_tokenizer import (
+    DistilBertTokenizer,
+)
 
 
 class DistilBertClassifierTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = DistilBertBackbone(
+        self.backbone = DistilBertBackbone(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
@@ -35,54 +41,87 @@ class DistilBertClassifierTest(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.vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
+        self.vocab += ["the", "quick", "brown", "fox", "."]
+        self.preprocessor = DistilBertPreprocessor(
+            DistilBertTokenizer(vocabulary=self.vocab),
+            sequence_length=8,
+        )
+        self.classifier = DistilBertClassifier(
+            self.backbone,
+            4,
+            preprocessor=self.preprocessor,
+        )
+        self.classifier_no_preprocessing = DistilBertClassifier(
+            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 smelly brown fox.",
+                "the old brown fox.",
+            ]
+        )
+        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 = DistilBertClassifier(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_distil_bert_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_distil_bert_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_distilbert_classifier_compile(self, jit_compile):
-        model = DistilBertClassifier(self.model, 4, 128, name="classifier")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_batch)
+    def test_bert_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_distilbert_classifier_compile_batched_ds(self, jit_compile):
-        model = DistilBertClassifier(self.model, 4, 128, name="classifier")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_dataset)
+    def test_bert_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 = DistilBertClassifier(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=save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, DistilBertClassifier)
 
         # 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/distil_bert/distil_bert_presets_test.py b/keras_nlp/models/distil_bert/distil_bert_presets_test.py
index 024710ca..5da425df 100644
--- a/keras_nlp/models/distil_bert/distil_bert_presets_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_presets_test.py
@@ -75,14 +75,27 @@ class DistilBertPresetSmokeTest(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 = DistilBertClassifier.from_preset(
+            "distil_bert_base_en_uncased",
+            load_weights=load_weights,
+        )
+        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([[101, 1996, 4248, 102]]),
             "padding_mask": tf.constant([[1, 1, 1, 1]]),
         }
         model = DistilBertClassifier.from_preset(
-            "distil_bert_base_en_uncased", load_weights=load_weights
+            "distil_bert_base_en_uncased",
+            load_weights=load_weights,
+            preprocessor=None,
         )
-        model(input_data)
+        model.predict(input_data)
 
     @parameterized.named_parameters(
         ("distilbert_tokenizer", DistilBertTokenizer),
@@ -139,7 +152,23 @@ class DistilBertPresetFullTest(tf.test.TestCase, parameterized.TestCase):
     def test_load_distilbert_classifier(self, load_weights):
         for preset in DistilBertClassifier.presets:
             classifier = DistilBertClassifier.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_distilbert_classifier_no_preprocessing(self, load_weights):
+        for preset in DistilBertClassifier.presets:
+            classifier = DistilBertClassifier.from_preset(
+                preset,
+                num_classes=4,
+                load_weights=load_weights,
+                preprocessor=None,
             )
             input_data = {
                 "token_ids": tf.random.uniform(
@@ -149,7 +178,7 @@ class DistilBertPresetFullTest(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 DistilBertTokenizer.presets: