Repository: keras-team/keras-nlp

Issue:
Create unique string id for each BERT pretrained checkpoint
Currently BERT pretrained encoders are specified by their vocabulary alone. This is possible because checkpoints must be called from the correct class (e.g., `BertBase('uncased_en')`).

However, if we want to request checkpoints from any combination of architectures and vocabularies in a task-specific model (e.g., `BertClassifier`) it will be much easier if each checkpoint has a unique name.

Other benefits:
* Will also make it easy to document all of our checkpoints in keras.io. There are many possible checkpoints that could be specified by `'uncased_en'` but only one by `'bert_base_uncased_en'`.
* Currently the `backbone` model can only be specified by instantiating a Bert model. This change allows string ids and defaults for the backbone models.


Discussion:

Resolving patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -13,13 +13,13 @@
 # limitations under the License.
 
 from keras_nlp.models.bert import BertBase
-from keras_nlp.models.bert import BertClassifier
 from keras_nlp.models.bert import BertCustom
 from keras_nlp.models.bert import BertLarge
 from keras_nlp.models.bert import BertMedium
 from keras_nlp.models.bert import BertPreprocessor
 from keras_nlp.models.bert import BertSmall
 from keras_nlp.models.bert import BertTiny
+from keras_nlp.models.bert_tasks import BertClassifier
 from keras_nlp.models.gpt2 import Gpt2Base
 from keras_nlp.models.gpt2 import Gpt2Custom
 from keras_nlp.models.gpt2 import Gpt2ExtraLarge
