Repository: keras-team/keras-nlp

Issue:
Unify modeling API around one class per model type
# Problem

Our current modeling API avoids creating an explicit config system by using config-in-code classes. These call a generic constructor with hardcoded arguments. In practice this requires us to maintain a bloated set of symbols in order to load a wide variety of checkpoints. If a user wants to add their own checkpoint with a different configuration they first need to create another config-in-code class.

```python
# The powerful API that we don"t use directly
# Does not allow checkpoint loading
keras_nlp.models.BertCustom(
    vocabulary_size=30522,
    num_layers=12,
    num_heads=12,     
    hidden_dim=768,
    intermediate_dim=3072,
    max_sequence_length=12,
    name="encoder",
)
# The API used by the rest of the library is a thin wrapper for the call above
keras_nlp.models.BertBase(weights="bert_base_uncased_en")
```

While this ostensibly avoids a config system, we reference this these thin wrappers in checkpoint metadata to essentially inline the config:
```python
checkpoints = {
    "bert_base_uncased_en": {
        "model": "BertBase",          # THIS IS THE NAME OF A HARDCODED CONFIG
        "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",
    },
}
```

The result is a loss of cohesion and flexibility in the library. `BertCustom` is clearly the powerful API but we all but hide it from users and don"t make it easy to load our checkpoints or make their own checkpoints with it. All configs must be anticipated by the package authors and the users have to navigate a large number of classes for example:
```python
# Why different classes? Why should the user care?
BertBase("bert_base_uncased_en")
BertLarge("bert_large_uncased_en")
```

