Repository: keras-team/keras-nlp
Pull request #360 (2022-09-15 15:46:12)

Use model.fit() for BERT Example

mattdangerw: Just FYI for anyone coming along this, @abheesht17 and I are still trying to figure out if what we want is actually working in Keras, not ready for review quite yet.

@abheesht17 one thing you can do for an unfinished PR like this is open them as a "draft" like https://github.com/keras-team/keras-nlp/pull/346.

Patch:
diff --git a/examples/bert/bert_train.py b/examples/bert/bert_train.py
index 9e870440..328841d6 100644
--- a/examples/bert/bert_train.py
+++ b/examples/bert/bert_train.py
@@ -202,19 +202,12 @@ class BertPretrainingModel(keras.Model):
         self.masked_lm_head = MaskedLMHead(
             embedding_table=encoder.token_embedding.embeddings,
             initializer=keras.initializers.TruncatedNormal(stddev=0.02),
+            name="mlm_layer",
         )
         self.next_sentence_head = keras.layers.Dense(
             encoder.num_segments,
             kernel_initializer=keras.initializers.TruncatedNormal(stddev=0.02),
-        )
-        self.loss_tracker = keras.metrics.Mean(name="loss")
-        self.lm_loss_tracker = keras.metrics.Mean(name="lm_loss")
-        self.nsp_loss_tracker = keras.metrics.Mean(name="nsp_loss")
-        self.lm_accuracy = keras.metrics.SparseCategoricalAccuracy(
-            name="lm_accuracy"
-        )
-        self.nsp_accuracy = keras.metrics.SparseCategoricalAccuracy(
-            name="nsp_accuracy"
+            name="nsp_layer",
         )
 
     def call(self, data):
@@ -233,40 +226,7 @@ class BertPretrainingModel(keras.Model):
             sequence_output, data["masked_lm_positions"]
         )
         nsp_preds = self.next_sentence_head(pooled_output)
-        return lm_preds, nsp_preds
-
-    def train_step(self, data):
-        with tf.GradientTape() as tape:
-            lm_preds, nsp_preds = self(data, training=True)
-            lm_labels = data["masked_lm_ids"]
-            lm_weights = data["masked_lm_weights"]
-            nsp_labels = data["next_sentence_labels"]
-
-            lm_loss = keras.losses.sparse_categorical_crossentropy(
-                lm_labels, lm_preds, from_logits=True
-            )
-            lm_weights_summed = tf.reduce_sum(lm_weights, -1)
-            lm_loss = tf.reduce_sum(lm_loss * lm_weights, -1)
-            lm_loss = tf.math.divide_no_nan(lm_loss, lm_weights_summed)
-            nsp_loss = keras.losses.sparse_categorical_crossentropy(
-                nsp_labels, nsp_preds, from_logits=True
-            )
-            nsp_loss = tf.reduce_mean(nsp_loss)
-            loss = lm_loss + nsp_loss
-
-        # Compute gradients
-        trainable_vars = self.trainable_variables
-        gradients = tape.gradient(loss, trainable_vars)
-        # Update weights
-        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
-
-        # Update metrics
-        self.loss_tracker.update_state(loss)
-        self.lm_loss_tracker.update_state(lm_loss)
-        self.nsp_loss_tracker.update_state(nsp_loss)
-        self.lm_accuracy.update_state(lm_labels, lm_preds, lm_weights)
-        self.nsp_accuracy.update_state(nsp_labels, nsp_preds)
-        return {m.name: m.result() for m in self.metrics}
+        return {"mlm": lm_preds, "nsp": nsp_preds}
 
 
 class LinearDecayWithWarmup(keras.optimizers.schedules.LearningRateSchedule):
@@ -328,7 +288,20 @@ def decode_record(record):
         if value.dtype == tf.int64:
             value = tf.cast(value, tf.int32)
         example[name] = value