diff --git a/keras_nlp/models/bert.py b/keras_nlp/models/bert.py
--- a/keras_nlp/models/bert.py
+++ b/keras_nlp/models/bert.py
@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""BERT model configurable class, preconfigured versions, and task heads."""
+"""BERT model configurable class and preconfigured versions."""
 
 import os
+from collections import defaultdict
 
 import tensorflow as tf
 from tensorflow import keras
@@ -25,93 +26,111 @@
 from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
 
 
-def _bert_kernel_initializer(stddev=0.02):
+def bert_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
+# TODO(jbischof): document checkpoints in keras.io and use URL in docstrings
 # Metadata for loading pretrained model weights.
 checkpoints = {
-    "bert_tiny": {
-        "uncased_en": {
-            "description": (
-                "Tiny size of BERT where all input is lowercased. "
-                "Trained on English Wikipedia + BooksCorpus."
-            ),
-            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_tiny_uncased_en/model.h5",
-            "weights_hash": "c2b29fcbf8f814a0812e4ab89ef5c068",
-        }
+    "bert_tiny_uncased_en": {
+        "model": "BertTiny",
+        "vocabulary": "uncased_en",
+        "description": (
+            "Tiny size of BERT where all input is lowercased. "
+            "Trained on English Wikipedia + BooksCorpus."
+        ),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_tiny_uncased_en/model.h5",
+        "weights_hash": "c2b29fcbf8f814a0812e4ab89ef5c068",
     },
-    "bert_small": {
-        "uncased_en": {
-            "description": (
-                "Small size of BERT where all input is lowercased. "
-                "Trained on English Wikipedia + BooksCorpus."
-            ),
-            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_small_uncased_en/model.h5",
-            "weights_hash": "08632c9479b034f342ba2c2b7afba5f7",
-        }
+    "bert_small_uncased_en": {
+        "model": "BertSmall",
+        "vocabulary": "uncased_en",
+        "description": (
+            "Small size of BERT where all input is lowercased. "
+            "Trained on English Wikipedia + BooksCorpus."
+        ),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_small_uncased_en/model.h5",
+        "weights_hash": "08632c9479b034f342ba2c2b7afba5f7",
     },
-    "bert_medium": {
-        "uncased_en": {
-            "description": (
-                "Medium size of BERT where all input is lowercased. "
-                "Trained on English Wikipedia + BooksCorpus."
-            ),
-            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_medium_uncased_en/model.h5",
-            "weights_hash": "bb990e1184ec6b6185450c73833cd661",
-        }
+    "bert_medium_uncased_en": {
+        "model": "BertMedium",
+        "vocabulary": "uncased_en",
+        "description": (
+            "Medium size of BERT where all input is lowercased. "
+            "Trained on English Wikipedia + BooksCorpus."
+        ),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_medium_uncased_en/model.h5",
+        "weights_hash": "bb990e1184ec6b6185450c73833cd661",
     },
-    "bert_base": {
-        "uncased_en": {
-            "description": (
-                "Base size of BERT where all input is lowercased. "
-                "Trained on English Wikipedia + BooksCorpus."
-            ),
-            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_uncased_en/model.h5",
-            "weights_hash": "9b2b2139f221988759ac9cdd17050b31",
-        },
-        "cased_en": {
-            "description": (
-                "Base size of BERT where case is maintained. "
-                "Trained on English Wikipedia + BooksCorpus."
-            ),
-            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_cased_en/model.h5",
-            "weights_hash": "f94a6cb012e18f4fb8ec92abb91864e9",
-        },
-        "zh": {
-            "description": ("Base size of BERT. Trained on Chinese Wikipedia."),
-            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_zh/model.h5",
-            "weights_hash": "79afa421e386076e62ab42dad555ab0c",
-        },
-        "multi_cased": {
-            "description": (
-                "Base size of BERT. Trained on Wikipedias of 104 languages."
-            ),
-            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_multi_cased/model.h5",
-            "weights_hash": "b0631cec0a1f2513c6cfd75ba29c33aa",
-        },
+    "bert_base_uncased_en": {
+        "model": "BertBase",
+        "vocabulary": "uncased_en",
+        "description": (
+            "Base size of BERT where all input is lowercased. "
+            "Trained on English Wikipedia + BooksCorpus."
+        ),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_uncased_en/model.h5",
+        "weights_hash": "9b2b2139f221988759ac9cdd17050b31",
+    },
+    "bert_base_cased_en": {
+        "model": "BertBase",
+        "vocabulary": "cased_en",
+        "description": (
+            "Base size of BERT where case is maintained. "
+            "Trained on English Wikipedia + BooksCorpus."
+        ),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_cased_en/model.h5",
+        "weights_hash": "f94a6cb012e18f4fb8ec92abb91864e9",
     },
-    "bert_large": {
-        "uncased_en": {
-            "description": (
-                "Large size of BERT where all input is lowercased. "
-                "Trained on English Wikipedia + BooksCorpus."
-            ),
-            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_large_uncased_en/model.h5",
-            "weights_hash": "cc5cacc9565ef400ee4376105f40ddae",
-        },
-        "cased_en": {
-            "description": (
-                "Large size of BERT where case is maintained. "
-                "Trained on English Wikipedia + BooksCorpus."
-            ),
-            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_large_cased_en/model.h5",
-            "weights_hash": "8b8ab82290bbf4f8db87d4f100648890",
-        },
+    "bert_base_zh": {
+        "model": "BertBase",
+        "vocabulary": "zh",
+        "description": ("Base size of BERT. Trained on Chinese Wikipedia."),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_zh/model.h5",
+        "weights_hash": "79afa421e386076e62ab42dad555ab0c",
+    },
+    "bert_base_multi_cased": {
+        "model": "BertBase",
+        "vocabulary": "multi_cased",
+        "description": ("Base size of BERT. Trained on Chinese Wikipedia."),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_multi_cased/model.h5",
+        "weights_hash": "b0631cec0a1f2513c6cfd75ba29c33aa",
+    },
+    "bert_large_uncased_en": {
+        "model": "BertLarge",
+        "vocabulary": "uncased_en",
+        "description": (
+            "Large size of BERT where all input is lowercased. "
+            "Trained on English Wikipedia + BooksCorpus."
+        ),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_large_uncased_en/model.h5",
+        "weights_hash": "cc5cacc9565ef400ee4376105f40ddae",
+    },
+    "bert_large_cased_en": {
+        "model": "BertLarge",
+        "vocabulary": "cased_en",
+        "description": (
+            "Base size of BERT where case is maintained. "
+            "Trained on English Wikipedia + BooksCorpus."
+        ),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_large_cased_en/model.h5",
+        "weights_hash": "8b8ab82290bbf4f8db87d4f100648890",
     },
 }
 
 
+# Index checkpoints by arch compatibility and create lookup function
+checkpoints_per_arch = defaultdict(set)
+for arch, metadata in checkpoints.items():
+    checkpoints_per_arch[metadata["model"]].add(arch)
+
+
+def compatible_checkpoints(arch):
+    """Returns a list of compatible checkpoints per arch"""
+    return checkpoints_per_arch[arch]
+
+
 # Metadata for loading pretrained tokenizer vocabularies.
 # We need the vocabulary_size hardcoded so we can instantiate a BERT network
 # with the right embedding size without downloading the matching vocabulary.
@@ -167,6 +186,7 @@ def _handle_pretrained_tokenizer_arguments(vocabulary, lowercase):
     fully resolve them in the case we are loading pretrained weights.
     """
 