This gets especially silly when we think of the `BertClassifier` model. We want to offer checkpoints for sentiment classifiers (#375), but in order to avoid providing configuration information we would have to create a config-in-code class for every possible number of classes!

```python
classifier_checkpoints = {
    "toxic_bert": {
        "model": "BertClassifierWith6Classes",          # THIS IS SILLY
        "backbone_model": "BertBase"
        "vocabulary": "uncased_en",
        "description": (
            "A multi-headed model capable of detecting different types "
            "of toxicity like threats, obscenity, insults, and identity-based hate."
        ),
        "weights_url": "https://storage.googleapis.com/keras-nlp/models/toxic_bert.h5",
        "weights_hash": "9b2b2139f221988759ac9cdd17050b31",
    },
}
# WHY!?!?!?
BertClassifierWith6Classes(weights="toxic_bert")
```


# Proposal

We can unify the API around one flexible, cohesive symbol using the native `keras.Model` methods of `from_config` and `to_config`.  With this functionality we do not need a maintain a parallel set of config classes. The config is just the constructor signature in most cases, and advanced functionality can be added by explicitly overriding `from_config`.

Now the same symbol can be called by all users:
```python
# Random graph
keras_nlp.models.Bert(
    vocabulary_size=30522,
    num_layers=12,
    num_heads=12,     
    hidden_dim=768,
    intermediate_dim=3072,
    max_sequence_length=12,
    name="encoder"
)
# Pretrained graph
keras_nlp.models.Bert.from_checkpoint("bert_base_uncased_en")
```

The checkpoint structs will now have the native Keras config for the model:
```python
#===> bert.py:
checkpoints = {
    "bert_base_uncased_en": {
        "config": {
            "vocabulary_size": 10000, 
            "hidden_dim": 768, 
            "intermediate_dim": 3072, 
            "num_layers": 12, 
            "num_heads": 12, 
            "max_sequence_length": 512, 
            "num_segments": 2, 
            "dropout": 0.1
        },
        "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_tasks.py
classifier_checkpoints = {
    "toxic_bert": {
        "config": {
            "backbone_config": {
                "vocabulary_size": 10000, 
                "hidden_dim": 768, 
                "intermediate_dim": 3072, 
                "num_layers": 12, 
                "num_heads": 12, 
                "max_sequence_length": 512, 
                "num_segments": 2, 
                "dropout": 0.1
        },
            "num_classes": 6,
        },
        "description": (
            "A multi-headed model capable of detecting different types "
            "of toxicity like threats, obscenity, insults, and identity-based hate."
        ),
        "weights_url": "https://storage.googleapis.com/keras-nlp/models/toxic_bert.h5",
        "weights_hash": "9b2b2139f221988759ac9cdd17050b31",
    },
}
```

The main difference between `Bert` and the old `BertCustom` API is the ability to load checkpoints using a `classmethod`:

```python
class Bert(keras.Model):
    ...
    @classmethod
    def from_checkpoint(checkpoint="bert_base_uncased_en") :
        model = Bert.from_config(checkpoints[name]["config"])
        model.load_weights(...)
...
```

Here `BertClassifier` constructor will work similarly, although we will have to implement a custom `from_config` to handle the nested structure:
```python
class BertClassifier(keras.Model):
    ...
    @classmethod
    def from_checkpoint(checkpoint="toxic_bert") :
        model = BertClassifier.from_config(classifier_checkpoints[name]["config"])
        model.load_weights(...)
...
```

# Architecture aliases
Since model architectures generally conform to a small number of templates, we can maintain a struct of hardcoded configs instead of hardcoded classes. In this setup the `checkpoints` struct would look identical but we would maintain another struct of architectures:

```python
architectures = {
    "bert_base" : {
        "vocabulary_size": 10000, 
        "hidden_dim": 768, 
        "intermediate_dim": 3072, 
        "num_layers": 12, 
        "num_heads": 12, 
        "max_sequence_length": 512, 
        "num_segments": 2, 
        "dropout": 0.1
    },
}
```

The checkpoint structs are now significantly cleaner:
```python
#===> bert.py:
checkpoints = {
    "bert_base_uncased_en": {
        "config": "bert_base",
        "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_tasks.py
classifier_checkpoints = {
    "toxic_bert": {
        "config": {
            "backbone_config": "bert_base"
            "num_classes": 6,
        },
        "description": (
            "A multi-headed model capable of detecting different types "
            "of toxicity like threats, obscenity, insults, and identity-based hate."
        ),
        "weights_url": "https://storage.googleapis.com/keras-nlp/models/toxic_bert.h5",
        "weights_hash": "9b2b2139f221988759ac9cdd17050b31",
    },
}
```

This would also allow us to set up convenience constructors for random graphs such as
```python
Bert.from_architecture(
    architecture="bert_base"
    vocabulary_size=10000
)
```

# User supplied weights
Since many power users will not want to rely on our GCP bucket to load weights, we should provide a `weights` argument that overrides our RCP weight loading but still uses our `checkpoints` config for all other information. If the supplied weights are incompatible with the constructed graph we can throw a `ValueError`.

```python
# User supplied weights for entire checkpoint---no RCPs involved!
keras_nlp.models.Bert.from_checkpoint(
    checkpoint="bert_base_uncased_en",
    weights="./my_weights.h5",
)
# User supplied weights for arch constructor.
# Interesting for consistency but less important since no RCPs to avoid
keras_nlp.models.Bert.from_architecture(    
    architecture="bert_base",
    vocabulary_size=10000,
    weights="./my_weights.h5",
)
# Equivalent to above, also for consistency
keras_nlp.models.Bert(
    weights="./my_weights.h5",
    vocabulary_size=30522,
    num_layers=12,
    num_heads=12,     
    hidden_dim=768,
    intermediate_dim=3072,
    max_sequence_length=12,
    name="encoder",
)
```

Note that in this design `weights` takes on a much more consistent meaning as "user-supplied weights" and does not allow for checkpoint or architecture overloading.

# API summary
Here are examples of the different API surfaces for interacting with a BERT model. Note that all can take a `weights` argument which pertains to a local file of compatible weights only.

```python
# Random encoder
keras_nlp.models.Bert(
    vocabulary_size=30522,
    num_layers=12,
    num_heads=12,     
    hidden_dim=768,
    intermediate_dim=3072,
    max_sequence_length=12,
)
# Random encoder from architecture alias
keras_nlp.models.Bert.from_architecture(
    architecture="bert_base",
    vocabulary_size=30522,    
)
# Pretrained encoder checkpoint
keras_nlp.models.Bert.from_checkpoint(
    checkpoint="bert_base_uncased_en",
)
# Classifier for fine-tuning
keras_nlp.models.BertClassifier(
    backbone="bert_base_uncased_en",
    num_classes=6
)
# Pretrained classifier
keras_nlp.models.BertClassifier.from_checkpoint(
    checkpoint="toxic_bert",
)
```

# Drawbacks
The biggest drawback to this proposal is that we are deviating significantly from the `keras.applications` and `keras_cv` style by using `classmethod` constructors rather than having special handling in `__init__` if `weights` is not `None`. It is technically possible to create the functionality described above using `__init__` and `__new__` but it will lead to meaningfully more complicated code. We are also deviating by avoiding the hardcoded classes, but I hope that this is a usability win. 

Discussion:
jbischof: Thanks @mattdangerw. Addressed your comments.
jbischof: Thanks @mattdangerw! Will resolve the merge conflicts once you approve.

Patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index 06766d79..cb77fead 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -12,12 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from keras_nlp.models.bert.bert_models import BertBase
-from keras_nlp.models.bert.bert_models import BertCustom
-from keras_nlp.models.bert.bert_models import BertLarge
-from keras_nlp.models.bert.bert_models import BertMedium
-from keras_nlp.models.bert.bert_models import BertSmall
-from keras_nlp.models.bert.bert_models import BertTiny
+from keras_nlp.models.bert.bert_models import Bert
 from keras_nlp.models.bert.bert_preprocessing import BertPreprocessor
 from keras_nlp.models.bert.bert_tasks import BertClassifier
 from keras_nlp.models.distilbert.distilbert_models import DistilBertBase
diff --git a/keras_nlp/models/bert/bert_models.py b/keras_nlp/models/bert/bert_models.py
index 1586e271..0fb2b16b 100644
--- a/keras_nlp/models/bert/bert_models.py
+++ b/keras_nlp/models/bert/bert_models.py
@@ -14,6 +14,7 @@
 
 """BERT backbone models."""
 
+import copy
 import os
 
 import tensorflow as tf
@@ -21,69 +22,32 @@ from tensorflow import keras
 
 from keras_nlp.layers.position_embedding import PositionEmbedding
 from keras_nlp.layers.transformer_encoder import TransformerEncoder
-from keras_nlp.models.bert.bert_checkpoints import checkpoints
-from keras_nlp.models.bert.bert_checkpoints import compatible_checkpoints
-from keras_nlp.models.bert.bert_checkpoints import vocabularies
+from keras_nlp.models.bert.bert_presets import backbone_presets
 
 
 def bert_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
-def _handle_pretrained_model_arguments(bert_variant, weights, vocabulary_size):
-    """Look up pretrained defaults for model arguments.
-
-    This helper will validate the `weights` and `vocabulary_size` arguments, and
-    fully resolve them in the case we are loading pretrained weights.
-    """
-    if (vocabulary_size is None and weights is None) or (
-        vocabulary_size and weights
-    ):
-        raise ValueError(
-            "One of `vocabulary_size` or `weights` must be specified "
-            "(but not both). "
-            f"Received: weights={weights}, "
-            f"vocabulary_size={vocabulary_size}"
-        )
-
-    if weights:
-        arch_checkpoints = compatible_checkpoints(bert_variant)
-        if weights not in arch_checkpoints:
-            raise ValueError(
-                "`weights` must be one of "
-                f"""{", ".join(arch_checkpoints)}. """
-                f"Received: {weights}"
-            )
-        metadata = checkpoints[weights]
-        vocabulary = metadata["vocabulary"]
-        vocabulary_size = vocabularies[vocabulary]["vocabulary_size"]
-
-        # TODO(jbischof): consider changing format from `h5` to
-        # `tf.train.Checkpoint` once
-        # https://github.com/keras-team/keras/issues/16946 is resolved.
-        weights = keras.utils.get_file(
-            "model.h5",
-            metadata["weights_url"],
-            cache_subdir=os.path.join("models", weights),
-            file_hash=metadata["weights_hash"],
-        )
-
-    return weights, vocabulary_size
+# TODO(jbischof): refactor into util file
+class classproperty(property):
+    def __get__(self, _, owner_cls):
+        return self.fget(owner_cls)
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class BertCustom(keras.Model):
-    """BERT encoder network with custom hyperparmeters.
+class Bert(keras.Model):
+    """BERT encoder network.
 
-    This network implements a bi-directional Transformer-based encoder as
+    This class implements a bi-directional Transformer-based encoder as
     described in ["BERT: Pre-training of Deep Bidirectional Transformers for
     Language Understanding"](https://arxiv.org/abs/1810.04805). It includes the
     embedding lookups and transformer layers, but not the masked language model
-    or classification task networks.
+    or next sentence prediction heads.
 
-    This class gives a fully customizable BERT model with any number of layers,
-    heads, and embedding dimensions. For specific BERT architectures defined in
-    the paper, see, for example, `keras_nlp.models.BertBase`.
+    The default constructor gives a fully customizable, randomly initalized BERT
+    encoder with any number of layers, heads, and embedding dimensions. To load
+    preset architectures and weights, use the `from_presets` constructor.
 
     Args:
         vocabulary_size: int. The size of the token vocabulary.
@@ -100,24 +64,9 @@ class BertCustom(keras.Model):
             embeddings.
         num_segments: int. The number of types that the 'segment_ids' input can
             take.
-        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,
-        name="encoder",
-    )
-
-    # Call encoder on the inputs
     input_data = {
         "token_ids": tf.random.uniform(
             shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
@@ -129,13 +78,21 @@ class BertCustom(keras.Model):
             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
         ),
     }
+
+    # Randomly initialized BERT encoder
+    model = keras_nlp.models.Bert(
+        vocabulary_size=30522,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=12,
+        name="encoder",
+    )
     output = model(input_data)
     ```
     """
 
-    # TODO(jbischof): consider changing `intermediate_dim` and `hidden_dim` to
-    # less confusing name here and in TransformerEncoder (`feed_forward_dim`?)
-
     def __init__(
         self,
         vocabulary_size,
@@ -146,8 +103,7 @@ class BertCustom(keras.Model):
         dropout=0.1,
         max_sequence_length=512,
         num_segments=2,
-        name=None,
-        trainable=True,
+        **kwargs,
     ):
 
         # Index of classification token in the vocabulary
@@ -232,11 +188,9 @@ class BertCustom(keras.Model):
                 "sequence_output": sequence_output,
                 "pooled_output": pooled_output,
             },
-            name=name,
-            trainable=trainable,
+            **kwargs,
         )
         # All references to `self` below this line
-        self.token_embedding = token_embedding_layer
         self.vocabulary_size = vocabulary_size
         self.hidden_dim = hidden_dim
         self.intermediate_dim = intermediate_dim
@@ -244,8 +198,8 @@ class BertCustom(keras.Model):
         self.num_heads = num_heads
         self.max_sequence_length = max_sequence_length
         self.num_segments = num_segments
-        self.intermediate_dim = intermediate_dim
         self.dropout = dropout
+        self.token_embedding = token_embedding_layer
         self.cls_token_index = cls_token_index
 
     def get_config(self):
@@ -266,210 +220,74 @@ class BertCustom(keras.Model):
     def from_config(cls, config):
         return cls(**config)
 
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy(backbone_presets)
 
-MODEL_DOCSTRING = """Bert "{type}" architecture.
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        load_weights=True,
+        **kwargs,
+    ):
 
-    This network implements a bi-directional Transformer-based encoder as
-    described in ["BERT: Pre-training of Deep Bidirectional Transformers for
-    Language Understanding"](https://arxiv.org/abs/1810.04805). It includes the
-    embedding lookups and transformer layers, but not the masked language model
-    or classification task networks.
+        if preset not in cls.presets:
+            raise ValueError(
+                "`preset` must be one of "
+                f"""{", ".join(cls.presets)}. Received: {preset}."""
+            )
+        metadata = cls.presets[preset]
+        config = metadata["config"]
+        model = cls.from_config({**config, **kwargs})
+
+        if not load_weights:
+            return model
+
+        weights = keras.utils.get_file(
+            "model.h5",
+            metadata["weights_url"],
+            cache_subdir=os.path.join("models", preset),
+            file_hash=metadata["weights_hash"],
+        )
+
+        model.load_weights(weights)
+        return model
+
+
+FROM_PRESET_DOCSTRING = """Instantiate BERT model from preset architecture and weights.
 
     Args:
-        weights: string, optional. Name of pretrained model to load weights.
-            Should be one of {names}.
-            If None, model is randomly initialized. Either `weights` or
-            `vocabulary_size` must be specified, but not both.
-        vocabulary_size: Int, optional. The size of the token vocabulary. Either
-            `weights` or `vocabulary_size` must be specified, but not both.
-        name: string, optional. Name of the model.
-        trainable: boolean, optional. If the model's variables should be
-            trainable.
+        preset: string. Must be one of {names}.
+        load_weights: Whether to load pre-trained weights into model. Defaults
+            to `True`.
 
     Examples:
     ```python
-    # Randomly initialized Bert{type} encoder
-    model = keras_nlp.models.Bert{type}(vocabulary_size=10000)
-
-    # Call encoder on the inputs.
     input_data = {{
         "token_ids": tf.random.uniform(
-            shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
+            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)
         ),
-        "segment_ids": tf.constant([0] * 200 + [1] * 312, shape=(1, 512)),
-        "padding_mask": tf.constant([1] * 512, shape=(1, 512)),
     }}
