Repository: keras-team/keras-nlp

Issue:
Add `AlbertClassifier`
Similar to `BertClassifier`.



Discussion:
shivance: @jbischof This PR is ready for review.

Please consider following points while reviewing
1. Albert Presets are not implemented yet, so in classproperty `from_preset()` of AlbertClassifier, I've raised a `NotImplementedError()`, please let me know your views on it!
2. The Tests pass locally.
shivance: @jbischof Sorry for the delay, I've been traveling for a week, and have 2 #648 #436 pending. Will soon follow up with them.
shivance: @jbischof This PR is ready for review as well !
shivance: CI is green @jbischof 
shivance: 1. `Albert` doesn't have classifier presets yet, thus I had to remove usage of classifier_presets from albert_classifier.py
2. when I am running tests for albert_presets_test.py , all are getting skipped on local
shivance: @jbischof This PR is ready for review.
shivance: @mattdangerw @jbischof Please review !

shivance: Still WIP
shivance: Ready for review @mattdangerw @jbischof 
shivance: CI is almost green.
Not sure why nightly is failing @mattdangerw @jbischof 
mattdangerw: Thanks so much @shivance ! Not sure what is going on with nightly, but I suspect it is unrelated to your PR.

Nightly testing is important for us to have, but it will definitely surface errors in the upstream tf release as often as our own. I will take a closer look today. Thanks!
shivance: @mattdangerw How does nightly differ from dev and stable release?
mattdangerw: @shivance only difference is that nightly will test against the [tf-nightly](https://pypi.org/project/tf-nightly/) and [tensorflow-text-nightly](https://pypi.org/project/tensorflow-text-nightly/) packages. Which are nightly "unstable" releases of our upstream dependencies. Often what will happen is the `tf-nightly` package will break our testing because of a bug in a recent tf change.
shivance: Got it ! Thanks ! @mattdangerw 
By the way, I'm really interested in GSoC this year at KerasNLP ( I love working here and luckily kerasnlp participate in gsoc last year under tensorflow). Will it participate this year too?

What are potential projects ? (My guess is, ideally it should align with the current roadmap of KerasNLP)
mattdangerw: > By the way, I'm really interested in GSoC this year at KerasNLP

I'm pretty sure yes! But we are just figuring out all the details right now actually, will follow up when we have more.

Also, confirmed that this nightly test is unrelated to your change. Will pull this in tomorrow, and thank you!

Patch:
diff --git a/keras_nlp/models/albert/albert_classifier.py b/keras_nlp/models/albert/albert_classifier.py
new file mode 100644
index 00000000..20f6e1a8
--- /dev/null
+++ b/keras_nlp/models/albert/albert_classifier.py
@@ -0,0 +1,197 @@
+# 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.
+"""ALBERT classification model."""
+
+import copy
+
+from tensorflow import keras
+
+from keras_nlp.models.albert.albert_backbone import AlbertBackbone
+from keras_nlp.models.albert.albert_backbone import albert_kernel_initializer
+from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
+from keras_nlp.models.albert.albert_presets import backbone_presets
+from keras_nlp.models.task import Task
+from keras_nlp.utils.python_utils import classproperty
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class AlbertClassifier(Task):
+    """An end-to-end ALBERT model for classification tasks
+
+    This model attaches a classification head to a `keras_nlp.model.AlbertBackbone`
+    backbone, 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.
+
+    Args:
+        backbone: A `keras_nlp.models.AlertBackbone` instance.
+        num_classes: int. Number of classes to predict.
+        dropout: float. The dropout probability value, applied after the dense
+            layer.
+        preprocessor: A `keras_nlp.models.AlbertPreprocessor` or `None`. If
+            `None`, this model will not apply preprocessing, and inputs should
+            be preprocessed before calling the model.
+
+    Examples:
+
+    Example usage.
+    ```python
+    # Define the preprocessed inputs.
+    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]
+
+    # Randomly initialize a ALBERT backbone.
+    backbone = AlbertBackbone(
+        vocabulary_size=1000,
+        num_layers=2,
+        num_heads=2,
+        embedding_dim=8,
+        hidden_dim=64,
+        intermediate_dim=128,
+        max_sequence_length=128,
+        name="encoder",
+    )
+
+    # Create a ALBERT classifier and fit your data.
+    classifier = keras_nlp.models.AlbertClassifier(
+        backbone,
+        num_classes=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
+
+    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.AlbertPreprocessor.from_preset(
+        "albert_base_en_uncased",
+        sequence_length=128,
+    )
+
+    # Create a AlbertClassifier and fit your data.
+    classifier = keras_nlp.models.AlbertClassifier.from_preset(
+        "albert_base_en_uncased",
+        num_classes=4,
+        preprocessor=preprocessor,
+    )
+    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 ALBERT classifier and fit your data.
+    classifier = keras_nlp.models.AlbertClassifier.from_preset(
+        "albert_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)
+    ```
+    """
+
+    def __init__(
+        self,
+        backbone,
+        num_classes=2,
+        dropout=0.1,
+        preprocessor=None,
+        **kwargs,
+    ):
+        inputs = backbone.input
+        pooled = backbone(inputs)["pooled_output"]
+        pooled = keras.layers.Dropout(dropout)(pooled)
+        outputs = keras.layers.Dense(
+            num_classes,
+            kernel_initializer=albert_kernel_initializer(),
+            name="logits",
+        )(pooled)
+        # Instantiate using Functional API Model constructor
+        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.dropout = dropout
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_classes": self.num_classes,
+                "dropout": self.dropout,
+            }
+        )
+
+        return config
+
+    @classproperty
+    def backbone_cls(cls):
+        return AlbertBackbone
+
+    @classproperty
+    def preprocessor_cls(cls):
+        return AlbertPreprocessor
+
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy({**backbone_presets})
diff --git a/setup.py b/setup.py
index ee064ba4..b7235bc7 100644
--- a/setup.py
+++ b/setup.py
@@ -38,27 +38,17 @@ setup(
         "absl-py",
         "numpy",
         "packaging",
-        # Don't require tensorflow on MacOS; tensorflow-macos will not
-        # satisfy the requirement.
-        "tensorflow; platform_system != 'Darwin'",
+        # Don't require tensorflow-text on MacOS, there are no binaries for ARM.
+        # Also, we rely on tensorflow *transitively* through tensorflow-text.
+        # This avoid a slowdown during `pip install keras-nlp` where pip would
+        # download many version of both libraries to find compatible versions.
         "tensorflow-text; platform_system != 'Darwin'",
     ],
     extras_require={
-        "tests": [
-            "black",
-            "flake8",
-            "isort",
-            "pytest",
-            "pytest-cov",
+        "extras": [
             "rouge-score",
             "sentencepiece",
         ],
-        "examples": [
-            "tensorflow_datasets",  # For GLUE in BERT example.
-            "nltk",
-            "wikiextractor",
-            "keras-tuner",
-        ],
     },
     # Supported Python versions
     python_requires=">=3.8",


Test patch:
diff --git a/keras_nlp/models/albert/albert_classifier_test.py b/keras_nlp/models/albert/albert_classifier_test.py
new file mode 100644
index 00000000..6f586b9c
--- /dev/null
+++ b/keras_nlp/models/albert/albert_classifier_test.py
@@ -0,0 +1,146 @@
+# 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 BERT classification model."""
+
+import io
+import os
+
+import sentencepiece
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.albert.albert_backbone import AlbertBackbone
+from keras_nlp.models.albert.albert_classifier import AlbertClassifier
+from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
+from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
+
+
+class AlbertClassifierTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        self.backbone = AlbertBackbone(
+            vocabulary_size=1000,
+            num_layers=2,
+            num_heads=2,
+            embedding_dim=8,
+            hidden_dim=64,
+            intermediate_dim=128,
+            max_sequence_length=128,
+            name="encoder",
+        )
+
+        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",
+            pad_id=0,
+            unk_id=1,
+            bos_id=2,
+            eos_id=3,
+            pad_piece="<pad>",
+            unk_piece="<unk>",
+            bos_piece="[CLS]",
+            eos_piece="[SEP]",
+        )
+        self.proto = bytes_io.getvalue()
+
+        tokenizer = AlbertTokenizer(proto=self.proto)
+
+        self.preprocessor = AlbertPreprocessor(
+            tokenizer=tokenizer,
+            sequence_length=8,
+        )
+        self.classifier = AlbertClassifier(
+            self.backbone,
+            4,
+            preprocessor=self.preprocessor,
+        )
+        self.classifier_no_preprocessing = AlbertClassifier(
+            self.backbone,
+            4,
+            preprocessor=None,
+        )
+
+        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):
+        self.classifier(self.preprocessed_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_albert_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_albert_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_albert_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_albert_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_output = self.classifier.predict(self.raw_batch)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        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, AlbertClassifier)
+
+        # Check that output matches.
+        restored_output = restored_model.predict(self.raw_batch)
+        self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/models/albert/albert_presets_test.py b/keras_nlp/models/albert/albert_presets_test.py
index 2ed61447..45862b16 100644
--- a/keras_nlp/models/albert/albert_presets_test.py
+++ b/keras_nlp/models/albert/albert_presets_test.py
@@ -18,6 +18,7 @@ import tensorflow as tf
 from absl.testing import parameterized
 
 from keras_nlp.models.albert.albert_backbone import AlbertBackbone