+    # TODO(jbischof): allow user to lookup vocabulary from checkpoint name
     if isinstance(vocabulary, str) and vocabulary in vocabularies:
         metadata = vocabularies[vocabulary]
         vocabulary = keras.utils.get_file(
@@ -197,14 +217,16 @@ def _handle_pretrained_model_arguments(bert_variant, weights, vocabulary_size):
         )
 
     if weights:
-        if weights not in checkpoints[bert_variant]:
+        arch_checkpoints = compatible_checkpoints(bert_variant)
+        if weights not in arch_checkpoints:
             raise ValueError(
                 "`weights` must be one of "
-                f"""{", ".join(checkpoints[bert_variant])}. """
+                f"""{", ".join(arch_checkpoints)}. """
                 f"Received: {weights}"
             )
-        metadata = checkpoints[bert_variant][weights]
-        vocabulary_size = vocabularies[weights]["vocabulary_size"]
+        metadata = checkpoints[weights]
+        vocabulary = metadata["vocabulary"]
+        vocabulary_size = vocabularies[vocabulary]["vocabulary_size"]
 
         # TODO(jbischof): consider changing format from `h5` to
         # `tf.train.Checkpoint` once
@@ -212,7 +234,7 @@ def _handle_pretrained_model_arguments(bert_variant, weights, vocabulary_size):
         weights = keras.utils.get_file(
             "model.h5",
             metadata["weights_url"],
-            cache_subdir=os.path.join("models", "bert", weights, bert_variant),
+            cache_subdir=os.path.join("models", weights),
             file_hash=metadata["weights_hash"],
         )
 
@@ -314,19 +336,19 @@ def __init__(
         token_embedding_layer = keras.layers.Embedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
-            embeddings_initializer=_bert_kernel_initializer(),
+            embeddings_initializer=bert_kernel_initializer(),
             name="token_embedding",
         )
         token_embedding = token_embedding_layer(token_id_input)
         position_embedding = PositionEmbedding(
-            initializer=_bert_kernel_initializer(),
+            initializer=bert_kernel_initializer(),
             sequence_length=max_sequence_length,
             name="position_embedding",
         )(token_embedding)
         segment_embedding = keras.layers.Embedding(
             input_dim=num_segments,
             output_dim=hidden_dim,
-            embeddings_initializer=_bert_kernel_initializer(),
+            embeddings_initializer=bert_kernel_initializer(),
             name="segment_embedding",
         )(segment_id_input)
 
@@ -354,7 +376,7 @@ def __init__(
                     x, approximate=True
                 ),
                 dropout=dropout,
-                kernel_initializer=_bert_kernel_initializer(),
+                kernel_initializer=bert_kernel_initializer(),
                 name=f"transformer_layer_{i}",
             )(x, padding_mask=padding_mask)
 
@@ -363,7 +385,7 @@ def __init__(
         sequence_output = x
         pooled_output = keras.layers.Dense(
             hidden_dim,
-            kernel_initializer=_bert_kernel_initializer(),
+            kernel_initializer=bert_kernel_initializer(),
             activation="tanh",
             name="pooled_dense",
         )(x[:, cls_token_index, :])
@@ -561,69 +583,7 @@ def call(self, inputs):
 )
 
 
