Repository: keras-team/keras-nlp

Issue:
Extend default compilation to other classifier models
In #695 we added default compilation for `BertClassifier` and `RobertaClassifier` models. Extend this to the other classifier tasks currently in our library:
* [ ] `DistilBertClassifier`
* [ ] `DebertaV3Classifier`
* [ ] `AlbertClassifier`
* [ ] `FNetClassifier`
* [ ] `XLMRobertaClassifier`

For now, we will use `Adam` with a constant learning rate of `5e-5`:
```python
# Default compilation
self.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.Adam(5e-5),
    metrics=keras.metrics.SparseCategoricalAccuracy(),
    jit_compile=is_xla_compatible(self),
)
```

Discussion:
google-cla[bot]: Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this [failed invocation](https://github.com/keras-team/keras-nlp/pull/714/checks?check_run_id=11058758001) of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Patch:
diff --git a/keras_nlp/models/albert/albert_classifier.py b/keras_nlp/models/albert/albert_classifier.py
index 20f6e1a8..a6b3eb41 100644
--- a/keras_nlp/models/albert/albert_classifier.py
+++ b/keras_nlp/models/albert/albert_classifier.py
@@ -22,6 +22,7 @@ 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.keras_utils import is_xla_compatible
 from keras_nlp.utils.python_utils import classproperty
 
 
@@ -173,6 +174,14 @@ class AlbertClassifier(Task):
         self.num_classes = num_classes
         self.dropout = dropout
 
+        # Default compilation
+        self.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=keras.optimizers.Adam(5e-5),
+            metrics=keras.metrics.SparseCategoricalAccuracy(),
+            jit_compile=is_xla_compatible(self),
+        )
+
     def get_config(self):
         config = super().get_config()
         config.update(
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_classifier.py b/keras_nlp/models/deberta_v3/deberta_v3_classifier.py
index 3ea01dd1..660d0a99 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_classifier.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_classifier.py
@@ -26,6 +26,7 @@ from keras_nlp.models.deberta_v3.deberta_v3_preprocessor import (
 )
 from keras_nlp.models.deberta_v3.deberta_v3_presets import backbone_presets
 from keras_nlp.models.task import Task
+from keras_nlp.utils.keras_utils import is_xla_compatible
 from keras_nlp.utils.python_utils import classproperty
 
 
@@ -199,6 +200,14 @@ class DebertaV3Classifier(Task):
         self.hidden_dim = hidden_dim
         self.dropout = dropout
 
+        # Default compilation
+        self.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=keras.optimizers.Adam(5e-5),
+            metrics=keras.metrics.SparseCategoricalAccuracy(),
+            jit_compile=is_xla_compatible(self),
+        )
+
     def get_config(self):
         config = super().get_config()
         config.update(
diff --git a/keras_nlp/models/distil_bert/distil_bert_classifier.py b/keras_nlp/models/distil_bert/distil_bert_classifier.py
index 6ed81d86..b20c556d 100644
--- a/keras_nlp/models/distil_bert/distil_bert_classifier.py
+++ b/keras_nlp/models/distil_bert/distil_bert_classifier.py
@@ -26,6 +26,7 @@ from keras_nlp.models.distil_bert.distil_bert_preprocessor import (
 )
 from keras_nlp.models.distil_bert.distil_bert_presets import backbone_presets
 from keras_nlp.models.task import Task
+from keras_nlp.utils.keras_utils import is_xla_compatible
 from keras_nlp.utils.python_utils import classproperty
 
 
@@ -202,6 +203,13 @@ class DistilBertClassifier(Task):
         self.hidden_dim = hidden_dim
         self.dropout = dropout
 
+        self.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=keras.optimizers.Adam(5e-5),
+            metrics=keras.metrics.SparseCategoricalAccuracy(),
+            jit_compile=is_xla_compatible(self),
+        )
+
     def get_config(self):
         config = super().get_config()
         config.update(
diff --git a/keras_nlp/models/f_net/f_net_classifier.py b/keras_nlp/models/f_net/f_net_classifier.py
index 8f7f08a0..6e43b76f 100644
--- a/keras_nlp/models/f_net/f_net_classifier.py
+++ b/keras_nlp/models/f_net/f_net_classifier.py
@@ -23,6 +23,7 @@ from keras_nlp.models.f_net.f_net_backbone import f_net_kernel_initializer
 from keras_nlp.models.f_net.f_net_preprocessor import FNetPreprocessor
 from keras_nlp.models.f_net.f_net_presets import backbone_presets
 from keras_nlp.models.task import Task
+from keras_nlp.utils.keras_utils import is_xla_compatible
 from keras_nlp.utils.python_utils import classproperty
 
 
@@ -121,6 +122,13 @@ class FNetClassifier(Task):
         self.num_classes = num_classes
         self.dropout = dropout
 
+        self.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=keras.optimizers.Adam(5e-5),
+            metrics=keras.metrics.SparseCategoricalAccuracy(),
+            jit_compile=is_xla_compatible(self),
+        )
+
     def get_config(self):
         config = super().get_config()
         config.update(
diff --git a/keras_nlp/models/opt/opt_backbone.py b/keras_nlp/models/opt/opt_backbone.py
index c88f208b..cdcb5062 100644
--- a/keras_nlp/models/opt/opt_backbone.py
+++ b/keras_nlp/models/opt/opt_backbone.py
@@ -30,7 +30,7 @@ def opt_kernel_initializer(stddev=0.02):
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
 class OPTBackbone(Backbone):
-    """OPT encoder-decoder network.
+    """OPT decoder network.
 
     This class implements a Transformer-based decoder model as described in
     ["OPT: Open Pre-trained Transformer Language Models"](https://arxiv.org/abs/2205.01068).
@@ -45,15 +45,14 @@ class OPTBackbone(Backbone):
 
     Args:
         vocabulary_size: int. The size of the token vocabulary.
-        num_layers: int. The number of transformer encoder layers and
-            transformer decoder layers.
+        num_layers: int. The number of transformer decoder layers.
         num_heads: int. The number of attention heads for each transformer.
             The hidden size must be divisible by the number of attention heads.
         hidden_dim: int. The hidden size of the transformer decoder layers.
         intermediate_dim: int. The output dimension of the first Dense layer in
             a two-layer feedforward network for each transformer decoder layer.
-        dropout: float. Dropout probability for the Transformer encoder.
-        max_sequence_length: int. The maximum sequence length that this encoder
+        dropout: float. Dropout probability for the Transformer decoder.
+        max_sequence_length: int. The maximum sequence length that this decoder
             can consume. If None, `max_sequence_length` uses the value from
             sequence length. This determines the variable shape for positional
             embeddings.
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
index 0e250ff8..c50ea42b 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
@@ -24,6 +24,7 @@ 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.keras_utils import is_xla_compatible
 from keras_nlp.utils.python_utils import classproperty
 
 
@@ -196,6 +197,13 @@ class XLMRobertaClassifier(Task):
         self.hidden_dim = hidden_dim
         self.dropout = dropout
 
+        self.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=keras.optimizers.Adam(5e-5),
+            metrics=keras.metrics.SparseCategoricalAccuracy(),
+            jit_compile=is_xla_compatible(self),
+        )
+
     def preprocess_samples(self, x, y=None, sample_weight=None):
         return self.preprocessor(x, y=y, sample_weight=sample_weight)
 


Test patch:
diff --git a/keras_nlp/models/albert/albert_classifier_test.py b/keras_nlp/models/albert/albert_classifier_test.py
index 6f586b9c..622e8bae 100644
--- a/keras_nlp/models/albert/albert_classifier_test.py
+++ b/keras_nlp/models/albert/albert_classifier_test.py
@@ -108,6 +108,9 @@ class AlbertClassifierTest(tf.test.TestCase, parameterized.TestCase):
         self.classifier_no_preprocessing.compile(jit_compile=jit_compile)
         self.classifier_no_preprocessing.predict(self.preprocessed_batch)
 
+    def test_albert_classifier_fit_default_compile(self):
+        self.classifier.fit(self.raw_dataset)
+
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py b/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py
index 1b57b8aa..9266a901 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py
@@ -106,6 +106,9 @@ class DebertaV3ClassifierTest(tf.test.TestCase, parameterized.TestCase):
         self.classifier_no_preprocessing.compile(jit_compile=jit_compile)
         self.classifier_no_preprocessing.predict(self.preprocessed_batch)
 
+    def test_debertav3_classifier_fit_default_compile(self):
+        self.classifier.fit(self.raw_dataset)
+
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
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 e0d9f89c..a9a0fc06 100644
--- a/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
@@ -109,6 +109,9 @@ class DistilBertClassifierTest(tf.test.TestCase, parameterized.TestCase):
         )
         self.classifier_no_preprocessing.fit(self.preprocessed_dataset)
 
+    def test_distilbert_classifier_fit_default_compile(self):
+        self.classifier.fit(self.raw_dataset)
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
diff --git a/keras_nlp/models/f_net/f_net_classifier_test.py b/keras_nlp/models/f_net/f_net_classifier_test.py
index cf4120fe..7593b809 100644
--- a/keras_nlp/models/f_net/f_net_classifier_test.py
+++ b/keras_nlp/models/f_net/f_net_classifier_test.py
@@ -107,6 +107,9 @@ class FNetClassifierTest(tf.test.TestCase, parameterized.TestCase):
         self.classifier_no_preprocessing.compile(jit_compile=jit_compile)
         self.classifier_no_preprocessing.predict(self.preprocessed_batch)
 
+    def test_fnet_classifier_fit_default_compile(self):
+        self.classifier.fit(self.raw_dataset)
+
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
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 b1e7ea92..c8efb4fb 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
@@ -122,6 +122,9 @@ class XLMRobertaClassifierTest(tf.test.TestCase, parameterized.TestCase):
         )
         self.classifier_no_preprocessing.fit(self.preprocessed_dataset)
 
+    def test_xlmroberta_classifier_fit_default_compile(self):
+        self.classifier.fit(self.raw_dataset)
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),