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:
jbischof: @mattdangerw actually we could make this more maintainable by moving `backbone_checkpoints` to `bert.py` and making it the source of truth. Then the current checkpoints struct would have the sole job of telling us the which checkpoints work with each class (e.g., 'uncased_en' for `BertBase` is in `backbone_checkpoints['bert_base_uncased_en']`). 
jbischof: I am confident I can clean up and consolidate the structs in `bert.py`, but mainly we need to decide if we want unique checkpoint ids or the shortest strings possible.
```python
# Not unique, but short as possible
BertBase('uncased_en')
# Unique but not short
BertBase('bert_base_uncased_en')
```

mattdangerw: Re the question of unique vs shortest strings, it seems like both could work, with very interesting tradeoffs.

But it seems incomplete without considering `BertPreprocessor`. Unique identifiers for checkpoints, but shortest possible identifiers for preprocessing feels like we would be asking users to mix two different mental models.

The shortest possible string proposal with preprocessing:
```
BertPreprocessor('uncased_en')
BertBase('uncased_en')
BertClassifier(backbone='base_uncased_en')

# And down the line...
BertClassifier('sst_base_uncase_en')
```

The unique ID proposal, also applied to preprocessing:
```
BertPreprocessor('bert_base_uncased_en')
BertBase('bert_base_uncased_en')
BertClassifier(backbone='bert_base_uncased_en')

# And down the line...
BertPreprocessor('bert_sst_base_uncased_en')
BertClassifier('bert_sst_base_uncased_en')
```

It feels like a non-obvious choice. Honestly the main reason to go with an ID approach would seems to be thinking about this in a far off future state, where we want to handle tons of checkpoints, or even integrating with some (or many!) external hubs. I'm not really sure how to evaluate that.

The former option seems the most in line with Keras today. Not trying to say that should trump other considerations though.
jbischof: Excellent point @mattdangerw! I think the short answer is that the preprocessor construction needs to work from either a model name prompt or direct vocab prompt. The user may want to train their own model using a pretrained vocab or want us to find the right vocab for a checkpoint.

```python
# Direct request for vocab
BertPreprocessor(vocabulary='uncased_en')
# Request to for correct vocab for a checkpoint
BertPreprocessor(weights='bert_based_uncased_en')
# Especially helpful for task driven models where vocab may be obvious
BertPreprocess(weights="ToxicBERT")
```

Since using `weights` is a bit confusing in this context (is this the checkpoint for a vocab or the vocab for a checkpoint?), we could also allow either to be specified as `vocabulary`. A model will basically never have the same name as a vocab.

```python
# Direct request for vocab
BertPreprocessor(vocabulary='uncased_en')
# Request to for correct vocab for a checkpoint
BertPreprocessor(vocabulary='bert_based_uncased_en')
# Especially helpful for task driven models where vocab may be obvious
BertPreprocess(vocabulary="ToxicBERT")
```

Under the hood we first look in the `vocabularies` struct and then `checkpoints` one. Now a user can never misunderstand...we accept both! WDYT?
mattdangerw: In the ID world (if we go with that), why does `'uncased_en'` have to mean anything anymore? I'm not sure it does. You could either provide a checkpoint ID, or you provide a vocab yourself.

```python
BertPreprocessor(vocabulary='bert_base_uncased_en')
or
BertPreprocessor(vocabulary='./my-cool-vocab.txt')
```

Worry with your proposal above is two fold I guess...

1) The whole point of fully specified string name would be to give the name universal meaning throughout all symbols we offer. That alone is sufficient right? Why also introduce another namespace just for specifying vocabularies directly? We never need to support a preset vocabulary that is not in that universal ID namespace.
2) `weights` current in Keras applications is supposed to either be a "name shortcut" or a literal weights file (tf or h5 format). The analog for `vocabulary` would seem to be, either a "name shortcut" or a literal vocab file. It seems most consistent to keep that approach. If we want to diverge to avoid overloaded args, it seems better to do that with a totally different name. `BertPreprocessor(preset='bert_base_uncased_en')`, `BertBase(preset='bert_base_uncased_en')` or `BertPreprocessor(vocabulary='./some-vocab.txt')`, `BertPreprocessor(weights='./some-weights.h5')`. Using `weights` in the way you are proposing above seems like it is both inconsistent with the rest of Keras, and breaking the user mental model for the name (i.e. that `weights` refers to some literal floats somewhere).
jbischof: Sounds good @mattdangerw let's figure out the preprocessor in a followup PR! We were just sanity checking how unique string ids work for vocabs.
jbischof: Got offline approval from @fchollet 

Patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index e1635a4a..b33a79da 100644
--- 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
index 5e2efad1..05d3ca89 100644
--- 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.layers.transformer_encoder import TransformerEncoder
 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 @@ class BertCustom(keras.Model):
         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 @@ class BertCustom(keras.Model):
                     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 @@ class BertCustom(keras.Model):
         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 @@ setattr(
 )
 
 
-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 @@ MODEL_DOCSTRING = """BERT "{type}" architecture.
     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 @@ MODEL_DOCSTRING = """BERT "{type}" architecture.
 
 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 @@ setattr(
     BertSmall,
     "__doc__",
     MODEL_DOCSTRING.format(
-        type="Small", names=", ".join(checkpoints["bert_small"])
+        type="Small", names=", ".join(compatible_checkpoints("BertSmall"))
     ),
 )
 
@@ -802,7 +771,7 @@ setattr(
     BertMedium,
     "__doc__",
     MODEL_DOCSTRING.format(
-        type="Medium", names=", ".join(checkpoints["bert_medium"])
+        type="Medium", names=", ".join(compatible_checkpoints("BertMedium"))
     ),
 )
 
@@ -810,13 +779,13 @@ setattr(
     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
index 00000000..874c09cd
--- /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
index 00000000..c7fe613a
--- /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
index 2c0694fc..79d598cf 100644
--- a/keras_nlp/models/bert_test.py
+++ b/keras_nlp/models/bert_test.py
@@ -106,10 +106,6 @@ class BertTest(tf.test.TestCase, parameterized.TestCase):
             }
             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 @@ class BertTest(tf.test.TestCase, parameterized.TestCase):
         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 @@ class BertTest(tf.test.TestCase, parameterized.TestCase):
         # 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 @@ class BertTest(tf.test.TestCase, parameterized.TestCase):
         # 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
index 796d0f66..77f3ca87 100644
--- 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.")