+from keras_nlp.models.albert.albert_classifier import AlbertClassifier
 from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
 from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
 
@@ -26,7 +27,6 @@ from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
 class AlbertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
     """
     A smoke test for ALBERT presets we run continuously.
-
     This only tests the smallest weights we have available. Run with:
     `pytest keras_nlp/models/albert/albert_presets_test.py --run_large`
     """
@@ -48,6 +48,35 @@ class AlbertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
         expected_outputs = [2, 13, 1, 3]
         self.assertAllEqual(outputs, expected_outputs)
 
+    @parameterized.named_parameters(
+        ("load_weights", True), ("no_load_weights", False)
+    )
+    def test_classifier_output(self, load_weights):
+        input_data = tf.constant(["The quick brown fox."])
+        model = AlbertClassifier.from_preset(
+            "albert_base_en_uncased",
+            load_weights=load_weights,
+        )
+        # We don't 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([[101, 1996, 4248, 102]]),
+            "segment_ids": tf.constant([[0, 0, 0, 0]]),
+            "padding_mask": tf.constant([[1, 1, 1, 1]]),
+        }
+        model = AlbertClassifier.from_preset(
+            "albert_base_en_uncased",
+            load_weights=load_weights,
+            preprocessor=None,
+        )
+        # Never assert output values, as the head weights are random.
+        model.predict(input_data)
+
     @parameterized.named_parameters(
         ("preset_weights", True), ("random_weights", False)
     )