-    return example
+
+    inputs = {
+        "token_ids": example["token_ids"],
+        "padding_mask": example["padding_mask"],
+        "segment_ids": example["segment_ids"],
+        "masked_lm_positions": example["masked_lm_positions"],
+    }
+    labels = {
+        "mlm": example["masked_lm_ids"],
+        "nsp": example["next_sentence_labels"],
+    }
+    sample_weights = {"mlm": example["masked_lm_weights"], "nsp": tf.ones((1,))}
+    sample = (inputs, labels, sample_weights)
+    return sample
 
 
 def get_checkpoint_callback():
@@ -429,9 +402,23 @@ def main(_):
         )
         optimizer = keras.optimizers.Adam(learning_rate=learning_rate_schedule)
 
+        lm_loss = keras.losses.SparseCategoricalCrossentropy(
+            from_logits=True,
+            name="lm_loss",
+        )
+        nsp_loss = keras.losses.SparseCategoricalCrossentropy(
+            from_logits=True,
+            name="nsp_loss",
+        )
+
+        lm_accuracy = keras.metrics.SparseCategoricalAccuracy(name="accuracy")
+        nsp_accuracy = keras.metrics.SparseCategoricalAccuracy(name="accuracy")
+
         pretraining_model = BertPretrainingModel(encoder)
         pretraining_model.compile(
             optimizer=optimizer,
+            loss={"mlm": lm_loss, "nsp": nsp_loss},
+            weighted_metrics={"mlm": lm_accuracy, "nsp": nsp_accuracy},
         )
 
     epochs = TRAINING_CONFIG["epochs"]
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index 39cbf53a..b33a79da 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -13,13 +13,18 @@
 # 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
+from keras_nlp.models.gpt2 import Gpt2Large
+from keras_nlp.models.gpt2 import Gpt2Medium
 from keras_nlp.models.roberta import RobertaBase
 from keras_nlp.models.roberta import RobertaClassifier
 from keras_nlp.models.roberta import RobertaCustom
diff --git a/keras_nlp/models/bert.py b/keras_nlp/models/bert.py
index b4f1a1ee..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,95 +26,113 @@ 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": {
+        "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_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_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_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": {
-        "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_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
+# We need the vocabulary_size hardcoded so we can instantiate a BERT network
 # with the right embedding size without downloading the matching vocabulary.
 # TODO(mattdangerw): Update our bucket structure so the vocabularies are
 # stored in an independent way, rather than reading from the base model.
@@ -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"],
         )
 
@@ -229,26 +251,26 @@ class BertCustom(keras.Model):
     or classification task networks.
 
     This class gives a fully customizable BERT model with any number of layers,
-    heads, and embedding dimensions. For specific specific BERT architectures
-    defined in the paper, see for example `keras_nlp.models.BertBase`.
+    heads, and embedding dimensions. For specific BERT architectures defined in
+    the paper, see, for example, `keras_nlp.models.BertBase`.
 
     Args:
-        vocabulary_size: Int. The size of the token vocabulary.
-        num_layers: Int. The number of transformer layers.
-        num_heads: Int. The number of attention heads for each transformer.
+        vocabulary_size: int. The size of the token vocabulary.
+        num_layers: int. The number of transformer layers.
+        num_heads: int. The number of attention heads for each transformer.
             The hidden size must be divisible by the number of attention heads.
-        hidden_dim: Int. The size of the transformer encoding and pooler layers.
-        intermediate_dim: Int. The output dimension of the first Dense layer in
+        hidden_dim: int. The size of the transformer encoding and pooler layers.
+        intermediate_dim: int. The output dimension of the first Dense layer in
             a two-layer feedforward network for each transformer.
-        dropout: Float. Dropout probability for the Transformer encoder.
-        max_sequence_length: Int. The maximum sequence length that this encoder
+        dropout: float. Dropout probability for the Transformer encoder.
+        max_sequence_length: int. The maximum sequence length that this encoder
             can consume. If None, `max_sequence_length` uses the value from
             sequence length. This determines the variable shape for positional
             embeddings.
-        num_segments: Int. The number of types that the 'segment_ids' input can
+        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
+        name: string, optional. Name of the model.
+        trainable: boolean, optional. If the model's variables should be
             trainable.
 
     Examples:
@@ -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,68 +583,6 @@ 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.
 
     This network implements a bi-directional Transformer-based encoder as
@@ -632,14 +592,14 @@ MODEL_DOCSTRING = """Bert "{type}" architecture.
     or classification task networks.
 
     Args:
-        weights: String, optional. Name of pretrained model to load weights.
+        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 `vocabularly_size` must be specified, but not both.
-        name: String, optional. Name of the model.
-        trainable: Boolean, optional. If the model's variables should be
+            `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.
 
     Examples:
@@ -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())),
+)
diff --git a/keras_nlp/models/gpt2.py b/keras_nlp/models/gpt2.py
new file mode 100644
index 00000000..1dc4b8ac
--- /dev/null
+++ b/keras_nlp/models/gpt2.py
@@ -0,0 +1,286 @@
+# 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.
+
+"""GPT-2 model configurable class, preconfigured versions, and task heads."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers import PositionEmbedding
+from keras_nlp.layers import TransformerDecoder
+
+
+def _gpt_2_kernel_initializer(stddev=0.02):
+    return keras.initializers.RandomNormal(stddev=stddev)
+
+
+class Gpt2Custom(keras.Model):
+    """GPT-2 core network with customizable hyperparameters.
+
+    This network implements a Transformer-based decoder network,
+    Generative Pretrained Transformer-2 (GPT-2), as described in
+    ["Language Models are Unsupervised Multitask Learners"](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf).
+    It includes the embedding lookups and transformer layers.
+
+    This class gives a fully customizable GPT-2 model with any number of layers,
+    heads, and embedding dimensions. For specific GPT-2 architectures
+    defined in the paper, see, for example, `keras_nlp.models.Gpt2Base`.
+
+    Args:
+        vocabulary_size: int. The size of the token vocabulary.
+        num_layers: int. The number of transformer layers.
+        num_heads: int. The number of attention heads for each transformer.
+            The hidden size must be divisible by the number of attention heads.
+        hidden_dim: int. The size of the transformer encoding and pooler layers.
+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        dropout: float. Dropout probability for the Transformer encoder.
+        max_sequence_length: int. The maximum sequence length that this encoder
+            can consume. If None, `max_sequence_length` uses the value from
+            sequence length. This determines the variable shape for positional
+            embeddings.
+        name: string, optional. Name of the model.
+        trainable: boolean, optional. If the model's variables should be
+            trainable.
+
+    Example usage:
+    ```python
+    # Randomly initialized GPT-2 decoder
+    model = keras_nlp.models.Gpt2Custom(
+        vocabulary_size=50257,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=1024,
+        name="encoder",
+    )
+
+    # Call encoder on the inputs
+    input_data = {
+        "token_ids": tf.random.uniform(
+            shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
+        ),
+        "padding_mask": tf.constant(
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
+    }
+    output = model(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        num_layers,
+        num_heads,
+        hidden_dim,
+        intermediate_dim,
+        dropout=0.1,
+        max_sequence_length=1024,
+        name=None,
+        trainable=True,
+    ):
+
+        # Inputs
+        token_ids = keras.Input(shape=(None,), dtype="int32", name="token_ids")
+        padding_mask = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+
+        # Embed tokens, positions.
+        token_embedding = keras.layers.Embedding(
+            input_dim=vocabulary_size,
+            output_dim=hidden_dim,
+            embeddings_initializer=_gpt_2_kernel_initializer(stddev=0.01),
+            name="token_embedding",
+        )(token_ids)
+
+        # Can't use `TokenAndPositionEmbedding` layer here because of different
+        # initializers.
+        position_embedding = PositionEmbedding(
+            initializer=_gpt_2_kernel_initializer(stddev=0.02),
+            sequence_length=max_sequence_length,
+            name="position_embedding",
+        )(token_embedding)
+
+        # Sum and apply dropout to embeddings.
+        x = keras.layers.Add()((token_embedding, position_embedding))
+        x = keras.layers.Dropout(
+            dropout,
+            name="embeddings_dropout",
+        )(x)
+
+        # Apply successive transformer decoder blocks.
+        for i in range(num_layers):
+            x = TransformerDecoder(
+                intermediate_dim=intermediate_dim,
+                num_heads=num_heads,
+                dropout=dropout,
+                activation=lambda x: keras.activations.gelu(
+                    x, approximate=True
+                ),
+                layer_norm_epsilon=1e-05,
+                kernel_initializer=_gpt_2_kernel_initializer(stddev=0.02),
+                normalize_first=True,
+                name=f"transformer_layer_{i}",
+            )(x, decoder_padding_mask=padding_mask)
+
+        sequence_output = keras.layers.LayerNormalization(
+            name="layer_norm",
+            axis=-1,
+            epsilon=1e-05,
+            dtype=tf.float32,
+        )(x)
+
+        # Instantiate using Functional API Model constructor
+        super().__init__(
+            inputs={
+                "token_ids": token_ids,
+                "padding_mask": padding_mask,
+            },
+            outputs=sequence_output,
+            name=name,
+            trainable=trainable,
+        )
+        # All references to `self` below this line
+        self.vocabulary_size = vocabulary_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.dropout = dropout
+        self.max_sequence_length = max_sequence_length
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "num_layers": self.num_layers,
+                "num_heads": self.num_heads,
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "dropout": self.dropout,
+                "max_sequence_length": self.max_sequence_length,
+            }
+        )
+        return config
+
+
+MODEL_DOCSTRING = """GPT-2 "{type}" architecture.
+
+    This network implements a Transformer-based decoder as
+    described in
+    ["Language Models are Unsupervised Multitask Learners"](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf).
+    It includes the embedding lookups and transformer layers.
+
+    Args:
+        vocabulary_size: int, optional. The size of the token vocabulary.
+        name: String, optional. Name of the model.
+        trainable: boolean, optional. If the model's variables should be
+            trainable.
+
+    Example usage:
+    ```python
+    # Randomly initialized Gpt2{type} encoder
+    model = keras_nlp.models.Gpt2{type}(vocabulary_size=10000)
+
+    # Call encoder on the inputs.
+    input_data = {{
+        "token_ids": tf.random.uniform(
+            shape=(1, 1024), dtype=tf.int64, maxval=model.vocabulary_size
+        ),
+        "padding_mask": tf.constant([1] * 1024, shape=(1, 1024)),
+    }}
+    output = model(input_data)
+"""
+
+
+def Gpt2Base(vocabulary_size, name=None, trainable=True):
+    return Gpt2Custom(
+        vocabulary_size=vocabulary_size,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        dropout=0.1,
+        max_sequence_length=1024,
+        name=name,
+        trainable=trainable,
+    )
+
+
+def Gpt2Medium(vocabulary_size, name=None, trainable=True):
+    return Gpt2Custom(
+        vocabulary_size=vocabulary_size,
+        num_layers=24,
+        num_heads=16,
+        hidden_dim=1024,
+        intermediate_dim=4096,
+        dropout=0.1,
+        max_sequence_length=1024,
+        name=name,
+        trainable=trainable,
+    )
+
+
+def Gpt2Large(vocabulary_size, name=None, trainable=True):
+    return Gpt2Custom(
+        vocabulary_size=vocabulary_size,
+        num_layers=36,
+        num_heads=20,
+        hidden_dim=1280,
+        intermediate_dim=5120,
+        dropout=0.1,
+        max_sequence_length=1024,
+        name=name,
+        trainable=trainable,
+    )
+
+
+def Gpt2ExtraLarge(vocabulary_size, name=None, trainable=True):
+    return Gpt2Custom(
+        vocabulary_size=vocabulary_size,
+        num_layers=48,
+        num_heads=25,
+        hidden_dim=1600,
+        intermediate_dim=6400,
+        dropout=0.1,
+        max_sequence_length=1024,
+        name=name,
+        trainable=trainable,
+    )
+
+
+setattr(
+    Gpt2Base,
+    "__doc__",
+    MODEL_DOCSTRING.format(type="Base"),
+)
+setattr(
+    Gpt2Medium,
+    "__doc__",
+    MODEL_DOCSTRING.format(type="Medium"),
+)
+setattr(
+    Gpt2Large,
+    "__doc__",
+    MODEL_DOCSTRING.format(type="Large"),
+)
+setattr(
+    Gpt2ExtraLarge,
+    "__doc__",
+    MODEL_DOCSTRING.format(type="ExtraLarge"),
+)
diff --git a/keras_nlp/models/roberta.py b/keras_nlp/models/roberta.py
index bbc6acc2..9d140f1e 100644
--- a/keras_nlp/models/roberta.py
+++ b/keras_nlp/models/roberta.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Roberta model configurable class, preconfigured versions, and task heads."""
+"""RoBERTa model configurable class, preconfigured versions, and task heads."""
 
 import tensorflow as tf
 from tensorflow import keras
@@ -30,32 +30,32 @@ class RobertaCustom(keras.Model):
 
     This network implements a bi-directional Transformer-based encoder as
     described in ["RoBERTa: A Robustly Optimized BERT Pretraining Approach"](https://arxiv.org/abs/1907.11692).
-    It includes the embedding lookups and transformer layers, but not include
-    the masked language model network.
+    It includes the embedding lookups and transformer layers, but does not
+    include the masked language modeling head used during pretraining.
 
-    This class gives a fully configurable Roberta model with any number of
-    layers, heads, and embedding dimensions. For specific roberta architectures
-    defined in the paper, see for example `keras_nlp.models.RobertaBase`.
+    This class gives a fully configurable RoBERTa model with any number of
+    layers, heads, and embedding dimensions. For specific RoBERTa architectures
+    defined in the paper, see, for example, `keras_nlp.models.RobertaBase`.
 
     Args:
-        vocabulary_size: Int. The size of the token vocabulary.
-        num_layers: Int. The number of transformer layers.
-        num_heads: Int. The number of attention heads for each transformer.
+        vocabulary_size: int. The size of the token vocabulary.
+        num_layers: int. The number of transformer layers.
+        num_heads: int. The number of attention heads for each transformer.
             The hidden size must be divisible by the number of attention heads.
-        hidden_dim: Int. The size of the transformer encoding layer.
-        intermediate_dim: Int. The output dimension of the first Dense layer in
+        hidden_dim: int. The size of the transformer encoding layer.
+        intermediate_dim: int. The output dimension of the first Dense layer in
             a two-layer feedforward network for each transformer.
-        dropout: Float. Dropout probability for the Transformer encoder.
-        max_sequence_length: Int. The maximum sequence length this encoder can
+        dropout: float. Dropout probability for the Transformer encoder.
+        max_sequence_length: int. The maximum sequence length this encoder can
             consume. The sequence length of the input must be less than
             `max_sequence_length`.
-        name: String, optional. Name of the model.
-        trainable: Boolean, optional. If the model's variables should be
+        name: string, optional. Name of the model.
+        trainable: boolean, optional. If the model's variables should be
             trainable.
 
     Example usage:
     ```python
-    # Randomly initialized Roberta model
+    # Randomly initialized RoBERTa model
     model = keras_nlp.models.RobertaCustom(
         vocabulary_size=50265,
         num_layers=12,
@@ -171,19 +171,19 @@ class RobertaCustom(keras.Model):
 
 
 class RobertaClassifier(keras.Model):
-    """Roberta encoder model with a classification head.
+    """RoBERTa encoder model with a classification head.
 
     Args:
         base_model: A `keras_nlp.models.Roberta` to encode inputs.
-        num_classes: Int. Number of classes to predict.
-        hidden_dim: Int. The size of the pooler layer.
-        name: String, optional. Name of the model.
-        trainable: Boolean, optional. If the model's variables should be
+        num_classes: int. Number of classes to predict.
+        hidden_dim: int. The size of the pooler layer.
+        name: string, optional. Name of the model.
+        trainable: boolean, optional. If the model's variables should be
             trainable.
 
     Example usage:
     ```python
-    # Randomly initialized Roberta encoder
+    # Randomly initialized RoBERTa encoder
     model = keras_nlp.models.RobertaCustom(
         vocabulary_size=50265,
         num_layers=12,
@@ -245,13 +245,13 @@ def RobertaBase(vocabulary_size, name=None, trainable=True):
     This network implements a bi-directional Transformer-based encoder as
     described in ["RoBERTa: A Robustly Optimized BERT Pretraining
     Approach"](https://arxiv.org/abs/1907.11692). It includes the
-    embedding lookups and transformer layers, but not the masked language model
-    or classification task networks.
+    embedding lookups and transformer layers, but does not include the masked
+    language modeling head used during pretraining.
 
     Args:
-        vocabulary_size: Int, optional. The size of the token vocabulary.
-        name: String, optional. Name of the model.
-        trainable: Boolean, optional. If the model's variables should be
+        vocabulary_size: int, optional. The size of the token vocabulary.
+        name: string, optional. Name of the model.
+        trainable: boolean, optional. If the model's variables should be
             trainable.
 
     Example usage:


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/models/gpt2_test.py b/keras_nlp/models/gpt2_test.py
new file mode 100644
index 00000000..fa54bdf0
--- /dev/null
+++ b/keras_nlp/models/gpt2_test.py
@@ -0,0 +1,92 @@
+# 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 GPT-2 model."""
+
+import os
+
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models import gpt2
+
+
+class Gpt2Test(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        self.model = gpt2.Gpt2Custom(
+            vocabulary_size=1000,
+            num_layers=2,
+            num_heads=2,
+            hidden_dim=64,
+            intermediate_dim=128,
+            max_sequence_length=128,
+            name="gpt2_test",
+        )
+        self.batch_size = 8
+        self.input_batch = {
+            "token_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"
+            ),
+        }
+
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
+    def test_valid_call_gpt2(self):
+        self.model(self.input_batch)
+
+    def test_variable_sequence_length_call_gpt2(self):
+        for seq_length in (25, 50, 75):
+            input_data = {
+                "token_ids": tf.ones(
+                    (self.batch_size, seq_length), dtype="int32"
+                ),
+                "padding_mask": tf.ones(
+                    (self.batch_size, seq_length), dtype="int32"
+                ),
+            }
+            self.model(input_data)
+
+    def test_valid_call_gpt2_base(self):
+        model = gpt2.Gpt2Base(vocabulary_size=1000, name="gpt2_base_test")
+        model(self.input_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_gpt2_base_compile(self, jit_compile):
+        model = gpt2.Gpt2Base(vocabulary_size=1000, name="gpt2_base_test")
+        model.compile(jit_compile=jit_compile)
+        model.predict(self.input_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_gpt2_base_compile_batched_ds(self, jit_compile):
+        model = gpt2.Gpt2Base(vocabulary_size=1000, name="gpt2_base_test")
+        model.compile(jit_compile=jit_compile)
+        model.predict(self.input_dataset)
+
+    def test_saving_model(self):
+        model_output = self.model(self.input_batch)
+        save_path = os.path.join(self.get_temp_dir(), "model")
+        self.model.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/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.")