Repository: keras-team/keras-nlp

Issue:
Standardize on "backbone" naming for BERT
We've settled on "backbone" as the name for the pretraining submodel reused in task models (even generative models like GPT)! I propose two changes to help reduce confusion for users:
1. `bert_models.py` -> `bert_backbone.py`
2. `Bert` -> `BertBackbone`

This should make it more clear when we refer to a "backbone" model in our user guides.

Note: I dropped the plural in the filename since I think the backbone has to be unique.

Discussion:
mattdangerw: Let's make sure to hit the `examples/` and `tools/` directories too (e.g. this would break the glue script Chen added). There's probably a carefully constructed `sed` that could hit this in one go.
jbischof: Thanks for the reminder @mattdangerw. Fixed!

Patch:
diff --git a/examples/bert/bert_config.py b/examples/bert/bert_config.py
index 6a0b8902..951ee152 100644
--- a/examples/bert/bert_config.py
+++ b/examples/bert/bert_config.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# TODO(jbischof): remove in favor of BertBase, BertSmall, etc
+# TODO(jbischof): remove in favor of presets with load_weights=False
 MODEL_CONFIGS = {
     "tiny": {
         "num_layers": 2,
diff --git a/examples/bert/bert_train.py b/examples/bert/bert_train.py
index cf369039..646d970f 100644
--- a/examples/bert/bert_train.py
+++ b/examples/bert/bert_train.py
@@ -91,7 +91,7 @@ flags.DEFINE_integer(
 
 
 class MaskedLMHead(keras.layers.Layer):
-    """Masked language model network head for Bert.
+    """Masked language model network head for BERT.
 
     This layer implements a masked language model based on the provided
     transformer based encoder. It assumes that the encoder network being passed
@@ -99,7 +99,14 @@ class MaskedLMHead(keras.layers.Layer):
 
     Example:
     ```python
-    encoder = keras_nlp.models.BertBase()
+    encoder = keras_nlp.models.BertBackbone(
+        vocabulary_size=30552,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=12,
+    )
     lm_layer = MaskedLMHead(embedding_table=encoder.get_embedding_table())
     ```
 
@@ -380,7 +387,7 @@ def main(_):
 
     with strategy.scope():
         # Create a Bert model the input config.
-        encoder = keras_nlp.models.Bert(
+        encoder = keras_nlp.models.BertBackbone(
             vocabulary_size=len(vocab), **model_config
         )
         # Make sure model has been called.
diff --git a/examples/glue_benchmark/glue.py b/examples/glue_benchmark/glue.py
index 7251bfc4..a0f38bab 100644
--- a/examples/glue_benchmark/glue.py
+++ b/examples/glue_benchmark/glue.py
@@ -242,7 +242,9 @@ def main(_):
         # Commonly the classifier is simply your model + several dense layers,
         # please refer to "Make the Finetuning Model" section in README for
         # detailed instructions.
-        bert_model = keras_nlp.models.Bert.from_preset("bert_tiny_uncased_en")
+        bert_model = keras_nlp.models.BertBackbone.from_preset(
+            "bert_tiny_uncased_en"
+        )
         finetuning_model = keras_nlp.models.BertClassifier(
             backbone=bert_model,
             num_classes=num_classes,
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index 97e33411..90762083 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from keras_nlp.models.bert.bert_models import Bert
+from keras_nlp.models.bert.bert_backbone import BertBackbone
 from keras_nlp.models.bert.bert_preprocessing import BertPreprocessor
 from keras_nlp.models.bert.bert_preprocessing import BertTokenizer
 from keras_nlp.models.bert.bert_tasks import BertClassifier
diff --git a/keras_nlp/models/bert/bert_models.py b/keras_nlp/models/bert/bert_backbone.py
similarity index 97%
rename from keras_nlp/models/bert/bert_models.py
rename to keras_nlp/models/bert/bert_backbone.py
index 6bdfe0f9..de362652 100644
--- a/keras_nlp/models/bert/bert_models.py
+++ b/keras_nlp/models/bert/bert_backbone.py
@@ -32,7 +32,7 @@ def bert_kernel_initializer(stddev=0.02):
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class Bert(keras.Model):
+class BertBackbone(keras.Model):
     """BERT encoder network.
 
     This class implements a bi-directional Transformer-based encoder as
@@ -77,7 +77,7 @@ class Bert(keras.Model):
     }
 
     # Randomly initialized BERT encoder
-    model = keras_nlp.models.Bert(
+    model = keras_nlp.models.BertBackbone(
         vocabulary_size=30552,
         num_layers=12,
         num_heads=12,
@@ -252,11 +252,14 @@ class Bert(keras.Model):
         }
 
         # Load architecture and weights from preset
-        model = Bert.from_preset("bert_base_uncased_en")
+        model = BertBackbone.from_preset("bert_base_uncased_en")
         output = model(input_data)
 
         # Load randomly initalized model from preset architecture
-        model = Bert.from_preset("bert_base_uncased_en", load_weights=False)
+        model = BertBackbone.from_preset(
+            "bert_base_uncased_en",
+            load_weights=False
+        )
         output = model(input_data)
         ```
         """
diff --git a/keras_nlp/models/bert/bert_presets.py b/keras_nlp/models/bert/bert_presets.py
index 9a0487df..a0d36188 100644
--- a/keras_nlp/models/bert/bert_presets.py
+++ b/keras_nlp/models/bert/bert_presets.py
@@ -225,7 +225,7 @@ classifier_presets = {
     "bert_tiny_uncased_en_sst2": {
         "config": {
             "backbone": {
-                "class_name": "keras_nlp>Bert",
+                "class_name": "keras_nlp>BertBackbone",
                 "config": {
                     "vocabulary_size": 30522,
                     "hidden_dim": 128,
diff --git a/keras_nlp/models/bert/bert_tasks.py b/keras_nlp/models/bert/bert_tasks.py
index cb6baa00..3110a2c9 100644
--- a/keras_nlp/models/bert/bert_tasks.py
+++ b/keras_nlp/models/bert/bert_tasks.py
@@ -18,8 +18,8 @@ import os
 
 from tensorflow import keras
 
-from keras_nlp.models.bert.bert_models import Bert
-from keras_nlp.models.bert.bert_models import bert_kernel_initializer
+from keras_nlp.models.bert.bert_backbone import BertBackbone
+from keras_nlp.models.bert.bert_backbone import bert_kernel_initializer
 from keras_nlp.models.bert.bert_preprocessing import BertPreprocessor
 from keras_nlp.models.bert.bert_presets import backbone_presets
 from keras_nlp.models.bert.bert_presets import classifier_presets
@@ -36,7 +36,7 @@ CLASSIFIER_PRESET_NAMES = ", ".join(
 class BertClassifier(PipelineModel):
     """An end-to-end BERT model for classification tasks
 
-    This model attaches a classification head to a `keras_nlp.model.Bert`
+    This model attaches a classification head to a `keras_nlp.model.BertBackbone`
     backbone, mapping from the backbone outputs to logit output suitable for
     a classification task. For usage of this model with pre-trained weights, see
     the `from_preset()` method.
@@ -50,7 +50,7 @@ class BertClassifier(PipelineModel):
     warranties or conditions of any kind.
 
     Args:
-        backbone: A `keras_nlp.models.Bert` instance.
+        backbone: A `keras_nlp.models.BertBackbone` instance.
         num_classes: int. Number of classes to predict.
         preprocessor: A `keras_nlp.models.BertPreprocessor` or `None`. If
             `None`, this model will not apply preprocessing, and inputs should
@@ -72,7 +72,7 @@ class BertClassifier(PipelineModel):
     labels = [0, 3]
 
     # Randomly initialize a BERT backbone.
-    backbone = keras_nlp.models.Bert(
+    backbone = keras_nlp.models.BertBackbone(
         vocabulary_size=30552,
         num_layers=12,
         num_heads=12,
@@ -128,7 +128,9 @@ class BertClassifier(PipelineModel):
 
     @property
     def backbone(self):
-        """A `keras_nlp.models.Bert` instance providing the encoder submodel."""
+        """A `keras_nlp.models.BertBackbone` instance providing the encoder
+        submodel.
+        """
         return self._backbone
 
     @property
@@ -260,8 +262,8 @@ class BertClassifier(PipelineModel):
             kwargs["preprocessor"] = BertPreprocessor.from_preset(preset)
 
         # Check if preset is backbone-only model
-        if preset in Bert.presets:
-            backbone = Bert.from_preset(preset, load_weights)
+        if preset in BertBackbone.presets:
+            backbone = BertBackbone.from_preset(preset, load_weights)
             return cls(backbone, **kwargs)
 
         # Otherwise must be one of class presets


Test patch:
diff --git a/keras_nlp/models/bert/bert_models_test.py b/keras_nlp/models/bert/bert_backbone_test.py
similarity index 93%
rename from keras_nlp/models/bert/bert_models_test.py
rename to keras_nlp/models/bert/bert_backbone_test.py
index 8a39cb17..1b8680f4 100644
--- a/keras_nlp/models/bert/bert_models_test.py
+++ b/keras_nlp/models/bert/bert_backbone_test.py
@@ -19,12 +19,12 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.bert.bert_models import Bert
+from keras_nlp.models.bert.bert_backbone import BertBackbone
 
 
-class BertTest(tf.test.TestCase, parameterized.TestCase):
+class BertBackboneTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = Bert(
+        self.model = BertBackbone(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
@@ -93,7 +93,7 @@ class BertTest(tf.test.TestCase, parameterized.TestCase):
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
-        self.assertIsInstance(restored_model, Bert)
+        self.assertIsInstance(restored_model, BertBackbone)
 
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
diff --git a/keras_nlp/models/bert/bert_presets_test.py b/keras_nlp/models/bert/bert_presets_test.py
index 1d368800..f66880aa 100644
--- a/keras_nlp/models/bert/bert_presets_test.py
+++ b/keras_nlp/models/bert/bert_presets_test.py
@@ -17,7 +17,7 @@ import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 
-from keras_nlp.models.bert.bert_models import Bert
+from keras_nlp.models.bert.bert_backbone import BertBackbone
 from keras_nlp.models.bert.bert_preprocessing import BertPreprocessor
 from keras_nlp.models.bert.bert_preprocessing import BertTokenizer
 from keras_nlp.models.bert.bert_tasks import BertClassifier
@@ -58,7 +58,7 @@ class BertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
             "segment_ids": tf.constant([[0, 0, 0, 0]]),
             "padding_mask": tf.constant([[1, 1, 1, 1]]),
         }
-        model = Bert.from_preset(
+        model = BertBackbone.from_preset(
             "bert_tiny_uncased_en", load_weights=load_weights
         )
         outputs = model(input_data)["sequence_output"]
@@ -105,7 +105,7 @@ class BertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
     @parameterized.named_parameters(
         ("bert_tokenizer", BertTokenizer),
         ("bert_preprocessor", BertPreprocessor),
-        ("bert", Bert),
+        ("bert", BertBackbone),
         ("bert_classifier", BertClassifier),
     )
     def test_preset_mutability(self, cls):
@@ -120,15 +120,17 @@ class BertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
         self.assertEqual(config["num_layers"], 1)
         self.assertEqual(obj.presets[preset]["config"]["num_layers"], 2)
         # Cannot mutate presets in the class
-        config = Bert.presets[preset]["config"]
+        config = BertBackbone.presets[preset]["config"]
         config["num_layers"] = 1
         self.assertEqual(config["num_layers"], 1)
-        self.assertEqual(Bert.presets[preset]["config"]["num_layers"], 2)
+        self.assertEqual(
+            BertBackbone.presets[preset]["config"]["num_layers"], 2
+        )
 
     @parameterized.named_parameters(
         ("bert_tokenizer", BertTokenizer),
         ("bert_preprocessor", BertPreprocessor),
-        ("bert", Bert),
+        ("bert", BertBackbone),
         ("bert_classifier", BertClassifier),
     )
     def test_preset_docstring(self, cls):
@@ -139,7 +141,7 @@ class BertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
     @parameterized.named_parameters(
         ("bert_tokenizer", BertTokenizer),
         ("bert_preprocessor", BertPreprocessor),
-        ("bert", Bert),
+        ("bert", BertBackbone),
         ("bert_classifier", BertClassifier),
     )
     def test_unknown_preset_error(self, cls):
@@ -179,8 +181,8 @@ class BertPresetFullTest(tf.test.TestCase, parameterized.TestCase):
         ("load_weights", True), ("no_load_weights", False)
     )
     def test_load_bert(self, load_weights):
-        for preset in Bert.presets:
-            model = Bert.from_preset(preset, load_weights=load_weights)
+        for preset in BertBackbone.presets:
+            model = BertBackbone.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/bert/bert_tasks_test.py b/keras_nlp/models/bert/bert_tasks_test.py
index 09619fcd..377c53e3 100644
--- a/keras_nlp/models/bert/bert_tasks_test.py
+++ b/keras_nlp/models/bert/bert_tasks_test.py
@@ -19,7 +19,7 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.bert.bert_models import Bert
+from keras_nlp.models.bert.bert_backbone import BertBackbone
 from keras_nlp.models.bert.bert_preprocessing import BertPreprocessor
 from keras_nlp.models.bert.bert_preprocessing import BertTokenizer
 from keras_nlp.models.bert.bert_tasks import BertClassifier
@@ -27,7 +27,7 @@ from keras_nlp.models.bert.bert_tasks import BertClassifier
 
 class BertClassifierTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.backbone = Bert(
+        self.backbone = BertBackbone(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,