Repository: keras-team/keras-nlp

Issue:
Add compilation defaults for the MaskedLM task models
In https://github.com/keras-team/keras-nlp/issues/709 we added compilation defaults to all classification tasks. We should also do the same for all `XXMaskedLM` models. E.g. [bert](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/bert/bert_masked_lm.py). Here's the full list of models.

- [ ] `AlbertMaskedLM`
- [ ] `BertMaskedLM`
- [ ] `DebertV3MaskedLM`
- [ ] `DistilBertMaskedLM`
- [ ] `FNetMaskedLM`
- [ ] `RobertaMaskedLM`

https://github.com/keras-team/keras-nlp/pull/714 is a good template PR. We can probably just use Adam and `1e-5` or `2e-5` as a learning rate for now. Though a little experimentation for each model to make sure the task does converge would be helpful. This [colab](https://gist.github.com/mattdangerw/b16c257973762a0b4ab9a34f6a932cc1) may be a helpful starting place.


Discussion:
I am new here to contribute open source and looking for issues. Can I try this issue, I have programming knowledge on python.
/assign
Hi @mattdangerw , I can take this up !
@mattdangerw How about this ... @Ujjal99kumar seems a beginner, but this issue is not good first issue. We can let him have crack at `BertMaskedLM` while I will raise PRs for rest of models ! I can review his pr as peer review, helping him.
@mattdangerw I will like to contribute in this issue.
@Ujjal99kumar @shivance we can split up models and raise PR.

Resolving patch:
diff --git a/keras_nlp/models/albert/albert_masked_lm.py b/keras_nlp/models/albert/albert_masked_lm.py
--- a/keras_nlp/models/albert/albert_masked_lm.py
+++ b/keras_nlp/models/albert/albert_masked_lm.py
@@ -27,6 +27,7 @@
 )
 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
 
 
@@ -142,6 +143,13 @@ def __init__(self, backbone, preprocessor=None, **kwargs):
         self.backbone = backbone
         self.preprocessor = preprocessor
 
+        self.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=keras.optimizers.Adam(5e-5),
+            weighted_metrics=keras.metrics.SparseCategoricalAccuracy(),
+            jit_compile=is_xla_compatible(self),
+        )
+
     @classproperty
     def backbone_cls(cls):
         return AlbertBackbone
diff --git a/keras_nlp/models/distil_bert/distil_bert_masked_lm.py b/keras_nlp/models/distil_bert/distil_bert_masked_lm.py
--- a/keras_nlp/models/distil_bert/distil_bert_masked_lm.py
+++ b/keras_nlp/models/distil_bert/distil_bert_masked_lm.py
@@ -28,6 +28,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
 
 
@@ -143,6 +144,13 @@ def __init__(
         self.backbone = backbone
         self.preprocessor = preprocessor
 
+        self.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=keras.optimizers.Adam(5e-5),
+            weighted_metrics=keras.metrics.SparseCategoricalAccuracy(),
+            jit_compile=is_xla_compatible(self),
+        )
+
     @classproperty
     def backbone_cls(cls):
         return DistilBertBackbone
diff --git a/keras_nlp/models/roberta/roberta_masked_lm.py b/keras_nlp/models/roberta/roberta_masked_lm.py
--- a/keras_nlp/models/roberta/roberta_masked_lm.py
+++ b/keras_nlp/models/roberta/roberta_masked_lm.py
@@ -26,6 +26,7 @@
 )
 from keras_nlp.models.roberta.roberta_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
 
 
@@ -141,6 +142,13 @@ def __init__(
         self.backbone = backbone
         self.preprocessor = preprocessor
 
+        self.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=keras.optimizers.Adam(5e-5),
+            weighted_metrics=keras.metrics.SparseCategoricalAccuracy(),
+            jit_compile=is_xla_compatible(self),
+        )
+
     @classproperty
     def backbone_cls(cls):
         return RobertaBackbone


Test patch:
diff --git a/keras_nlp/models/albert/albert_masked_lm_test.py b/keras_nlp/models/albert/albert_masked_lm_test.py
--- a/keras_nlp/models/albert/albert_masked_lm_test.py
+++ b/keras_nlp/models/albert/albert_masked_lm_test.py
@@ -114,6 +114,9 @@ def test_albert_masked_lm_predict_no_preprocessing(self, jit_compile):
         self.masked_lm_no_preprocessing.compile(jit_compile=jit_compile)
         self.masked_lm_no_preprocessing.predict(self.preprocessed_batch)
 
+    def test_albert_masked_lm_fit_default_compile(self):
+        self.masked_lm.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_masked_lm_test.py b/keras_nlp/models/distil_bert/distil_bert_masked_lm_test.py
--- a/keras_nlp/models/distil_bert/distil_bert_masked_lm_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_masked_lm_test.py
@@ -89,6 +89,9 @@ def test_distilbert_masked_lm_predict_no_preprocessing(self, jit_compile):
         self.masked_lm_no_preprocessing.compile(jit_compile=jit_compile)
         self.masked_lm_no_preprocessing.predict(self.preprocessed_batch)
 
+    def test_distil_bert_masked_lm_fit_default_compile(self):
+        self.masked_lm.fit(self.raw_dataset)
+
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
diff --git a/keras_nlp/models/roberta/roberta_masked_lm_test.py b/keras_nlp/models/roberta/roberta_masked_lm_test.py
--- a/keras_nlp/models/roberta/roberta_masked_lm_test.py
+++ b/keras_nlp/models/roberta/roberta_masked_lm_test.py
@@ -103,6 +103,9 @@ def test_roberta_masked_lm_predict_no_preprocessing(self, jit_compile):
         self.masked_lm_no_preprocessing.compile(jit_compile=jit_compile)
         self.masked_lm_no_preprocessing.predict(self.preprocessed_batch)
 
+    def test_roberta_masked_lm_fit_default_compile(self):
+        self.masked_lm.fit(self.raw_dataset)
+
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )