Repository: keras-team/keras-nlp

Issue:
Replicate #536 changing `DistilBert` -> `DistilBertBackbone`
1. `distilbert_models.py` -> `distilbert_backbone.py`
2. `distilbert_models_test.py` -> `distilbert_backbone_test.py`
3. `DistilBert` -> `DistilBertBackbone`

See #526 for more background.

Discussion:

Patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index 1dac8816..7d2cdbd1 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -16,7 +16,7 @@ from keras_nlp.models.bert.bert_backbone import BertBackbone
 from keras_nlp.models.bert.bert_classifier import BertClassifier
 from keras_nlp.models.bert.bert_preprocessor import BertPreprocessor
 from keras_nlp.models.bert.bert_tokenizer import BertTokenizer
-from keras_nlp.models.distil_bert.distil_bert_backbone import DistilBert
+from keras_nlp.models.distil_bert.distil_bert_backbone import DistilBertBackbone
 from keras_nlp.models.distil_bert.distil_bert_classifier import (
     DistilBertClassifier,
 )
diff --git a/keras_nlp/models/distil_bert/distil_bert_backbone.py b/keras_nlp/models/distil_bert/distil_bert_backbone.py
index 3e399044..b75ae1db 100644
--- a/keras_nlp/models/distil_bert/distil_bert_backbone.py
+++ b/keras_nlp/models/distil_bert/distil_bert_backbone.py
@@ -34,7 +34,7 @@ def distilbert_kernel_initializer(stddev=0.02):
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class DistilBert(keras.Model):
+class DistilBertBackbone(keras.Model):
     """DistilBERT encoder network.
 
     This network implements a bi-directional Transformer-based encoder as
@@ -77,7 +77,7 @@ class DistilBert(keras.Model):
     }
 
     # Randomly initialized DistilBERT encoder
-    model = keras_nlp.models.DistilBert(
+    model = keras_nlp.models.DistilBertBackbone(
         vocabulary_size=30552,
         num_layers=6,
         num_heads=12,
@@ -193,7 +193,7 @@ class DistilBert(keras.Model):
         load_weights=True,
         **kwargs,
     ):
-        """Instantiate DistilBERT model from preset architecture and weights.
+        """Instantiate a DistilBERT model from preset architecture and weights.
 
         Args:
             preset: string. Must be one of {{names}}.
@@ -210,13 +210,13 @@ class DistilBert(keras.Model):
         }
 
         # Load architecture and weights from preset
-        model = keras_nlp.models.DistilBert.from_preset(
+        model = keras_nlp.models.DistilBertBackbone.from_preset(
             "distil_bert_base_en_uncased"
         )
         output = model(input_data)
 
         # Load randomly initalized model from preset architecture
-        model = keras_nlp.models.DistilBert.from_preset(
+        model = keras_nlp.models.DistilBertBackbone.from_preset(
             "distil_bert_base_en_uncased", load_weights=False
         )
         output = model(input_data)
diff --git a/keras_nlp/models/distil_bert/distil_bert_classifier.py b/keras_nlp/models/distil_bert/distil_bert_classifier.py
index cba353de..d6292c29 100644
--- a/keras_nlp/models/distil_bert/distil_bert_classifier.py
+++ b/keras_nlp/models/distil_bert/distil_bert_classifier.py
@@ -17,7 +17,7 @@ import copy
 
 from tensorflow import keras
 
-from keras_nlp.models.distil_bert.distil_bert_backbone import DistilBert
+from keras_nlp.models.distil_bert.distil_bert_backbone import DistilBertBackbone
 from keras_nlp.models.distil_bert.distil_bert_backbone import (
     distilbert_kernel_initializer,
 )
@@ -51,7 +51,7 @@ class DistilBertClassifier(keras.Model):
     }
 
     # Randomly initialized DistilBERT encoder
-    model = keras_nlp.models.DistilBert(
+    model = keras_nlp.models.DistilBertBackbone(
         vocabulary_size=30552,
         num_layers=6,
         num_heads=12,
@@ -171,8 +171,8 @@ class DistilBertClassifier(keras.Model):
         ```
         """
         # Check if preset is backbone-only model
-        if preset in DistilBert.presets:
-            backbone = DistilBert.from_preset(preset, load_weights)
+        if preset in DistilBertBackbone.presets:
+            backbone = DistilBertBackbone.from_preset(preset, load_weights)
             return cls(backbone, **kwargs)
 
         # Otherwise must be one of class presets
