Address comments

keras_nlp/models/roberta/roberta_masked_lm.py

@@ -24,9 +24,8 @@
     RobertaMaskedLMPreprocessor,
 )
 from keras_nlp.models.roberta.roberta_presets import backbone_presets
-from keras_nlp.models import Task
+from keras_nlp.models.task import Task
 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")
@@ -56,19 +55,37 @@ class RobertaMaskedLM(Task):
             inputs should be preprocessed before calling the model.
 
     Example usage:
+
+    Raw string inputs and pretrained backbone.
     ```python
-    # Create a dataset where 0 is the mask token.
+    # Create a dataset with raw string features. Labels are inferred.
+    features = ["The quick brown fox jumped.", "I forgot my homework."]
+
+    # Create a RobertaMaskedLM with a pretrained backbone and further train
+    # on an MLM task.
+    classifier = keras_nlp.models.RobertaMaskedLM.from_preset(
+        "roberta_base_en",
+    )
+    classifier.compile(
+        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+    )
+    classifier.fit(x=features, batch_size=2)
+    ```
+
+    Preprocessed inputs and custom backbone.
+    ```python
+    # Create a preprocessed dataset where 0 is the mask token.
     preprocessed_features = {
         "token_ids": tf.constant(
-            [[1, 2, 0, 4, 0, 6, 7, 8, 9]] * 2, shape=(2, 12)
+            [[1, 2, 0, 4, 0, 6, 7, 8]] * 2, shape=(2, 8)
         ),
         "padding_mask": tf.constant(
-            [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]] * 2, shape=(2, 12)
+            [[1, 1, 1, 1, 1, 1, 1, 1]] * 2, shape=(2, 8)
         ),
-        "mask_positions": tf.constant([[1, 2]] * 2, shape=(2, 12))
+        "mask_positions": tf.constant([[2, 4]] * 2, shape=(2, 2))
     }
     # Labels are the original masked values.
-    labels = [3, 5]
+    labels = [[3, 5]] * 2
 
     # Randomly initialize a RoBERTa encoder
     backbone = keras_nlp.models.RobertaBackbone(
@@ -122,3 +139,15 @@ def __init__(
         # All references to `self` below this line
         self._backbone = backbone
         self._preprocessor = preprocessor
+
+    @classproperty
+    def backbone_cls(cls):
+        return RobertaBackbone
+
+    @classproperty
+    def preprocessor_cls(cls):
+        return RobertaMaskedLMPreprocessor
+
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy(backbone_presets)

keras_nlp/models/roberta/roberta_masked_lm_test.py

@@ -128,16 +128,17 @@ def test_roberta_masked_lm_fit_no_preprocessing(self, jit_compile):
         ("keras_format", "keras_v3", "model.keras"),
     )
     def test_saving_model(self, save_format, filename):
-        keras.utils.set_random_seed(42)
-        model_output = self.masked_lm.predict(self.raw_batch)
         save_path = os.path.join(self.get_temp_dir(), filename)
         self.masked_lm.save(save_path, save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
         self.assertIsInstance(restored_model, RobertaMaskedLM)
 
+        keras.utils.set_random_seed(42)
+        model_output = self.masked_lm.predict(self.raw_batch)
         # Check that output matches.
         keras.utils.set_random_seed(42)
         restored_output = restored_model.predict(self.raw_batch)
+
         self.assertAllClose(model_output, restored_output)

keras_nlp/models/roberta/roberta_preprocessor.py

@@ -79,7 +79,9 @@ class RobertaPreprocessor(Preprocessor):
     Examples:
     ```python
     # Load the preprocessor from a preset.
-    preprocessor = keras_nlp.models.RobertaPreprocessor.from_preset("roberta_base_en")
+    preprocessor = keras_nlp.models.RobertaPreprocessor.from_preset(
+        "roberta_base_en",
+    )
 
     # Tokenize and pack a single sentence.
     sentence = tf.constant(" afternoon sun")
@@ -140,7 +142,7 @@ class RobertaPreprocessor(Preprocessor):
         vocabulary=vocab,
         merges=merges,
     )
-    preprocessor = keras_nlp.models.RobertaMaskedLMPreprocessor(
+    preprocessor = keras_nlp.models.RobertaPreprocessor(
         tokenizer=tokenizer,
         sequence_length=20,
     )