Repository: keras-team/keras-nlp

Issue:
Replicate #536 changing `Roberta` and `XLMRoberta` -> `RobertaBackbone`
1. `roberta_models.py` -> `roberta_backbone.py`
2. `roberta_models_test.py` -> `roberta_backbone_test.py`
3. `Roberta` -> `RobertaBackbone`

Same idea for `XLMRoberta`.

See #526 for more background.

Discussion:

Patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index 7d2cdbd1..5ee315b3 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -28,11 +28,11 @@ from keras_nlp.models.distil_bert.distil_bert_tokenizer import (
 )
 from keras_nlp.models.gpt2.gpt2_backbone import GPT2
 from keras_nlp.models.gpt2.gpt2_tokenizer import GPT2Tokenizer
-from keras_nlp.models.roberta.roberta_backbone import Roberta
+from keras_nlp.models.roberta.roberta_backbone import RobertaBackbone
 from keras_nlp.models.roberta.roberta_classifier import RobertaClassifier
 from keras_nlp.models.roberta.roberta_preprocessor import RobertaPreprocessor
 from keras_nlp.models.roberta.roberta_tokenizer import RobertaTokenizer
-from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRoberta
+from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRobertaBackbone
 from keras_nlp.models.xlm_roberta.xlm_roberta_preprocessor import (
     XLMRobertaPreprocessor,
 )
diff --git a/keras_nlp/models/roberta/roberta_backbone.py b/keras_nlp/models/roberta/roberta_backbone.py
index 3bdb469e..6094391a 100644
--- a/keras_nlp/models/roberta/roberta_backbone.py
+++ b/keras_nlp/models/roberta/roberta_backbone.py
@@ -32,7 +32,7 @@ def roberta_kernel_initializer(stddev=0.02):
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class Roberta(keras.Model):
+class RobertaBackbone(keras.Model):
     """RoBERTa encoder.
 
     This network implements a bi-directional Transformer-based encoder as
@@ -72,7 +72,7 @@ class Roberta(keras.Model):
     }
 
     # Randomly initialized RoBERTa model
-    model = keras_nlp.models.Roberta(
+    model = keras_nlp.models.RobertaBackbone(
         vocabulary_size=50265,
         num_layers=12,
         num_heads=12,
@@ -208,11 +208,11 @@ class Roberta(keras.Model):
         }
 
         # Load architecture and weights from preset
-        model = keras_nlp.models.Roberta.from_preset("roberta_base")
+        model = keras_nlp.models.RobertaBackbone.from_preset("roberta_base")
         output = model(input_data)
 
         # Load randomly initalized model from preset architecture
-        model = keras_nlp.models.Roberta.from_preset(
+        model = keras_nlp.models.RobertaBackbone.from_preset(
             "roberta_base", load_weights=False
         )
         output = model(input_data)
diff --git a/keras_nlp/models/roberta/roberta_classifier.py b/keras_nlp/models/roberta/roberta_classifier.py
index bf4aed56..2526d51c 100644
--- a/keras_nlp/models/roberta/roberta_classifier.py
+++ b/keras_nlp/models/roberta/roberta_classifier.py
@@ -17,7 +17,7 @@ import copy
 
 from tensorflow import keras
 
-from keras_nlp.models.roberta.roberta_backbone import Roberta
+from keras_nlp.models.roberta.roberta_backbone import RobertaBackbone
 from keras_nlp.models.roberta.roberta_backbone import roberta_kernel_initializer
 from keras_nlp.models.roberta.roberta_presets import backbone_presets
 from keras_nlp.utils.python_utils import classproperty
@@ -47,7 +47,7 @@ class RobertaClassifier(keras.Model):
     }
 
     # Randomly initialized RoBERTa encoder
-    model = keras_nlp.models.Roberta(
+    model = keras_nlp.models.RobertaBackbone(
         vocabulary_size=50265,
         num_layers=12,
         num_heads=12,
@@ -165,8 +165,8 @@ class RobertaClassifier(keras.Model):
         ```
         """
         # Check if preset is backbone-only model
