Default compilation other classifiers

keras_nlp/models/albert/albert_classifier.py

@@ -22,6 +22,7 @@
 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 @@ def __init__(
         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(

keras_nlp/models/albert/albert_classifier_test.py

@@ -108,6 +108,9 @@ 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)
 
+    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)
     )

keras_nlp/models/deberta_v3/deberta_v3_classifier.py

@@ -26,6 +26,7 @@
 )
 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 @@ def __init__(
         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(

keras_nlp/models/deberta_v3/deberta_v3_classifier_test.py

@@ -106,6 +106,9 @@ 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)
 
+    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)
     )

keras_nlp/models/distil_bert/distil_bert_classifier.py

@@ -26,6 +26,7 @@
 )
 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 @@ def __init__(
         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(

keras_nlp/models/distil_bert/distil_bert_classifier_test.py

@@ -109,6 +109,9 @@ def test_classifier_fit_no_preprocessing(self, jit_compile):
         )
         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"),

keras_nlp/models/f_net/f_net_classifier.py

@@ -23,6 +23,7 @@
 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 @@ def __init__(
         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(

keras_nlp/models/f_net/f_net_classifier_test.py

@@ -107,6 +107,9 @@ def test_fnet_classifier_predict_no_preprocessing(self, jit_compile):
         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)
     )

keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py

@@ -24,6 +24,7 @@
     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 @@ def __init__(
         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)
 

keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py

@@ -122,6 +122,9 @@ def test_classifier_fit_no_preprocessing(self, jit_compile):
         )
         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"),