diff --git a/tools/checkpoint_conversion/convert_distilbert_checkpoints.py b/tools/checkpoint_conversion/convert_distilbert_checkpoints.py
index c78e9833..ca2f5a59 100644
--- a/tools/checkpoint_conversion/convert_distilbert_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_distilbert_checkpoints.py
@@ -288,7 +288,7 @@ def main(_):
     )
 
     print("\n-> Load KerasNLP model.")
-    keras_nlp_model = keras_nlp.models.DistilBert.from_preset(
+    keras_nlp_model = keras_nlp.models.DistilBertBackbone.from_preset(
         FLAGS.preset, load_weights=False
     )
 


Test patch:
diff --git a/keras_nlp/models/distil_bert/distil_bert_backbone_test.py b/keras_nlp/models/distil_bert/distil_bert_backbone_test.py
index f5864e7b..4a50bf45 100644
--- a/keras_nlp/models/distil_bert/distil_bert_backbone_test.py
+++ b/keras_nlp/models/distil_bert/distil_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.distil_bert.distil_bert_backbone import DistilBert
+from keras_nlp.models.distil_bert.distil_bert_backbone import DistilBertBackbone
 
 
 class DistilBertTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = DistilBert(
+        self.model = DistilBertBackbone(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
@@ -86,7 +86,7 @@ class DistilBertTest(tf.test.TestCase, parameterized.TestCase):
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
-        self.assertIsInstance(restored_model, DistilBert)
+        self.assertIsInstance(restored_model, DistilBertBackbone)
 
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
diff --git a/keras_nlp/models/distil_bert/distil_bert_classifier_test.py b/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
index be2370cd..7ffc70d8 100644
--- a/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
@@ -19,7 +19,7 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.distil_bert.distil_bert_backbone import DistilBert
+from keras_nlp.models.distil_bert.distil_bert_backbone import DistilBertBackbone
 from keras_nlp.models.distil_bert.distil_bert_classifier import (
     DistilBertClassifier,
 )
@@ -27,7 +27,7 @@ from keras_nlp.models.distil_bert.distil_bert_classifier import (
 
 class DistilBertClassifierTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = DistilBert(
+        self.model = DistilBertBackbone(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
diff --git a/keras_nlp/models/distil_bert/distil_bert_presets_test.py b/keras_nlp/models/distil_bert/distil_bert_presets_test.py
index b25ac989..024710ca 100644
--- a/keras_nlp/models/distil_bert/distil_bert_presets_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_presets_test.py
@@ -17,7 +17,7 @@ import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 
-from keras_nlp.models.distil_bert.distil_bert_backbone import DistilBert
+from keras_nlp.models.distil_bert.distil_bert_backbone import DistilBertBackbone
 from keras_nlp.models.distil_bert.distil_bert_classifier import (
     DistilBertClassifier,
 )
@@ -63,7 +63,7 @@ class DistilBertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
             "token_ids": tf.constant([[101, 1996, 4248, 102]]),
             "padding_mask": tf.constant([[1, 1, 1, 1]]),
         }
-        model = DistilBert.from_preset(
+        model = DistilBertBackbone.from_preset(
             "distil_bert_base_en_uncased", load_weights=load_weights
         )
         outputs = model(input_data)[0, 0, :5]
@@ -87,7 +87,7 @@ class DistilBertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
     @parameterized.named_parameters(
         ("distilbert_tokenizer", DistilBertTokenizer),
         ("distilbert_preprocessor", DistilBertPreprocessor),
-        ("distilbert", DistilBert),
+        ("distilbert", DistilBertBackbone),
         ("distilbert_classifier", DistilBertClassifier),
     )
     def test_preset_docstring(self, cls):
@@ -98,7 +98,7 @@ class DistilBertPresetSmokeTest(tf.test.TestCase, parameterized.TestCase):
     @parameterized.named_parameters(
         ("distilbert_tokenizer", DistilBertTokenizer),
         ("distilbert_preprocessor", DistilBertPreprocessor),
-        ("distilbert", DistilBert),
+        ("distilbert", DistilBertBackbone),
         ("distilbert_classifier", DistilBertClassifier),
     )
     def test_unknown_preset_error(self, cls):
@@ -121,8 +121,10 @@ class DistilBertPresetFullTest(tf.test.TestCase, parameterized.TestCase):
         ("preset_weights", True), ("random_weights", False)
     )
     def test_load_distilbert(self, load_weights):
-        for preset in DistilBert.presets:
-            model = DistilBert.from_preset(preset, load_weights=load_weights)
+        for preset in DistilBertBackbone.presets:
+            model = DistilBertBackbone.from_preset(
+                preset, load_weights=load_weights
+            )
             input_data = {
                 "token_ids": tf.random.uniform(
                     shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size