-        if preset in Roberta.presets:
-            backbone = Roberta.from_preset(preset, load_weights)
+        if preset in RobertaBackbone.presets:
+            backbone = RobertaBackbone.from_preset(preset, load_weights)
             return cls(backbone, **kwargs)
 
         # Otherwise must be one of class presets
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py
index 9a6ff5b5..6c178037 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py
@@ -26,7 +26,7 @@ from keras_nlp.utils.python_utils import format_docstring
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class XLMRoberta(roberta_backbone.Roberta):
+class XLMRobertaBackbone(roberta_backbone.RobertaBackbone):
     """XLM-RoBERTa encoder.
 
     This network implements a bi-directional Transformer-based encoder as
@@ -67,7 +67,7 @@ class XLMRoberta(roberta_backbone.Roberta):
     }
 
     # Randomly initialized XLM-R model
-    model = keras_nlp.models.XLMRoberta(
+    model = keras_nlp.models.XLMRobertaBackbone(
         vocabulary_size=250002,
         num_layers=12,
         num_heads=12,
@@ -110,11 +110,11 @@ class XLMRoberta(roberta_backbone.Roberta):
         }
 
         # Load architecture and weights from preset
-        model = keras_nlp.models.XLMRoberta.from_preset("xlm_roberta_base")
+        model = keras_nlp.models.XLMRobertaBackbone.from_preset("xlm_roberta_base")
         output = model(input_data)
 
         # Load randomly initalized model from preset architecture
-        model = keras_nlp.models.XLMRoberta.from_preset(
+        model = keras_nlp.models.XLMRobertaBackbone.from_preset(
             "xlm_roberta_base", load_weights=False
         )
         output = model(input_data)
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
index 0f71c48e..6ce05e2a 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
@@ -18,7 +18,7 @@ import copy
 from tensorflow import keras
 
 from keras_nlp.models.roberta.roberta_backbone import roberta_kernel_initializer
-from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRoberta
+from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRobertaBackbone
 from keras_nlp.models.xlm_roberta.xlm_roberta_presets import backbone_presets
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
@@ -48,7 +48,7 @@ class XLMRobertaClassifier(keras.Model):
     }
 
     # Randomly initialized XLM-RoBERTa encoder
-    model = keras_nlp.models.XLMRoberta(
+    model = keras_nlp.models.XLMRobertaBackbone(
         vocabulary_size=250002,
         num_layers=12,
         num_heads=12,
@@ -171,8 +171,8 @@ class XLMRobertaClassifier(keras.Model):
         ```
         """
         # Check if preset is backbone-only model
-        if preset in XLMRoberta.presets:
-            backbone = XLMRoberta.from_preset(preset, load_weights)
+        if preset in XLMRobertaBackbone.presets:
+            backbone = XLMRobertaBackbone.from_preset(preset, load_weights)
             return cls(backbone, **kwargs)
 
         # Otherwise must be one of class presets
diff --git a/tools/checkpoint_conversion/convert_roberta_checkpoints.py b/tools/checkpoint_conversion/convert_roberta_checkpoints.py
index bc4d1d9d..90342f88 100644
--- a/tools/checkpoint_conversion/convert_roberta_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_roberta_checkpoints.py
@@ -93,7 +93,7 @@ def convert_checkpoints(preset, size):
     }
     print("Config:", cfg)
 
-    keras_nlp_model = keras_nlp.models.Roberta.from_preset(
+    keras_nlp_model = keras_nlp.models.RobertaBackbone.from_preset(
         preset, load_weights=False
     )
 
diff --git a/tools/checkpoint_conversion/convert_xlm_roberta_checkpoints.py b/tools/checkpoint_conversion/convert_xlm_roberta_checkpoints.py
index 1f95e482..eed375fc 100644
--- a/tools/checkpoint_conversion/convert_xlm_roberta_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_xlm_roberta_checkpoints.py
@@ -78,7 +78,7 @@ def convert_checkpoints(preset, size):
     }
     print("Config:", cfg)
 