@@ -70,6 +99,7 @@ class AlbertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
         ("albert_tokenizer", AlbertTokenizer),
         ("albert_preprocessor", AlbertPreprocessor),
         ("albert", AlbertBackbone),
+        ("albert_classifier", AlbertClassifier),
     )
     def test_preset_docstring(self, cls):
         """Check we did our docstring formatting correctly."""
@@ -80,6 +110,7 @@ class AlbertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
         ("albert_tokenizer", AlbertTokenizer),
         ("albert_preprocessor", AlbertPreprocessor),
         ("albert", AlbertBackbone),
+        ("albert_classifier", AlbertClassifier),
     )
     def test_unknown_preset_error(self, cls):
         # Not a preset name
@@ -91,7 +122,6 @@ class AlbertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
 class AlbertPresetFullTest(tf.test.TestCase, parameterized.TestCase):
     """
     Test the full enumeration of our preset.
-
     This tests every ALBERT preset and is only run manually.
     Run with:
     `pytest keras_nlp/models/albert/albert_presets_test.py --run_extra_large`
@@ -116,6 +146,41 @@ class AlbertPresetFullTest(tf.test.TestCase, parameterized.TestCase):
             }
             model(input_data)
 
+    @parameterized.named_parameters(
+        ("load_weights", True), ("no_load_weights", False)
+    )
+    def test_load_albert_classifier(self, load_weights):
+        for preset in AlbertClassifier.presets:
+            classifier = AlbertClassifier.from_preset(
+                preset,
+                load_weights=load_weights,
+            )
+            input_data = tf.constant(["This quick brown fox"])
+            classifier.predict(input_data)
+
+    @parameterized.named_parameters(
+        ("load_weights", True), ("no_load_weights", False)
+    )
+    def test_load_albert_classifier_without_preprocessing(self, load_weights):
+        for preset in AlbertClassifier.presets:
+            classifier = AlbertClassifier.from_preset(
+                preset,
+                preprocessor=None,
+                load_weights=load_weights,
+            )
+            input_data = {
+                "token_ids": tf.random.uniform(
+                    shape=(1, 512),
+                    dtype=tf.int64,
+                    maxval=classifier.backbone.vocabulary_size,
+                ),
+                "segment_ids": tf.constant(
+                    [0] * 200 + [1] * 312, shape=(1, 512)
+                ),
+                "padding_mask": tf.constant([1] * 512, shape=(1, 512)),
+            }
+            classifier.predict(input_data)
+
     def test_load_tokenizers(self):
         for preset in AlbertTokenizer.presets:
             tokenizer = AlbertTokenizer.from_preset(preset)