+
+    # Load architecture and weights from preset
+    model = Bert.from_preset("bert_base_uncased_en")
     output = model(input_data)
 
-    # Load a pretrained model
-    model = keras_nlp.models.Bert{type}(weights="bert_base_uncased_en")
-    # Call encoder on the inputs.
+    # Load randomly initalized model from preset architecture
+    model = Bert.from_preset("bert_base_uncased_en", load_weights=False)
     output = model(input_data)
     ```
-"""
-
-
-def BertTiny(weights=None, vocabulary_size=None, name=None, trainable=True):
-    weights, vocabulary_size = _handle_pretrained_model_arguments(
-        "BertTiny", weights, vocabulary_size
-    )
-
-    model = BertCustom(
-        vocabulary_size=vocabulary_size,
-        num_layers=2,
-        num_heads=2,
-        hidden_dim=128,
-        intermediate_dim=512,
-        dropout=0.1,
-        max_sequence_length=512,
-        name=name,
-        trainable=trainable,
-    )
-
-    if weights:
-        model.load_weights(weights)
-
-    return model
-
-
-def BertSmall(weights=None, vocabulary_size=None, name=None, trainable=True):
-    weights, vocabulary_size = _handle_pretrained_model_arguments(
-        "BertSmall", weights, vocabulary_size
-    )
-
-    model = BertCustom(
-        vocabulary_size=vocabulary_size,
-        num_layers=4,
-        num_heads=8,
-        hidden_dim=512,
-        intermediate_dim=2048,
-        dropout=0.1,
-        max_sequence_length=512,
-        name=name,
-        trainable=trainable,
-    )
-
-    if weights:
-        model.load_weights(weights)
-
-    return model
-
-
-def BertMedium(weights=None, vocabulary_size=None, name=None, trainable=True):
-    weights, vocabulary_size = _handle_pretrained_model_arguments(
-        "BertMedium", weights, vocabulary_size
-    )
-
-    model = BertCustom(
-        vocabulary_size=vocabulary_size,
-        num_layers=8,
-        num_heads=8,
-        hidden_dim=512,
-        intermediate_dim=2048,
-        dropout=0.1,
-        max_sequence_length=512,
-        name=name,
-        trainable=trainable,
-    )
-
-    if weights:
-        model.load_weights(weights)
-
-    return model
-
-
-def BertBase(weights=None, vocabulary_size=None, name=None, trainable=True):
-    weights, vocabulary_size = _handle_pretrained_model_arguments(
-        "BertBase", weights, vocabulary_size
-    )
-
-    model = BertCustom(
-        vocabulary_size=vocabulary_size,
-        num_layers=12,
-        num_heads=12,
-        hidden_dim=768,
-        intermediate_dim=3072,
-        dropout=0.1,
-        max_sequence_length=512,
-        name=name,
-        trainable=trainable,
-    )
-
-    if weights is not None:
-        model.load_weights(weights)
-
-    return model
-
-
-def BertLarge(weights=None, vocabulary_size=None, name=None, trainable=True):
-    weights, vocabulary_size = _handle_pretrained_model_arguments(
-        "BertLarge", weights, vocabulary_size
-    )
-
-    model = BertCustom(
-        vocabulary_size=vocabulary_size,
-        num_layers=24,
-        num_heads=16,
-        hidden_dim=1024,
-        intermediate_dim=4096,
-        dropout=0.1,
-        max_sequence_length=512,
-        name=name,
-        trainable=trainable,
-    )
-
-    if weights is not None:
-        model.load_weights(weights)
-
-    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(compatible_checkpoints("BertTiny"))
-    ),
-)
-
-setattr(
-    BertSmall,
-    "__doc__",
-    MODEL_DOCSTRING.format(
-        type="Small", names=", ".join(compatible_checkpoints("BertSmall"))
-    ),
-)
-
-setattr(
-    BertMedium,
-    "__doc__",
-    MODEL_DOCSTRING.format(
-        type="Medium", names=", ".join(compatible_checkpoints("BertMedium"))
-    ),
-)
+    """
 
 setattr(
-    BertBase,
-    "__doc__",
-    MODEL_DOCSTRING.format(
-        type="Base", names=", ".join(compatible_checkpoints("BertBase"))
-    ),
-)
-setattr(
-    BertLarge,
+    Bert.from_preset.__func__,
     "__doc__",
-    MODEL_DOCSTRING.format(
-        type="Large", names=", ".join(compatible_checkpoints("BertLarge"))
-    ),
+    FROM_PRESET_DOCSTRING.format(names=", ".join(backbone_presets)),
 )
diff --git a/keras_nlp/models/bert/bert_preprocessing.py b/keras_nlp/models/bert/bert_preprocessing.py
index a6765591..ae461bab 100644
--- a/keras_nlp/models/bert/bert_preprocessing.py
+++ b/keras_nlp/models/bert/bert_preprocessing.py
@@ -18,7 +18,7 @@ import os
 from tensorflow import keras
 
 from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
-from keras_nlp.models.bert.bert_checkpoints import vocabularies
+from keras_nlp.models.bert.bert_presets import vocabularies
 from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
 
 
@@ -115,6 +115,7 @@ ds = ds.map(
 """
 
 
+# TODO(jbischof): add `from_preset` classmethod
 class BertPreprocessor(keras.layers.Layer):
     def __init__(
         self,
diff --git a/keras_nlp/models/bert/bert_checkpoints.py b/keras_nlp/models/bert/bert_presets.py
similarity index 68%
rename from keras_nlp/models/bert/bert_checkpoints.py
rename to keras_nlp/models/bert/bert_presets.py
index 327c84ac..2f448cd5 100644
--- a/keras_nlp/models/bert/bert_checkpoints.py
+++ b/keras_nlp/models/bert/bert_presets.py
@@ -12,13 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections import defaultdict
-
-# TODO(jbischof): document checkpoints in keras.io and use URL in docstrings
+# TODO(jbischof): document presets in keras.io and use URL in docstrings
 # Metadata for loading pretrained model weights.
-checkpoints = {
+backbone_presets = {
     "bert_tiny_uncased_en": {
-        "model": "BertTiny",
+        "config": {
+            "vocabulary_size": 30522,
+            "num_layers": 2,
+            "num_heads": 2,
+            "hidden_dim": 128,
+            "intermediate_dim": 512,
+            "dropout": 0.1,
+            "max_sequence_length": 512,
+            "num_segments": 2,
+        },
         "vocabulary": "uncased_en",
         "description": (
             "Tiny size of BERT where all input is lowercased. "
@@ -28,7 +35,16 @@ checkpoints = {
         "weights_hash": "c2b29fcbf8f814a0812e4ab89ef5c068",
     },
     "bert_small_uncased_en": {
-        "model": "BertSmall",
+        "config": {
+            "vocabulary_size": 30522,
+            "num_layers": 4,
+            "num_heads": 8,
+            "hidden_dim": 512,
+            "intermediate_dim": 2048,
+            "dropout": 0.1,
+            "max_sequence_length": 512,
+            "num_segments": 2,
+        },
         "vocabulary": "uncased_en",
         "description": (
             "Small size of BERT where all input is lowercased. "
@@ -38,7 +54,16 @@ checkpoints = {
         "weights_hash": "08632c9479b034f342ba2c2b7afba5f7",
     },
     "bert_medium_uncased_en": {
-        "model": "BertMedium",
+        "config": {
+            "vocabulary_size": 30522,
+            "num_layers": 8,
+            "num_heads": 8,
+            "hidden_dim": 512,
+            "intermediate_dim": 2048,
+            "dropout": 0.1,
+            "max_sequence_length": 512,
+            "num_segments": 2,
+        },
         "vocabulary": "uncased_en",
         "description": (
             "Medium size of BERT where all input is lowercased. "
@@ -48,7 +73,16 @@ checkpoints = {
         "weights_hash": "bb990e1184ec6b6185450c73833cd661",
     },
     "bert_base_uncased_en": {
-        "model": "BertBase",
+        "config": {
+            "vocabulary_size": 30522,
+            "num_layers": 12,
+            "num_heads": 12,
+            "hidden_dim": 768,
+            "intermediate_dim": 3072,
+            "dropout": 0.1,
+            "max_sequence_length": 512,
+            "num_segments": 2,
+        },
         "vocabulary": "uncased_en",
         "description": (
             "Base size of BERT where all input is lowercased. "
@@ -58,7 +92,16 @@ checkpoints = {
         "weights_hash": "9b2b2139f221988759ac9cdd17050b31",
     },
     "bert_base_cased_en": {
-        "model": "BertBase",
+        "config": {
+            "vocabulary_size": 28996,
+            "num_layers": 12,
+            "num_heads": 12,
+            "hidden_dim": 768,
+            "intermediate_dim": 3072,
+            "dropout": 0.1,
+            "max_sequence_length": 512,
+            "num_segments": 2,
+        },
         "vocabulary": "cased_en",
         "description": (
             "Base size of BERT where case is maintained. "
@@ -68,21 +111,48 @@ checkpoints = {
         "weights_hash": "f94a6cb012e18f4fb8ec92abb91864e9",
     },
     "bert_base_zh": {
-        "model": "BertBase",
+        "config": {
+            "vocabulary_size": 21128,
+            "num_layers": 12,
+            "num_heads": 12,
+            "hidden_dim": 768,
+            "intermediate_dim": 3072,
+            "dropout": 0.1,
+            "max_sequence_length": 512,
+            "num_segments": 2,
+        },
         "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",
+        "config": {
+            "vocabulary_size": 119547,
+            "num_layers": 12,
+            "num_heads": 12,
+            "hidden_dim": 768,
+            "intermediate_dim": 3072,
+            "dropout": 0.1,
+            "max_sequence_length": 512,
+            "num_segments": 2,
+        },
         "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",
+        "config": {
+            "vocabulary_size": 30522,
+            "num_layers": 24,
+            "num_heads": 16,
+            "hidden_dim": 1024,
+            "intermediate_dim": 4096,
+            "dropout": 0.1,
+            "max_sequence_length": 512,
+            "num_segments": 2,
+        },
         "vocabulary": "uncased_en",
         "description": (
             "Large size of BERT where all input is lowercased. "
@@ -92,7 +162,16 @@ checkpoints = {
         "weights_hash": "cc5cacc9565ef400ee4376105f40ddae",
     },
     "bert_large_cased_en": {
-        "model": "BertLarge",
+        "config": {
+            "vocabulary_size": 28996,
+            "num_layers": 24,
+            "num_heads": 16,
+            "hidden_dim": 1024,
+            "intermediate_dim": 4096,
+            "dropout": 0.1,
+            "max_sequence_length": 512,
+            "num_segments": 2,
+        },
         "vocabulary": "cased_en",
         "description": (
             "Base size of BERT where case is maintained. "
@@ -150,13 +229,3 @@ vocabularies = {
         "lowercase": False,
     },
 }
-
-# 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]
diff --git a/keras_nlp/models/bert/bert_tasks.py b/keras_nlp/models/bert/bert_tasks.py
index e1f1e570..fd1c3f53 100644
--- a/keras_nlp/models/bert/bert_tasks.py
+++ b/keras_nlp/models/bert/bert_tasks.py
@@ -15,17 +15,15 @@
 
 from tensorflow import keras
 
-from keras_nlp.models.bert.bert_checkpoints import checkpoints
+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_models import model_class_by_name
 
 # 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}.
+        backbone: A string or `keras_nlp.models.Bert` instance. 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
@@ -80,15 +78,12 @@ class BertClassifier(keras.Model):
         # 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 checkpoints:
+            if backbone not in Bert.presets:
                 raise ValueError(
                     "`backbone` must be one of "
-                    f"""{", ".join(checkpoints.keys())}. """
-                    f"Received: {backbone}"
+                    f"""{", ".join(Bert.presets)}. Received: {backbone}."""
                 )
-            backbone_class_str = checkpoints[backbone]["model"]
-            backbone_class = model_class_by_name(backbone_class_str)
-            backbone = backbone_class(backbone)
+            backbone = Bert.from_preset(backbone)
 
         inputs = backbone.input
         pooled = backbone(inputs)["pooled_output"]
@@ -123,5 +118,5 @@ class BertClassifier(keras.Model):
 setattr(
     BertClassifier,
     "__doc__",
-    CLASSIFIER_DOCSTRING.format(names=", ".join(checkpoints.keys())),
+    CLASSIFIER_DOCSTRING.format(names=", ".join(Bert.presets)),
 )


Test patch:
diff --git a/keras_nlp/models/bert/bert_models_test.py b/keras_nlp/models/bert/bert_models_test.py
index 740d225f..803d4b7b 100644
--- a/keras_nlp/models/bert/bert_models_test.py
+++ b/keras_nlp/models/bert/bert_models_test.py
@@ -19,13 +19,12 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.bert.bert_models import BertBase
-from keras_nlp.models.bert.bert_models import BertCustom
+from keras_nlp.models.bert.bert_models import Bert
 
 
 class BertTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = BertCustom(
+        self.model = Bert(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
@@ -53,6 +52,7 @@ class BertTest(tf.test.TestCase, parameterized.TestCase):
 
     def test_valid_call_bert(self):
         self.model(self.input_batch)
+        self.assertEqual(self.model.name, "encoder")
 
     def test_variable_sequence_length_call_bert(self):
         for seq_length in (25, 50, 75):
@@ -69,56 +69,68 @@ class BertTest(tf.test.TestCase, parameterized.TestCase):
             }
             self.model(input_data)
 
-    def test_valid_call_bert_base(self):
-        model = BertBase(vocabulary_size=1000, name="encoder")
-        input_data = {
-            "token_ids": tf.ones(
-                (self.batch_size, self.model.max_sequence_length), dtype="int32"
-            ),
-            "segment_ids": tf.ones(
-                (self.batch_size, self.model.max_sequence_length), dtype="int32"
-            ),
-            "padding_mask": tf.ones(
-                (self.batch_size, self.model.max_sequence_length), dtype="int32"
-            ),
-        }
-        model(input_data)
+    def test_valid_call_presets(self):
+        # Test preset loading without weights
+        for preset in Bert.presets:
+            model = Bert.from_preset(preset, load_weights=False, name="encoder")
+            input_data = {
+                "token_ids": tf.ones(
+                    (self.batch_size, self.model.max_sequence_length),
+                    dtype="int32",
+                ),
+                "segment_ids": tf.ones(
+                    (self.batch_size, self.model.max_sequence_length),
+                    dtype="int32",
+                ),
+                "padding_mask": tf.ones(
+                    (self.batch_size, self.model.max_sequence_length),
+                    dtype="int32",
+                ),
+            }
+            model(input_data)
+
+    def test_unknown_preset_error(self):
+        # Not a preset name
+        with self.assertRaises(ValueError):
+            Bert.from_preset(
+                "bert_base_uncased_clowntown",
+                load_weights=False,
+                name="encoder",
+            )
+
+    def test_preset_mutability(self):
+        preset = "bert_base_uncased_en"
+        # Cannot overwrite the presents attribute in an object
+        with self.assertRaises(AttributeError):
+            self.model.presets = {"my_model": "clowntown"}
+        # Cannot mutate presents in an object
+        config = self.model.presets[preset]["config"]
+        config["max_sequence_length"] = 1
+        self.assertEqual(config["max_sequence_length"], 1)
+        self.assertEqual(
+            self.model.presets[preset]["config"]["max_sequence_length"], 512
+        )
+        # Cannot mutate presets in the class
+        config = Bert.presets[preset]["config"]
+        config["max_sequence_length"] = 1
+        self.assertEqual(config["max_sequence_length"], 1)
+        self.assertEqual(
+            Bert.presets[preset]["config"]["max_sequence_length"], 512
+        )
 
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
-    def test_bert_base_compile(self, jit_compile):
-        model = BertBase(vocabulary_size=1000, name="encoder")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_batch)
+    def test_compile(self, jit_compile):
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_batch)
 
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
-    def test_bert_base_compile_batched_ds(self, jit_compile):
-        model = BertBase(vocabulary_size=1000, name="encoder")
-        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):
-            BertBase(name="encoder")
-
-        # Only one of `vocabulary_size` or `weights`
-        with self.assertRaises(ValueError):
-            BertBase(
-                weights="bert_base_uncased_en",
-                vocabulary_size=1000,
-                name="encoder",
-            )
-
-        # Not a checkpoint name
-        with self.assertRaises(ValueError):
-            BertBase(
-                weights="bert_base_uncased_clowntown",
-                name="encoder",
-            )
+    def test_compile_batched_ds(self, jit_compile):
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_dataset)
 
     @parameterized.named_parameters(
         ("save_format_tf", "tf"), ("save_format_h5", "h5")
@@ -130,7 +142,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, BertCustom)
+        self.assertIsInstance(restored_model, Bert)
 
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
diff --git a/keras_nlp/models/bert/bert_tasks_test.py b/keras_nlp/models/bert/bert_tasks_test.py
index eca12ebc..0fe3abda 100644
--- a/keras_nlp/models/bert/bert_tasks_test.py
+++ b/keras_nlp/models/bert/bert_tasks_test.py
@@ -19,13 +19,13 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.bert.bert_models import BertCustom
+from keras_nlp.models.bert.bert_models import Bert
 from keras_nlp.models.bert.bert_tasks import BertClassifier
 
 
 class BertClassifierTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.backbone = BertCustom(
+        self.backbone = Bert(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
diff --git a/keras_nlp/tests/network_tests/test_load_ckpts.py b/keras_nlp/tests/network_tests/test_load_presets.py
similarity index 79%
rename from keras_nlp/tests/network_tests/test_load_ckpts.py
rename to keras_nlp/tests/network_tests/test_load_presets.py
index 76ff3b2d..2718fec3 100644
--- a/keras_nlp/tests/network_tests/test_load_ckpts.py
+++ b/keras_nlp/tests/network_tests/test_load_presets.py
@@ -11,24 +11,22 @@
 # 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 loading pretrained model checkpoints."""
+"""Tests for loading pretrained model presets."""
 
 import pytest
 import tensorflow as tf
 
-from keras_nlp.models.bert import bert_checkpoints
 from keras_nlp.models.bert import bert_models
 from keras_nlp.models.bert import bert_preprocessing
+from keras_nlp.models.bert import bert_presets
 from keras_nlp.models.bert import bert_tasks
 
 
 @pytest.mark.slow
-class BertCkptTest(tf.test.TestCase):
+class BertPresetTest(tf.test.TestCase):
     def test_load_bert(self):
-        for checkpoint in bert_checkpoints.checkpoints:
-            bert_variant = bert_checkpoints.checkpoints[checkpoint]["model"]
-            bert_class = bert_models.model_class_by_name(bert_variant)
-            model = bert_class(weights=checkpoint)
+        for preset in bert_models.Bert.presets:
+            model = bert_models.Bert.from_preset(preset, load_weights=True)
             input_data = {
                 "token_ids": tf.random.uniform(
                     shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
@@ -41,10 +39,8 @@ class BertCkptTest(tf.test.TestCase):
             model(input_data)
 
     def test_load_bert_backbone_string(self):
-        for checkpoint in bert_checkpoints.checkpoints:
-            classifier = bert_tasks.BertClassifier(
-                checkpoint, 4, name="classifier"
-            )
+        for preset in bert_models.Bert.presets:
+            classifier = bert_tasks.BertClassifier(preset, 4, name="classifier")
             input_data = {
                 "token_ids": tf.random.uniform(
                     shape=(1, 512),
@@ -72,7 +68,7 @@ class BertCkptTest(tf.test.TestCase):
         classifier(input_data)
 
     def test_load_vocabularies(self):
-        for vocabulary in bert_checkpoints.vocabularies:
+        for vocabulary in bert_presets.vocabularies:
             tokenizer = bert_preprocessing.BertPreprocessor(
                 vocabulary=vocabulary,
             )