-    keras_nlp_model = keras_nlp.models.XLMRoberta.from_preset(
+    keras_nlp_model = keras_nlp.models.XLMRobertaBackbone.from_preset(
         preset, load_weights=False
     )
 


Test patch:
diff --git a/keras_nlp/models/roberta/roberta_backbone_test.py b/keras_nlp/models/roberta/roberta_backbone_test.py
index df464476..da3984f2 100644
--- a/keras_nlp/models/roberta/roberta_backbone_test.py
+++ b/keras_nlp/models/roberta/roberta_backbone_test.py
@@ -20,12 +20,12 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.roberta.roberta_backbone import Roberta
+from keras_nlp.models.roberta.roberta_backbone import RobertaBackbone
 
 
 class RobertaTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = Roberta(
+        self.model = RobertaBackbone(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
@@ -93,7 +93,7 @@ class RobertaTest(tf.test.TestCase, parameterized.TestCase):
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
-        self.assertIsInstance(restored_model, Roberta)
+        self.assertIsInstance(restored_model, RobertaBackbone)
 
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
diff --git a/keras_nlp/models/roberta/roberta_classifier_test.py b/keras_nlp/models/roberta/roberta_classifier_test.py
index 0fc6f216..2d81e430 100644
--- a/keras_nlp/models/roberta/roberta_classifier_test.py
+++ b/keras_nlp/models/roberta/roberta_classifier_test.py
@@ -19,13 +19,13 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.roberta.roberta_backbone import Roberta
+from keras_nlp.models.roberta.roberta_backbone import RobertaBackbone
 from keras_nlp.models.roberta.roberta_classifier import RobertaClassifier
 
 
 class RobertaClassifierTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = Roberta(
+        self.model = RobertaBackbone(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
diff --git a/keras_nlp/models/roberta/roberta_presets_test.py b/keras_nlp/models/roberta/roberta_presets_test.py
index 1841fa6b..75ec4c20 100644
--- a/keras_nlp/models/roberta/roberta_presets_test.py
+++ b/keras_nlp/models/roberta/roberta_presets_test.py
@@ -17,7 +17,7 @@ import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 
-from keras_nlp.models.roberta.roberta_backbone import Roberta
+from keras_nlp.models.roberta.roberta_backbone import RobertaBackbone
 from keras_nlp.models.roberta.roberta_classifier import RobertaClassifier
 from keras_nlp.models.roberta.roberta_preprocessor import RobertaPreprocessor
 from keras_nlp.models.roberta.roberta_tokenizer import RobertaTokenizer
@@ -57,7 +57,9 @@ class RobertaPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
             "token_ids": tf.constant([[0, 133, 2119, 2]]),
             "padding_mask": tf.constant([[1, 1, 1, 1]]),
         }
-        model = Roberta.from_preset("roberta_base", load_weights=load_weights)
+        model = RobertaBackbone.from_preset(
+            "roberta_base", load_weights=load_weights
+        )
         outputs = model(input_data)
         if load_weights:
             outputs = outputs[0, 0, :5]
@@ -81,7 +83,7 @@ class RobertaPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
     @parameterized.named_parameters(
         ("roberta_tokenizer", RobertaTokenizer),
         ("roberta_preprocessor", RobertaPreprocessor),
-        ("roberta", Roberta),
+        ("roberta", RobertaBackbone),
         ("roberta_classifier", RobertaClassifier),
     )
     def test_preset_docstring(self, cls):
@@ -92,7 +94,7 @@ class RobertaPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
     @parameterized.named_parameters(
         ("roberta_tokenizer", RobertaTokenizer),
         ("roberta_preprocessor", RobertaPreprocessor),
-        ("roberta", Roberta),
+        ("roberta", RobertaBackbone),
         ("roberta_classifier", RobertaClassifier),
     )
     def test_unknown_preset_error(self, cls):
@@ -115,8 +117,10 @@ class RobertaPresetFullTest(tf.test.TestCase, parameterized.TestCase):
         ("preset_weights", True), ("random_weights", False)
     )
     def test_load_roberta(self, load_weights):
-        for preset in Roberta.presets:
-            model = Roberta.from_preset(preset, load_weights=load_weights)
+        for preset in RobertaBackbone.presets:
+            model = RobertaBackbone.from_preset(
+                preset, load_weights=load_weights
+            )
             input_data = {
                 "token_ids": tf.random.uniform(
                     shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py
index 64a121f3..7c2b8efd 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py
@@ -20,12 +20,12 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRoberta
+from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRobertaBackbone
 
 
 class XLMRobertaTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = XLMRoberta(
+        self.model = XLMRobertaBackbone(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
@@ -93,7 +93,7 @@ class XLMRobertaTest(tf.test.TestCase, parameterized.TestCase):
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
-        self.assertIsInstance(restored_model, XLMRoberta)
+        self.assertIsInstance(restored_model, XLMRobertaBackbone)
 
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
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 22c07e17..29cc2f0d 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
@@ -19,7 +19,7 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRoberta
+from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRobertaBackbone
 from keras_nlp.models.xlm_roberta.xlm_roberta_classifier import (
     XLMRobertaClassifier,
 )
@@ -27,7 +27,7 @@ from keras_nlp.models.xlm_roberta.xlm_roberta_classifier import (
 
 class XLMRobertaClassifierTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = XLMRoberta(
+        self.model = XLMRobertaBackbone(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_presets_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_presets_test.py
index 46bcb0f3..fa578673 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_presets_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_presets_test.py
@@ -17,7 +17,7 @@ import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 
-from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRoberta
+from keras_nlp.models.xlm_roberta.xlm_roberta_backbone import XLMRobertaBackbone
 from keras_nlp.models.xlm_roberta.xlm_roberta_classifier import (
     XLMRobertaClassifier,
 )
@@ -63,7 +63,7 @@ class XLMRobertaPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
             "token_ids": tf.constant([[0, 581, 63773, 2]]),
             "padding_mask": tf.constant([[1, 1, 1, 1]]),
         }
-        model = XLMRoberta.from_preset(
+        model = XLMRobertaBackbone.from_preset(
             "xlm_roberta_base", load_weights=load_weights
         )
         outputs = model(input_data)
@@ -89,7 +89,7 @@ class XLMRobertaPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
     @parameterized.named_parameters(
         ("xlm_roberta_tokenizer", XLMRobertaTokenizer),
         ("xlm_roberta_preprocessor", XLMRobertaPreprocessor),
-        ("xlm_roberta", XLMRoberta),
+        ("xlm_roberta", XLMRobertaBackbone),
         ("xlm_roberta_classifier", XLMRobertaClassifier),
     )
     def test_preset_docstring(self, cls):
@@ -100,7 +100,7 @@ class XLMRobertaPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
     @parameterized.named_parameters(
         ("xlm_roberta_tokenizer", XLMRobertaTokenizer),
         ("xlm_roberta_preprocessor", XLMRobertaPreprocessor),
-        ("xlm_roberta", XLMRoberta),
+        ("xlm_roberta", XLMRobertaBackbone),
         ("xlm_roberta_classifier", XLMRobertaClassifier),
     )
     def test_unknown_preset_error(self, cls):
@@ -123,8 +123,10 @@ class XLMRobertaPresetFullTest(tf.test.TestCase, parameterized.TestCase):
         ("preset_weights", True), ("random_weights", False)
     )
     def test_load_xlm_roberta(self, load_weights):
-        for preset in XLMRoberta.presets:
-            model = XLMRoberta.from_preset(preset, load_weights=load_weights)
+        for preset in XLMRobertaBackbone.presets:
+            model = XLMRobertaBackbone.from_preset(
+                preset, load_weights=load_weights
+            )
             input_data = {
                 "token_ids": tf.random.uniform(
                     shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size