-class BertClassifier(keras.Model):
-    """BERT encoder model with a classification head.
-
-    Args:
-        base_model: A `keras_nlp.models.BertCustom` to encode inputs.
-        num_classes: int. Number of classes to predict.
-        name: string, optional. Name of the model.
-        trainable: boolean, optional. If the model's variables should be
-            trainable.
-
-    Examples:
-    ```python
-    # Randomly initialized BERT encoder
-    model = keras_nlp.models.BertCustom(
-        vocabulary_size=30522,
-        num_layers=12,
-        num_heads=12,
-        hidden_dim=768,
-        intermediate_dim=3072,
-        max_sequence_length=12
-    )
-
-    # Call classifier on the inputs.
-    input_data = {
-        "token_ids": tf.random.uniform(
-            shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
-        ),
-        "segment_ids": tf.constant(
-            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-        ),
-        "padding_mask": tf.constant(
-            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-        ),
-    }
-    classifier = bert.BertClassifier(model, 4, name="classifier")
-    logits = classifier(input_data)
-    ```
-    """
-
-    def __init__(
-        self,
-        base_model,
-        num_classes,
-        name=None,
-        trainable=True,
-    ):
-        inputs = base_model.input
-        pooled = base_model(inputs)["pooled_output"]
-        outputs = keras.layers.Dense(
-            num_classes,
-            kernel_initializer=_bert_kernel_initializer(),
-            name="logits",
-        )(pooled)
-        # Instantiate using Functional API Model constructor
-        super().__init__(
-            inputs=inputs, outputs=outputs, name=name, trainable=trainable
-        )
-        # All references to `self` below this line
-        self.base_model = base_model
-        self.num_classes = num_classes
-
-
-MODEL_DOCSTRING = """BERT "{type}" architecture.
+MODEL_DOCSTRING = """Bert "{type}" architecture.
 
     This network implements a bi-directional Transformer-based encoder as
     described in ["BERT: Pre-training of Deep Bidirectional Transformers for
@@ -658,7 +618,7 @@ def __init__(
     output = model(input_data)
 
     # Load a pretrained model
-    model = keras_nlp.models.Bert{type}(weights="uncased_en")
+    model = keras_nlp.models.Bert{type}(weights="bert_base_uncased_en")
     # Call encoder on the inputs.
     output = model(input_data)
     ```
@@ -667,7 +627,7 @@ def __init__(
 
 def BertTiny(weights=None, vocabulary_size=None, name=None, trainable=True):
     weights, vocabulary_size = _handle_pretrained_model_arguments(
-        "bert_tiny", weights, vocabulary_size
+        "BertTiny", weights, vocabulary_size
     )
 
     model = BertCustom(
@@ -685,14 +645,12 @@ def BertTiny(weights=None, vocabulary_size=None, name=None, trainable=True):
     if weights:
         model.load_weights(weights)
 
-    # TODO(jbischof): attach the tokenizer or create separate tokenizer class.
-    # This comment applies to other variants as well.
     return model
 
 
 def BertSmall(weights=None, vocabulary_size=None, name=None, trainable=True):
     weights, vocabulary_size = _handle_pretrained_model_arguments(
-        "bert_small", weights, vocabulary_size
+        "BertSmall", weights, vocabulary_size
     )
 
     model = BertCustom(
@@ -715,7 +673,7 @@ def BertSmall(weights=None, vocabulary_size=None, name=None, trainable=True):
 
 def BertMedium(weights=None, vocabulary_size=None, name=None, trainable=True):
     weights, vocabulary_size = _handle_pretrained_model_arguments(
-        "bert_medium", weights, vocabulary_size
+        "BertMedium", weights, vocabulary_size
     )
 
     model = BertCustom(
@@ -738,7 +696,7 @@ def BertMedium(weights=None, vocabulary_size=None, name=None, trainable=True):
 
 def BertBase(weights=None, vocabulary_size=None, name=None, trainable=True):
     weights, vocabulary_size = _handle_pretrained_model_arguments(
-        "bert_base", weights, vocabulary_size
+        "BertBase", weights, vocabulary_size
     )
 
     model = BertCustom(
@@ -761,7 +719,7 @@ def BertBase(weights=None, vocabulary_size=None, name=None, trainable=True):
 
 def BertLarge(weights=None, vocabulary_size=None, name=None, trainable=True):
     weights, vocabulary_size = _handle_pretrained_model_arguments(
-        "bert_large", weights, vocabulary_size
+        "BertLarge", weights, vocabulary_size
     )
 
     model = BertCustom(
@@ -782,11 +740,22 @@ def BertLarge(weights=None, vocabulary_size=None, name=None, trainable=True):
     return model
 
 
+def model_class_by_name(classname):
+    """Return model class given the class name."""
+    return {
+        "BertTiny": BertTiny,
+        "BertSmall": BertSmall,
+        "BertMedium": BertMedium,
+        "BertBase": BertBase,
+        "BertLarge": BertLarge,
+    }[classname]
+
+
 setattr(
     BertTiny,
     "__doc__",
     MODEL_DOCSTRING.format(
-        type="Tiny", names=", ".join(checkpoints["bert_tiny"])
+        type="Tiny", names=", ".join(compatible_checkpoints("BertTiny"))
     ),
 )
 
@@ -794,7 +763,7 @@ def BertLarge(weights=None, vocabulary_size=None, name=None, trainable=True):
     BertSmall,
     "__doc__",
     MODEL_DOCSTRING.format(
-        type="Small", names=", ".join(checkpoints["bert_small"])
+        type="Small", names=", ".join(compatible_checkpoints("BertSmall"))
     ),
 )
 
@@ -802,7 +771,7 @@ def BertLarge(weights=None, vocabulary_size=None, name=None, trainable=True):
     BertMedium,
     "__doc__",
     MODEL_DOCSTRING.format(
-        type="Medium", names=", ".join(checkpoints["bert_medium"])
+        type="Medium", names=", ".join(compatible_checkpoints("BertMedium"))
     ),
 )
 
@@ -810,13 +779,13 @@ def BertLarge(weights=None, vocabulary_size=None, name=None, trainable=True):
     BertBase,
     "__doc__",
     MODEL_DOCSTRING.format(
-        type="Base", names=", ".join(checkpoints["bert_base"])
+        type="Base", names=", ".join(compatible_checkpoints("BertBase"))
     ),
 )
 setattr(
     BertLarge,
     "__doc__",
     MODEL_DOCSTRING.format(
-        type="Large", names=", ".join(checkpoints["bert_large"])
+        type="Large", names=", ".join(compatible_checkpoints("BertLarge"))
     ),
 )
diff --git a/keras_nlp/models/bert_tasks.py b/keras_nlp/models/bert_tasks.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/models/bert_tasks.py
@@ -0,0 +1,113 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""BERT model task heads and preconfigured versions."""
+
+from tensorflow import keras
+
+from keras_nlp.models import bert
+from keras_nlp.models.bert import bert_kernel_initializer
+from keras_nlp.models.bert import checkpoints
+
+# TODO(jbischof): Find more scalable way to list checkpoints.
+CLASSIFIER_DOCSTRING = """BERT encoder model with a classification head.
+
+    Args:
+        backbone: A string, `keras_nlp.models.BertCustom` or derivative such as
+            `keras_nlp.models.BertBase` to encode inputs. If a string, should be
+            one of {names}.
+        num_classes: int. Number of classes to predict.
+        name: string, optional. Name of the model.
+        trainable: boolean, optional. If the model's variables should be
+            trainable.
+
+    Examples:
+    ```python
+    # Randomly initialized BERT encoder
+    model = keras_nlp.models.BertCustom(
+        vocabulary_size=30522,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=12
+    )
+
+    # Call classifier on the inputs.
+    input_data = {{
+        "token_ids": tf.random.uniform(
+            shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
+        ),
+        "segment_ids": tf.constant(
+            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
+        "padding_mask": tf.constant(
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
+    }}
+    classifier = keras_nlp.models.BertClassifier(model, 4, name="classifier")
+    logits = classifier(input_data)
+
+    # String backbone specification
+    classifier = keras_nlp.models.BertClassifier(
+        "bert_base_uncased_en", 4, name="classifier"
+    )
+    logits = classifier(input_data)
+    ```
+"""
+
+
+class BertClassifier(keras.Model):
+    def __init__(
+        self,
+        backbone="bert_base_uncased_en",
+        num_classes=2,
+        name=None,
+        trainable=True,
+    ):
+        # Load backbone from string identifier
+        # TODO(jbischof): create util function when ready to load backbones in
+        # other task classes (e.g., load_backbone_from_string())
+        if isinstance(backbone, str):
+            if backbone not in bert.checkpoints:
+                raise ValueError(
+                    "`backbone` must be one of "
+                    f"""{", ".join(bert.checkpoints.keys())}. """
+                    f"Received: {backbone}"
+                )
+            backbone_class_str = bert.checkpoints[backbone]["model"]
+            backbone_class = bert.model_class_by_name(backbone_class_str)
+            backbone = backbone_class(backbone)
+
+        inputs = backbone.input
+        pooled = backbone(inputs)["pooled_output"]
+        outputs = keras.layers.Dense(
+            num_classes,
+            kernel_initializer=bert_kernel_initializer(),
+            name="logits",
+        )(pooled)
+        # Instantiate using Functional API Model constructor
+        super().__init__(
+            inputs=inputs, outputs=outputs, name=name, trainable=trainable
+        )
+        # All references to `self` below this line
+        self.backbone = backbone
+        self.num_classes = num_classes
+
+
+setattr(
+    BertClassifier,
+    "__doc__",
+    CLASSIFIER_DOCSTRING.format(names=", ".join(checkpoints.keys())),
+)


Test patch:
diff --git a/keras_nlp/models/bert_tasks_test.py b/keras_nlp/models/bert_tasks_test.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/models/bert_tasks_test.py
@@ -0,0 +1,88 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for Bert model."""
+
+import os
+
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models import bert
+from keras_nlp.models import bert_tasks
+
+
+class BertTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        self.backbone = bert.BertCustom(
+            vocabulary_size=1000,
+            num_layers=2,
+            num_heads=2,
+            hidden_dim=64,
+            intermediate_dim=128,
+            max_sequence_length=128,
+            name="encoder",
+        )
+        self.classifier = bert_tasks.BertClassifier(
+            self.backbone, 4, name="classifier"
+        )
+        self.batch_size = 8
+        self.input_batch = {
+            "token_ids": tf.ones(
+                (self.batch_size, self.backbone.max_sequence_length),
+                dtype="int32",
+            ),
+            "segment_ids": tf.ones(
+                (self.batch_size, self.backbone.max_sequence_length),
+                dtype="int32",
+            ),
+            "padding_mask": tf.ones(
+                (self.batch_size, self.backbone.max_sequence_length),
+                dtype="int32",
+            ),
+        }
+
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
+    def test_valid_call_classifier(self):
+        self.classifier(self.input_batch)
+
+        # Not a checkpoint name
+        with self.assertRaises(ValueError):
+            bert_tasks.BertClassifier("bert_clowntown", 4, name="classifier")
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_bert_classifier_compile(self, jit_compile):
+        self.classifier.compile(jit_compile=jit_compile)
+        self.classifier.predict(self.input_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_bert_classifier_compile_batched_ds(self, jit_compile):
+        self.classifier.compile(jit_compile=jit_compile)
+        self.classifier.predict(self.input_dataset)
+
+    def test_saving_model(self):
+        model_output = self.classifier(self.input_batch)
+        save_path = os.path.join(self.get_temp_dir(), "model")
+        self.classifier.save(save_path)
+        restored_model = keras.models.load_model(save_path)
+
+        restored_output = restored_model(self.input_batch)
+        self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/models/bert_test.py b/keras_nlp/models/bert_test.py
--- a/keras_nlp/models/bert_test.py
+++ b/keras_nlp/models/bert_test.py
@@ -106,10 +106,6 @@ def test_variable_sequence_length_call_bert(self):
             }
             self.model(input_data)
 
-    def test_valid_call_classifier(self):
-        classifier = bert.BertClassifier(self.model, 4, name="classifier")
-        classifier(self.input_batch)
-
     def test_valid_call_bert_base(self):
         model = bert.BertBase(vocabulary_size=1000, name="encoder")
         input_data = {
@@ -141,22 +137,6 @@ def test_bert_base_compile_batched_ds(self, jit_compile):
         model.compile(jit_compile=jit_compile)
         model.predict(self.input_dataset)
 
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_bert_classifier_compile(self, jit_compile):
-        model = bert.BertClassifier(self.model, 4, name="classifier")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_batch)
-
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_bert_classifier_compile_batched_ds(self, jit_compile):
-        model = bert.BertClassifier(self.model, 4, name="classifier")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_dataset)
-
     def test_bert_base_vocab_error(self):
         # Need `vocabulary_size` or `weights`
         with self.assertRaises(ValueError):
@@ -165,7 +145,7 @@ def test_bert_base_vocab_error(self):
         # Only one of `vocabulary_size` or `weights`
         with self.assertRaises(ValueError):
             bert.BertBase(
-                weights="bert_base_uncased",
+                weights="bert_base_uncased_en",
                 vocabulary_size=1000,
                 name="encoder",
             )
@@ -173,7 +153,7 @@ def test_bert_base_vocab_error(self):
         # Not a checkpoint name
         with self.assertRaises(ValueError):
             bert.BertBase(
-                weights="bert_base_clowntown",
+                weights="bert_base_uncased_clowntown",
                 name="encoder",
             )
 
diff --git a/keras_nlp/network_tests/test_load_ckpts.py b/keras_nlp/network_tests/test_load_ckpts.py
--- a/keras_nlp/network_tests/test_load_ckpts.py
+++ b/keras_nlp/network_tests/test_load_ckpts.py
@@ -15,41 +15,65 @@
 
 import pytest
 import tensorflow as tf
-from absl.testing import parameterized
 
 import keras_nlp
+from keras_nlp.models.bert import checkpoints as bert_checkpoints
+from keras_nlp.models.bert import (
+    model_class_by_name as bert_model_class_by_name,
+)
+from keras_nlp.models.bert import vocabularies as bert_vocabularies
 
 
 @pytest.mark.slow
-class BertCkptTest(tf.test.TestCase, parameterized.TestCase):
-    @parameterized.named_parameters(
-        ("tiny_uncased_en", keras_nlp.models.BertTiny, "uncased_en"),
-        ("small_uncased_en", keras_nlp.models.BertSmall, "uncased_en"),
-        ("medium_uncased_en", keras_nlp.models.BertMedium, "uncased_en"),
-        ("base_uncased_en", keras_nlp.models.BertBase, "uncased_en"),
-        ("base_cased_en", keras_nlp.models.BertBase, "cased_en"),
-        ("base_zh", keras_nlp.models.BertBase, "zh"),
-        ("base_multi_cased", keras_nlp.models.BertBase, "multi_cased"),
-        ("large_uncased_en", keras_nlp.models.BertLarge, "uncased_en"),
-        ("large_cased_en", keras_nlp.models.BertLarge, "cased_en"),
-    )
-    def test_load_bert(self, bert_variant, weights):
-        model = bert_variant(weights=weights)
+class BertCkptTest(tf.test.TestCase):
+    def test_load_bert(self):
+        for checkpoint in bert_checkpoints:
+            bert_variant = bert_checkpoints[checkpoint]["model"]
+            bert_class = bert_model_class_by_name(bert_variant)
+            model = bert_class(weights=checkpoint)
+            input_data = {
+                "token_ids": tf.random.uniform(
+                    shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
+                ),
+                "segment_ids": tf.constant(
+                    [0] * 200 + [1] * 312, shape=(1, 512)
+                ),
+                "padding_mask": tf.constant([1] * 512, shape=(1, 512)),
+            }
+            model(input_data)
+
+    def test_load_bert_backbone_string(self):
+        for checkpoint in bert_checkpoints:
+            classifier = keras_nlp.models.bert_tasks.BertClassifier(
+                checkpoint, 4, name="classifier"
+            )
+            input_data = {
+                "token_ids": tf.random.uniform(
+                    shape=(1, 512),
+                    dtype=tf.int64,
+                    maxval=classifier.backbone.vocabulary_size,
+                ),
+                "segment_ids": tf.constant(
+                    [0] * 200 + [1] * 312, shape=(1, 512)
+                ),
+                "padding_mask": tf.constant([1] * 512, shape=(1, 512)),
+            }
+            classifier(input_data)
+
+    def test_classifier_default_args(self):
+        classifier = keras_nlp.models.bert_tasks.BertClassifier()
         input_data = {
             "token_ids": tf.random.uniform(
-                shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
+                shape=(1, 512),
+                dtype=tf.int64,
+                maxval=classifier.backbone.vocabulary_size,
             ),
             "segment_ids": tf.constant([0] * 200 + [1] * 312, shape=(1, 512)),
             "padding_mask": tf.constant([1] * 512, shape=(1, 512)),
         }
-        model(input_data)
+        classifier(input_data)
 
-    @parameterized.named_parameters(
-        ("uncased_en", "uncased_en"),
-        ("cased_en", "cased_en"),
-        ("zh", "zh"),
-        ("multi_cased", "multi_cased"),
-    )
-    def test_load_vocabularies(self, vocabulary):
-        tokenizer = keras_nlp.models.BertPreprocessor(vocabulary=vocabulary)
-        tokenizer("The quick brown fox.")
+    def test_load_vocabularies(self):
+        for vocabulary in bert_vocabularies:
+            tokenizer = keras_nlp.models.BertPreprocessor(vocabulary=vocabulary)
+            tokenizer("The quick brown fox.")