Repository: keras-team/keras-nlp
Pull request #304 (2022-08-22 09:21:28)

RoBERTa pretrained model forward pass
RoBERTa model and heads implementation

jessechancy: @jbischof 

Patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index 703f09ca..f07d8f1d 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -15,3 +15,6 @@
 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.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 143e7618..9b210acc 100644
--- a/keras_nlp/models/bert.py
+++ b/keras_nlp/models/bert.py
@@ -25,6 +25,27 @@ def _bert_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
+# Pretrained models
+BASE_PATH = "https://storage.googleapis.com/keras-nlp/models/"
+
+checkpoints = {
+    "bert_base": {
+        "uncased_en": {
+            "md5": "9b2b2139f221988759ac9cdd17050b31",
+            "description": "Base size of Bert where all input is lowercased. "
+            "Trained on English wikipedia + books corpora.",
+            "vocabulary_size": 30522,
+        },
+        "cased_en": {
+            "md5": "f94a6cb012e18f4fb8ec92abb91864e9",
+            "description": "Base size of Bert where case is maintained. "
+            "Trained on English wikipedia + books corpora.",
+            "vocabulary_size": 28996,
+        },
+    }
+}
+
+
 class BertCustom(keras.Model):
     """Bi-directional Transformer-based encoder network.
 
@@ -61,25 +82,29 @@ class BertCustom(keras.Model):
     Example usage:
     ```python
     # Randomly initialized Bert encoder
-    encoder = keras_nlp.models.BertCustom(
+    model = keras_nlp.models.BertCustom(
         vocabulary_size=30522,
         num_layers=12,
         num_heads=12,
         hidden_dim=768,
         intermediate_dim=3072,
-        max_sequence_length=12
+        max_sequence_length=12,
+        name="encoder",
     )
 
-     # Call encoder on the inputs.
+    # Call encoder on the inputs
     input_data = {
         "input_ids": tf.random.uniform(
-            shape=(1, 12), dtype=tf.int64, maxval=30522),
+            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)),
+            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
         "input_mask": tf.constant(
-            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
     }
-    output = encoder(input_data)
+    output = model(input_data)
     ```
     """
 
@@ -156,7 +181,7 @@ class BertCustom(keras.Model):
                 ),
                 dropout=dropout,
                 kernel_initializer=_bert_kernel_initializer(),
-                name=f"""transformer_layer_{i}""",
+                name=f"transformer_layer_{i}",
             )(x, padding_mask=input_mask)
 
         # Construct the two Bert outputs. The pooled output is a dense layer on
@@ -225,10 +250,9 @@ class BertClassifier(keras.Model):
             trainable.
 
     Example usage:
-    ```
-    python
+    ```python
     # Randomly initialized Bert encoder
-    encoder = keras_nlp.models.BertCustom(
+    model = keras_nlp.models.BertCustom(
         vocabulary_size=30522,
         num_layers=12,
         num_heads=12,
@@ -237,16 +261,19 @@ class BertClassifier(keras.Model):
         max_sequence_length=12
     )
 
-     # Call classifier on the inputs.
+    # Call classifier on the inputs.
     input_data = {
         "input_ids": tf.random.uniform(
-            shape=(1, 12), dtype=tf.int64, maxval=30522),
+            shape=(1, 12), dtype=tf.int64, maxval=30522
+        ),
         "segment_ids": tf.constant(
-            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
         "input_mask": tf.constant(
-            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
     }
-    classifier = keras_nlp.models.BertClassifier(encoder, 4)
+    classifier = bert.BertClassifier(model, 4, name="classifier")
     logits = classifier(input_data)
     ```
     """
@@ -274,9 +301,8 @@ class BertClassifier(keras.Model):
         self.num_classes = num_classes
 
 
-def BertBase(name=None, trainable=True):
-    """Bi-directional Transformer-based encoder network (Bert) using "Base"
-    architecture.
+MODEL_DOCSTRING = """Bi-directional Transformer-based encoder network (Bert)
+    using "{type}" architecture.
 
     This network implements a bi-directional Transformer-based encoder as
     described in ["BERT: Pre-training of Deep Bidirectional Transformers for
@@ -285,6 +311,12 @@ def BertBase(name=None, trainable=True):
     or classification task networks.
 
     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 `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
             trainable.
@@ -292,23 +324,49 @@ def BertBase(name=None, trainable=True):
     Example usage:
     ```python
     # Randomly initialized BertBase encoder
-    encoder = keras_nlp.models.BertBase()
+    model = keras_nlp.models.BertBase(vocabulary_size=10000)
 
-     # Call encoder on the inputs.
-    input_data = {
+    # Call encoder on the inputs.
+    input_data = {{
         "input_ids": tf.random.uniform(
-            shape=(1, 512), dtype=tf.int64, maxval=encoder.vocabulary_size),
-        "segment_ids": tf.constant(
-            [0] * 200 + [1] * 312, shape=(1, 512)),
-        "input_mask": tf.constant(
-            [1] * 512, shape=(1, 512)),
-    }
-    output = encoder(input_data)
+            shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
+        ),
+        "segment_ids": tf.constant([0] * 200 + [1] * 312, shape=(1, 512)),
+        "input_mask": tf.constant([1] * 512, shape=(1, 512)),
+    }}
+    output = model(input_data)
+
+    # Load a pretrained model
+    model = keras_nlp.models.BertBase(weights="uncased_en")
+    # Call encoder on the inputs.
+    output = model(input_data)
     ```
-    """
+"""
+
+
+def BertBase(weights=None, vocabulary_size=None, name=None, trainable=True):
+
+    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:
+        if weights not in checkpoints["bert_base"]:
+            raise ValueError(
+                "`weights` must be one of "
+                f"""{", ".join(checkpoints["bert_base"])}. """
+                f"Received: {weights}"
+            )
+        vocabulary_size = checkpoints["bert_base"][weights]["vocabulary_size"]
 
     model = BertCustom(
-        vocabulary_size=30522,
+        vocabulary_size=vocabulary_size,
         num_layers=12,
         num_heads=12,
         hidden_dim=768,
@@ -319,6 +377,26 @@ def BertBase(name=None, trainable=True):
         trainable=trainable,
     )
 
-    # TODO(jbischof): add some documentation or magic to load our checkpoints
-    # TODO(jbischof): attach the tokenizer
+    # TODO(jbischof): consider changing format from `h5` to
+    # `tf.train.Checkpoint` once
+    # https://github.com/keras-team/keras/issues/16946 is resolved
+    if weights:
+        filepath = keras.utils.get_file(
+            "model.h5",
+            BASE_PATH + "bert_base_" + weights + "/model.h5",
+            cache_subdir="models/bert_base/" + weights + "/",
+            file_hash=checkpoints["bert_base"][weights]["md5"],
+        )
+        model.load_weights(filepath)
+
+    # TODO(jbischof): attach the tokenizer or create separate tokenizer class
     return model
+
+
+setattr(
+    BertBase,
+    "__doc__",
+    MODEL_DOCSTRING.format(
+        type="Base", names=", ".join(checkpoints["bert_base"])
+    ),
+)
diff --git a/keras_nlp/models/roberta.py b/keras_nlp/models/roberta.py
new file mode 100644
index 00000000..4813e79c
--- /dev/null
+++ b/keras_nlp/models/roberta.py
@@ -0,0 +1,276 @@
+# 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.
+
+"""Roberta model configurable class, preconfigured versions, and task heads."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers import TokenAndPositionEmbedding
+from keras_nlp.layers import TransformerEncoder
+
+START_TOKEN_INDEX = 0
+
+
+class RobertaCustom(keras.Model):
+    """RoBERTa model with a customizable set of hyperparameters.
+
+    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.
+
+    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.
+            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
+            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
+            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
+            trainable.
+
+    Example usage:
+    ```python
+    # Randomly initialized Roberta encoder
+    encoder = keras_nlp.models.RobertaCustom(
+        vocabulary_size=50265,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=12
+    )
+
+    # Call encoder on the inputs.
+    input_data = {
+        "input_ids": tf.random.uniform(
+            shape=(1, 12), dtype=tf.int64, maxval=50265),
+        "input_mask": tf.constant(
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+    }
+    output = encoder(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        num_layers,
+        num_heads,
+        hidden_dim,
+        intermediate_dim,
+        dropout=0.1,
+        max_sequence_length=512,
+        name=None,
+        trainable=True,
+    ):
+
+        token_id_input = keras.Input(
+            shape=(None,), dtype=tf.int32, name="input_ids"
+        )
+        input_mask = keras.Input(
+            shape=(None,), dtype=tf.int32, name="input_mask"
+        )
+
+        # Embed tokens and positions.
+        embedding_layer = TokenAndPositionEmbedding(
+            vocabulary_size=vocabulary_size,
+            sequence_length=max_sequence_length,
+            embedding_dim=hidden_dim,
+            embeddings_initializer=keras.initializers.TruncatedNormal(0.02),
+            name="embeddings",
+        )
+        embedding = embedding_layer(token_id_input)
+
+        # Sum, normalize and apply dropout to embeddings.
+        x = keras.layers.LayerNormalization(
+            name="embeddings_layer_norm",
+            axis=-1,
+            epsilon=1e-5,  # Original paper uses this epsilon value
+            dtype=tf.float32,
+        )(embedding)
+        x = keras.layers.Dropout(
+            dropout,
+            name="embeddings_dropout",
+        )(x)
+
+        # Apply successive transformer encoder blocks.
+        for i in range(num_layers):
+            x = TransformerEncoder(
+                num_heads=num_heads,
+                intermediate_dim=intermediate_dim,
+                activation="gelu",
+                dropout=dropout,
+                kernel_initializer=keras.initializers.TruncatedNormal(0.02),
+                name=f"transformer_layer_{i}",
+            )(x, padding_mask=input_mask)
+
+        # Instantiate using Functional API Model constructor
+        super().__init__(
+            inputs={
+                "input_ids": token_id_input,
+                "input_mask": input_mask,
+            },
+            outputs={
+                "sequence_output": x,
+            },
+            name=name,
+            trainable=trainable,
+        )
+        # All references to `self` below this line
+        self.vocabulary_size = vocabulary_size
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.max_sequence_length = max_sequence_length
+        self.intermediate_dim = intermediate_dim
+        self.dropout = dropout
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "num_layers": self.num_layers,
+                "num_heads": self.num_heads,
+                "max_sequence_length": self.max_sequence_length,
+                "dropout": self.dropout,
+            }
+        )
+        return config
+
+
+class RobertaClassifier(keras.Model):
+    """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
+            trainable.
+
+    Example usage:
+    ```python
+    # Randomly initialized Roberta encoder
+    encoder = keras_nlp.models.RobertaCustom(
+        vocabulary_size=50265,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=12
+    )
+
+    # Call classifier on the inputs.
+    input_data = {
+        "input_ids": tf.random.uniform(
+            shape=(1, 12), dtype=tf.int64, maxval=50265),
+        "input_mask": tf.constant(
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+    }
+    classifier = keras_nlp.models.RobertaClassifier(encoder, 4)
+    logits = classifier(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        base_model,
+        num_classes,
+        hidden_dim=None,
+        dropout=0.0,
+        name=None,
+        trainable=True,
+    ):
+        inputs = base_model.input
+        if hidden_dim is None:
+            hidden_dim = base_model.hidden_dim
+
+        x = base_model(inputs)["sequence_output"][:, START_TOKEN_INDEX, :]
+        x = keras.layers.Dropout(dropout, name="pooled_dropout")(x)
+        x = keras.layers.Dense(
+            hidden_dim, activation="tanh", name="pooled_dense"
+        )(x)
+        x = keras.layers.Dropout(dropout, name="classifier_dropout")(x)
+        outputs = keras.layers.Dense(
+            num_classes,
+            kernel_initializer=keras.initializers.TruncatedNormal(0.02),
+            name="logits",
+        )(x)
+
+        # 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
+
+
+def RobertaBase(name=None, trainable=True):
+    """RoBERTa implementation using "Base" architecture.
+
+    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.
+
+    Args:
+        name: String, optional. Name of the model.
+        trainable: Boolean, optional. If the model's variables should be
+            trainable.
+
+    Example usage:
+    ```python
+    # Randomly initialized RobertaBase encoder
+    encoder = keras_nlp.models.RobertaBase()
+
+    # Call encoder on the inputs.
+    input_data = {
+        "input_ids": tf.random.uniform(
+            shape=(1, 512), dtype=tf.int64, maxval=encoder.vocabulary_size),
+        "input_mask": tf.ones((1, 512)),
+    }
+    output = encoder(input_data)
+    ```
+    """
+
+    return RobertaCustom(
+        vocabulary_size=50265,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        dropout=0.1,
+        max_sequence_length=512,
+        name=name,
+        trainable=trainable,
+    )
diff --git a/setup.cfg b/setup.cfg
index 67f100fb..8a046364 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -10,8 +10,8 @@ filterwarnings =
 
 addopts=-vv
 
-# Do not run tests in the build folder
-norecursedirs= build
+# Do not run tests in the `build` or `network_tests` folders
+norecursedirs = build network_tests
 
 [isort]
 known_first_party = keras_nlp,tests
diff --git a/shell/format.sh b/shell/format.sh
index 2eee93f9..b8b1e3b5 100755
--- a/shell/format.sh
+++ b/shell/format.sh
@@ -1,7 +1,7 @@
 #!/bin/bash -e
 
 base_dir=$(dirname $(dirname $0))
-targets="${base_dir}/*.py ${base_dir}/examples/ ${base_dir}/keras_nlp/"
+targets="${base_dir}/*.py ${base_dir}/examples/ ${base_dir}/keras_nlp/ ${base_dir}/network_tests/"
 
 isort --sp "${base_dir}/setup.cfg" --sl ${targets}
 black --line-length 80 ${targets}
diff --git a/tools/checkpoint_conversion/bert_base_cased.ipynb b/tools/checkpoint_conversion/bert_base_cased.ipynb
new file mode 100644
index 00000000..ce805b77
--- /dev/null
+++ b/tools/checkpoint_conversion/bert_base_cased.ipynb
@@ -0,0 +1,994 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/jbischof/keras-nlp/blob/bert_ckpt/bert_base_cased.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vGp_yrJi5Ehf"
+      },
+      "source": [
+        "## Install deps"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Szd6xKUd2tIE",
+        "outputId": "564b86a2-a7ed-4e22-f1fa-4246368d30a7"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "\u001b[K     |████████████████████████████████| 511.7 MB 6.8 kB/s \n",
+            "\u001b[K     |████████████████████████████████| 2.1 MB 49.8 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 4.6 MB 45.7 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 5.8 MB 49.9 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 1.6 MB 59.9 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 438 kB 57.4 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 116 kB 48.3 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 352 kB 67.8 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 43 kB 2.1 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 99 kB 11.3 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 1.3 MB 51.7 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 238 kB 73.0 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 1.1 MB 60.3 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 636 kB 71.3 MB/s \n",
+            "\u001b[?25h  Building wheel for keras-nlp (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
+            "  Building wheel for py-cpuinfo (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
+            "  Building wheel for seqeval (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
+          ]
+        }
+      ],
+      "source": [
+        "!pip install git+https://github.com/jbischof/keras-nlp.git@bert_ckpt tensorflow tf-models-official --upgrade --quiet"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {
+        "id": "JsbnAdSz5DzZ"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "\n",
+        "import keras_nlp\n",
+        "import tensorflow as tf\n",
+        "import tensorflow_models as tfm\n",
+        "from tensorflow import keras"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "TOKEN_TYPE = \"cased\"\n",
+        "MODEL_TYPE = \"bert_base\"\n",
+        "MODEL_NAME = MODEL_TYPE + \"_\" + TOKEN_TYPE\n",
+        "VOCAB_SIZE = 28996"
+      ],
+      "metadata": {
+        "id": "DmVlNiSexzR7"
+      },
+      "execution_count": 3,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "L3j5PFBt5JeR"
+      },
+      "source": [
+        "## Load the model garden checkpoints and weights"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {
+        "id": "JdXFWsMVEf-x",
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "outputId": "4d87dbf9-dda5-4afe-fc07-533c84ad768a"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Downloading data from https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/cased_L-12_H-768_A-12.tar.gz\n",
+            "401886519/401886519 [==============================] - 3s 0us/step\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Model garden BERT paths.\n",
+        "zip_path = f\"\"\"https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/{TOKEN_TYPE}_L-12_H-768_A-12.tar.gz\"\"\"\n",
+        "zip_file = keras.utils.get_file(\n",
+        "    f\"\"\"/content/{MODEL_NAME}\"\"\",\n",
+        "    zip_path,\n",
+        "    extract=True,\n",
+        "    archive_format=\"tar\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 5,
+      "metadata": {
+        "id": "LHYiSsvYtfEU",
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "outputId": "8c63f095-4a49-4de0-b61a-8b2fe4705ff8"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "tmp/temp_dir/raw/\n",
+            "tmp/temp_dir/raw/vocab.txt\n",
+            "tmp/temp_dir/raw/bert_model.ckpt.index\n",
+            "tmp/temp_dir/raw/bert_model.ckpt.data-00000-of-00001\n",
+            "tmp/temp_dir/raw/bert_config.json\n"
+          ]
+        }
+      ],
+      "source": [
+        "!tar -xvf \"\"\"{MODEL_NAME}\"\"\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 6,
+      "metadata": {
+        "id": "Q9QRSp47tnVo"
+      },
+      "outputs": [],
+      "source": [
+        "# Model garden BERT paths.\n",
+        "extract_dir = \"/content/tmp/temp_dir/raw/\"\n",
+        "vocab_path = extract_dir + \"vocab.txt\"\n",
+        "checkpoint_path = extract_dir + \"bert_model.ckpt\"\n",
+        "config_path = extract_dir + \"bert_config.json\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 7,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "vCUJk59B2rai",
+        "outputId": "5f3b159b-381e-48e8-a77b-138b251a458c"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "_CHECKPOINTABLE_OBJECT_GRAPH []\n",
+            "encoder/layer_with_weights-0/embeddings/.ATTRIBUTES/VARIABLE_VALUE [28996, 768]\n",
+            "encoder/layer_with_weights-1/embeddings/.ATTRIBUTES/VARIABLE_VALUE [512, 768]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-10/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-10/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-10/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-10/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-10/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-10/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-10/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-10/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-11/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-11/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-11/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-11/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-11/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-11/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-11/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-11/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-12/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-12/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-12/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-12/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-12/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-12/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-12/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-12/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-13/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-13/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-13/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-13/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-13/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-13/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-13/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-13/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-14/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-14/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-14/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-14/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-14/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-14/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-14/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-14/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-15/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-15/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-15/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-15/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-15/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-15/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-15/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-15/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-16/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-16/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 768]\n",
+            "encoder/layer_with_weights-2/embeddings/.ATTRIBUTES/VARIABLE_VALUE [2, 768]\n",
+            "encoder/layer_with_weights-3/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-3/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-4/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-4/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-4/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-4/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-4/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-4/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-4/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-4/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-5/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-5/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-5/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-5/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-5/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-5/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-5/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-5/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-6/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-6/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-6/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-6/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-6/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-6/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-6/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-6/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-7/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-7/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-7/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-7/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-7/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-7/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-7/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-7/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-8/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-8/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-8/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-8/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-8/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-8/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-8/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-8/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-9/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-9/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-9/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-9/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-9/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-9/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-9/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-9/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n"
+          ]
+        }
+      ],
+      "source": [
+        "vars = tf.train.list_variables(checkpoint_path)\n",
+        "weights = {}\n",
+        "for name, shape in vars:\n",
+        "    print(name, shape)\n",
+        "    weight = tf.train.load_variable(checkpoint_path, name)\n",
+        "    weights[name] = weight"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "XrwVEkDP5RjE"
+      },
+      "source": [
+        "## Load BertBase model with KerasNLP."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 8,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Y5gJgD2j5BTG",
+        "outputId": "220a2d92-27b0-41d7-a081-d39aea11d89a"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Model: \"bert\"\n",
+            "__________________________________________________________________________________________________\n",
+            " Layer (type)                   Output Shape         Param #     Connected to                     \n",
+            "==================================================================================================\n",
+            " input_ids (InputLayer)         [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " token_embedding (Embedding)    (None, None, 768)    22268928    ['input_ids[0][0]']              \n",
+            "                                                                                                  \n",
+            " segment_ids (InputLayer)       [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " position_embedding (PositionEm  (None, None, 768)   393216      ['token_embedding[0][0]']        \n",
+            " bedding)                                                                                         \n",
+            "                                                                                                  \n",
+            " segment_embedding (Embedding)  (None, None, 768)    1536        ['segment_ids[0][0]']            \n",
+            "                                                                                                  \n",
+            " add (Add)                      (None, None, 768)    0           ['token_embedding[0][0]',        \n",
+            "                                                                  'position_embedding[0][0]',     \n",
+            "                                                                  'segment_embedding[0][0]']      \n",
+            "                                                                                                  \n",
+            " embeddings_layer_norm (LayerNo  (None, None, 768)   1536        ['add[0][0]']                    \n",
+            " rmalization)                                                                                     \n",
+            "                                                                                                  \n",
+            " embeddings_dropout (Dropout)   (None, None, 768)    0           ['embeddings_layer_norm[0][0]']  \n",
+            "                                                                                                  \n",
+            " input_mask (InputLayer)        [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " transformer_layer_0 (Transform  (None, None, 768)   7087872     ['embeddings_dropout[0][0]',     \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_1 (Transform  (None, None, 768)   7087872     ['transformer_layer_0[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_2 (Transform  (None, None, 768)   7087872     ['transformer_layer_1[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_3 (Transform  (None, None, 768)   7087872     ['transformer_layer_2[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_4 (Transform  (None, None, 768)   7087872     ['transformer_layer_3[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_5 (Transform  (None, None, 768)   7087872     ['transformer_layer_4[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_6 (Transform  (None, None, 768)   7087872     ['transformer_layer_5[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_7 (Transform  (None, None, 768)   7087872     ['transformer_layer_6[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_8 (Transform  (None, None, 768)   7087872     ['transformer_layer_7[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_9 (Transform  (None, None, 768)   7087872     ['transformer_layer_8[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_10 (Transfor  (None, None, 768)   7087872     ['transformer_layer_9[0][0]',    \n",
+            " merEncoder)                                                      'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_11 (Transfor  (None, None, 768)   7087872     ['transformer_layer_10[0][0]',   \n",
+            " merEncoder)                                                      'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " tf.__operators__.getitem (Slic  (None, 768)         0           ['transformer_layer_11[0][0]']   \n",
+            " ingOpLambda)                                                                                     \n",
+            "                                                                                                  \n",
+            " pooled_dense (Dense)           (None, 768)          590592      ['tf.__operators__.getitem[0][0]'\n",
+            "                                                                 ]                                \n",
+            "                                                                                                  \n",
+            "==================================================================================================\n",
+            "Total params: 108,310,272\n",
+            "Trainable params: 108,310,272\n",
+            "Non-trainable params: 0\n",
+            "__________________________________________________________________________________________________\n"
+          ]
+        }
+      ],
+      "source": [
+        "model = keras_nlp.models.BertBase(vocabulary_size=VOCAB_SIZE)\n",
+        "model.summary()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tZitZxcFyvlb"
+      },
+      "source": [
+        "## Convert Weights"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 9,
+      "metadata": {
+        "id": "wmuCofTwBgfo"
+      },
+      "outputs": [],
+      "source": [
+        "model.get_layer(\"token_embedding\").embeddings.assign(\n",
+        "    weights[\"encoder/layer_with_weights-0/embeddings/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "model.get_layer(\"position_embedding\").position_embeddings.assign(\n",
+        "    weights[\"encoder/layer_with_weights-1/embeddings/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "model.get_layer(\"segment_embedding\").embeddings.assign(\n",
+        "    weights[\"encoder/layer_with_weights-2/embeddings/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "model.get_layer(\"embeddings_layer_norm\").gamma.assign(\n",
+        "    weights[\"encoder/layer_with_weights-3/gamma/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "model.get_layer(\"embeddings_layer_norm\").beta.assign(\n",
+        "    weights[\"encoder/layer_with_weights-3/beta/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "\n",
+        "for i in range(model.num_layers):\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._key_dense.kernel.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._key_dense.bias.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._query_dense.kernel.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._query_dense.bias.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._value_dense.kernel.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._value_dense.bias.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._output_dense.kernel.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._output_dense.bias.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._attention_layernorm.gamma.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._attention_layernorm.beta.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._intermediate_dense.kernel.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._intermediate_dense.bias.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._output_dense.kernel.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._output_dense.bias.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._feedforward_layernorm.gamma.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._feedforward_layernorm.beta.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "\n",
+        "model.get_layer(\"pooled_dense\").kernel.assign(\n",
+        "    weights[\"encoder/layer_with_weights-16/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "model.get_layer(\"pooled_dense\").bias.assign(\n",
+        "    weights[\"encoder/layer_with_weights-16/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "pass"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tNxzCIaF_-IG"
+      },
+      "source": [
+        "## Compare Output"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 10,
+      "metadata": {
+        "id": "plQq3yxI_ry_"
+      },
+      "outputs": [],
+      "source": [
+        "def preprocess(x):\n",
+        "    tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(\n",
+        "        vocabulary=vocab_path,\n",
+        "    )\n",
+        "    packer = keras_nlp.layers.MultiSegmentPacker(\n",
+        "        sequence_length=model.max_sequence_length,\n",
+        "        start_value=tokenizer.token_to_id(\"[CLS]\"),\n",
+        "        end_value=tokenizer.token_to_id(\"[SEP]\"),\n",
+        "    )\n",
+        "    return packer(tokenizer(x))\n",
+        "\n",
+        "token_ids, segment_ids = preprocess([\"the quick brown fox.\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 11,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "chg-o813CJNJ",
+        "outputId": "3627782d-92d2-4299-8a5e-55a74402d01e"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7f3205f12b90>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 11
+        }
+      ],
+      "source": [
+        "encoder_config = tfm.nlp.encoders.EncoderConfig(\n",
+        "    type=\"bert\",\n",
+        "    bert=json.load(tf.io.gfile.GFile(config_path)),\n",
+        ")\n",
+        "mg_model = tfm.nlp.encoders.build_encoder(encoder_config)\n",
+        "checkpoint = tf.train.Checkpoint(encoder=mg_model)\n",
+        "checkpoint.read(checkpoint_path).assert_consumed()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 12,
+      "metadata": {
+        "id": "oFG-tRIKEzer"
+      },
+      "outputs": [],
+      "source": [
+        "keras_nlp_output = model({\n",
+        "    \"input_ids\": token_ids,\n",
+        "    \"segment_ids\": segment_ids,\n",
+        "    \"input_mask\": token_ids != 0,\n",
+        "})[\"pooled_output\"]\n",
+        "\n",
+        "mg_output = mg_model({\n",
+        "    \"input_word_ids\": token_ids,\n",
+        "    \"input_type_ids\": segment_ids,\n",
+        "    \"input_mask\": token_ids != 0,\n",
+        "})[\"pooled_output\"]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 27,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Swp2Y0OvYIId",
+        "outputId": "f413f078-4d4a-42ee-a048-bcf77e655f4b"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(10,), dtype=float32, numpy=\n",
+              "array([-0.81249267,  0.54269457,  0.9999706 , -0.9970537 ,  0.97125727,\n",
+              "        0.6073321 ,  0.9917275 , -0.95177907, -0.9899817 , -0.61669564],\n",
+              "      dtype=float32)>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 27
+        }
+      ],
+      "source": [
+        "keras_nlp_output[0, 0:10]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 28,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "zmhZYi-1YMGH",
+        "outputId": "b46aab26-c8c2-467f-e8f8-6c3fa992d5b2"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(10,), dtype=float32, numpy=\n",
+              "array([-0.81249064,  0.54269207,  0.9999706 , -0.9970535 ,  0.97125703,\n",
+              "        0.6073294 ,  0.99172735, -0.9517791 , -0.9899815 , -0.6166962 ],\n",
+              "      dtype=float32)>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 28
+        }
+      ],
+      "source": [
+        "mg_output[0, 0:10]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 15,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "TvNj8DFBYNPT",
+        "outputId": "4447e234-bf17-46f2-f09f-86493a61e1cb"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(), dtype=float32, numpy=-3.0643907e-08>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 15
+        }
+      ],
+      "source": [
+        "# Very close! Though not 100% exact.\n",
+        "tf.reduce_mean(keras_nlp_output - mg_output)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 16,
+      "metadata": {
+        "id": "aXu7Y_zue70C"
+      },
+      "outputs": [],
+      "source": [
+        "# Save BertBase checkpoint\n",
+        "model.save_weights(f\"\"\"{MODEL_NAME}.h5\"\"\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 17,
+      "metadata": {
+        "id": "bwlLYTFeg1og"
+      },
+      "outputs": [],
+      "source": [
+        "model2 = keras_nlp.models.BertBase(vocabulary_size=VOCAB_SIZE)\n",
+        "model2.load_weights(f\"\"\"{MODEL_NAME}.h5\"\"\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 18,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "yTqPl39qhPMV",
+        "outputId": "42e126a0-92af-4779-f652-47b3396d38a9"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(), dtype=float32, numpy=0.0>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 18
+        }
+      ],
+      "source": [
+        "# Same output from loaded checkpoint\n",
+        "keras_nlp_output2 = model2({\n",
+        "    \"input_ids\": token_ids,\n",
+        "    \"segment_ids\": segment_ids,\n",
+        "    \"input_mask\": token_ids != 0,\n",
+        "})[\"pooled_output\"]\n",
+        "tf.reduce_mean(keras_nlp_output - keras_nlp_output2)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 19,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "cXiV4NvilLg6",
+        "outputId": "f506140a-217b-4f45-b331-793e47fc5bae"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "210441"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 19
+        }
+      ],
+      "source": [
+        "# Save vocab file as well\n",
+        "vocab_info = tf.io.gfile.GFile(vocab_path).read()\n",
+        "f = open(\"vocab.txt\", \"w\")\n",
+        "f.write(vocab_info)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 20,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "lVutnU-hB6IQ",
+        "outputId": "07debeae-6f54-4201-8d15-8cfa992fb121"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "f30ac6fac11115322e3d0c61e87e98b2  bert_base_cased.h5\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Get MD5 of model\n",
+        "!md5sum \"\"\"{MODEL_NAME}.h5\"\"\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 21,
+      "metadata": {
+        "id": "z_0iMTCdFl8t"
+      },
+      "outputs": [],
+      "source": [
+        "# Upload model to drive\n",
+        "#from google.colab import drive\n",
+        "#drive.mount('/content/drive')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 22,
+      "metadata": {
+        "id": "wTd-5vUyVG0Q",
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "outputId": "b96360bc-47d6-4611-e727-f861d6038726"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Downloading data from https://storage.googleapis.com/keras-nlp/models/bert_base_cased/model.h5\n",
+            "433474808/433474808 [==============================] - 5s 0us/step\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Check uploaded model once added to repo\n",
+        "model_cloud = keras_nlp.models.BertBase(weights=MODEL_NAME)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 23,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "zs5x_f6GVdNY",
+        "outputId": "f1aa7c12-6681-4ae2-c689-f4b7ddfd1701"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(), dtype=float32, numpy=0.0>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 23
+        }
+      ],
+      "source": [
+        "# Same output from cloud model\n",
+        "keras_nlp_output_cloud = model_cloud({\n",
+        "    \"input_ids\": token_ids,\n",
+        "    \"segment_ids\": segment_ids,\n",
+        "    \"input_mask\": token_ids != 0,\n",
+        "})[\"pooled_output\"]\n",
+        "tf.reduce_mean(keras_nlp_output - keras_nlp_output_cloud)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "keras_nlp_output_cloud[0, 0:10]"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "OKfvONXmzX_L",
+        "outputId": "ece23e28-78b2-415a-faba-be5a3825812f"
+      },
+      "execution_count": 26,
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(10,), dtype=float32, numpy=\n",
+              "array([-0.81249267,  0.54269457,  0.9999706 , -0.9970537 ,  0.97125727,\n",
+              "        0.6073321 ,  0.9917275 , -0.95177907, -0.9899817 , -0.61669564],\n",
+              "      dtype=float32)>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 26
+        }
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "name": "checkpoint convert model garden -> keras-nlp Bert cased",
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU",
+    "gpuClass": "standard"
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
\ No newline at end of file
diff --git a/tools/checkpoint_conversion/bert_base_uncased.ipynb b/tools/checkpoint_conversion/bert_base_uncased.ipynb
new file mode 100644
index 00000000..d8bf48b5
--- /dev/null
+++ b/tools/checkpoint_conversion/bert_base_uncased.ipynb
@@ -0,0 +1,994 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/jbischof/keras-nlp/blob/bert_ckpt/bert_based_uncased.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vGp_yrJi5Ehf"
+      },
+      "source": [
+        "## Install deps"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Szd6xKUd2tIE",
+        "outputId": "6fa43e61-f3ca-449b-a6ea-971bae8f5020"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "\u001b[K     |████████████████████████████████| 511.7 MB 6.0 kB/s \n",
+            "\u001b[K     |████████████████████████████████| 2.1 MB 46.1 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 4.6 MB 57.2 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 438 kB 66.9 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 1.6 MB 56.9 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 5.8 MB 26.7 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 99 kB 9.2 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 116 kB 70.6 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 1.3 MB 59.9 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 1.1 MB 52.8 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 636 kB 75.9 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 43 kB 2.0 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 238 kB 68.0 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 352 kB 76.8 MB/s \n",
+            "\u001b[?25h  Building wheel for keras-nlp (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
+            "  Building wheel for py-cpuinfo (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
+            "  Building wheel for seqeval (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
+          ]
+        }
+      ],
+      "source": [
+        "!pip install git+https://github.com/jbischof/keras-nlp.git@bert_ckpt tensorflow tf-models-official --upgrade --quiet"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {
+        "id": "JsbnAdSz5DzZ"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "\n",
+        "import keras_nlp\n",
+        "import tensorflow as tf\n",
+        "import tensorflow_models as tfm\n",
+        "from tensorflow import keras"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "TOKEN_TYPE = \"uncased\"\n",
+        "MODEL_TYPE = \"bert_base\"\n",
+        "MODEL_NAME = MODEL_TYPE + \"_\" + TOKEN_TYPE\n",
+        "VOCAB_SIZE = 30522"
+      ],
+      "metadata": {
+        "id": "DmVlNiSexzR7"
+      },
+      "execution_count": 3,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "L3j5PFBt5JeR"
+      },
+      "source": [
+        "## Load the model garden checkpoints and weights"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {
+        "id": "JdXFWsMVEf-x",
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "outputId": "fe21cd4d-7f08-4d49-ef07-1660e3e76fa4"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Downloading data from https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/uncased_L-12_H-768_A-12.tar.gz\n",
+            "405351189/405351189 [==============================] - 3s 0us/step\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Model garden BERT paths.\n",
+        "zip_path = f\"\"\"https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/{TOKEN_TYPE}_L-12_H-768_A-12.tar.gz\"\"\"\n",
+        "zip_file = keras.utils.get_file(\n",
+        "    f\"\"\"/content/{MODEL_NAME}\"\"\",\n",
+        "    zip_path,\n",
+        "    extract=True,\n",
+        "    archive_format=\"tar\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 5,
+      "metadata": {
+        "id": "LHYiSsvYtfEU",
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "outputId": "77580043-09a7-4950-99bb-473d71614365"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "tmp/temp_dir/raw/\n",
+            "tmp/temp_dir/raw/vocab.txt\n",
+            "tmp/temp_dir/raw/bert_model.ckpt.index\n",
+            "tmp/temp_dir/raw/bert_model.ckpt.data-00000-of-00001\n",
+            "tmp/temp_dir/raw/bert_config.json\n"
+          ]
+        }
+      ],
+      "source": [
+        "!tar -xvf \"\"\"{MODEL_NAME}\"\"\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 6,
+      "metadata": {
+        "id": "Q9QRSp47tnVo"
+      },
+      "outputs": [],
+      "source": [
+        "# Model garden BERT paths.\n",
+        "extract_dir = \"/content/tmp/temp_dir/raw/\"\n",
+        "vocab_path = extract_dir + \"vocab.txt\"\n",
+        "checkpoint_path = extract_dir + \"bert_model.ckpt\"\n",
+        "config_path = extract_dir + \"bert_config.json\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 7,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "vCUJk59B2rai",
+        "outputId": "67ca9d60-c88e-4ad1-c93a-a06eeaf2e631"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "_CHECKPOINTABLE_OBJECT_GRAPH []\n",
+            "encoder/layer_with_weights-0/embeddings/.ATTRIBUTES/VARIABLE_VALUE [30522, 768]\n",
+            "encoder/layer_with_weights-1/embeddings/.ATTRIBUTES/VARIABLE_VALUE [512, 768]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-10/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-10/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-10/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-10/_attention_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-10/_attention_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-10/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-10/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-10/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-10/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-10/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-10/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-11/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-11/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-11/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-11/_attention_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-11/_attention_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-11/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-11/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-11/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-11/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-11/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-11/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-12/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-12/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-12/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-12/_attention_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-12/_attention_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-12/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-12/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-12/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-12/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-12/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-12/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-13/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-13/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-13/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-13/_attention_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-13/_attention_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-13/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-13/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-13/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-13/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-13/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-13/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-14/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-14/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-14/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-14/_attention_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-14/_attention_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-14/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-14/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-14/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-14/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-14/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-14/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-15/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-15/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-15/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-15/_attention_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-15/_attention_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-15/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-15/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-15/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-15/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-15/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-15/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-2/embeddings/.ATTRIBUTES/VARIABLE_VALUE [2, 768]\n",
+            "encoder/layer_with_weights-3/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-3/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-4/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-4/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-4/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-4/_attention_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-4/_attention_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-4/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-4/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-4/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-4/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-4/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-4/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-5/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-5/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-5/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-5/_attention_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-5/_attention_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-5/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-5/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-5/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-5/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-5/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-5/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-6/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-6/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-6/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-6/_attention_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-6/_attention_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-6/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-6/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-6/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-6/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-6/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-6/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-7/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-7/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-7/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-7/_attention_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-7/_attention_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-7/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-7/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-7/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-7/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-7/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-7/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-8/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-8/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-8/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-8/_attention_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-8/_attention_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-8/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-8/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-8/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-8/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-8/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-8/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [12, 64]\n",
+            "encoder/layer_with_weights-9/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 12, 64]\n",
+            "encoder/layer_with_weights-9/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-9/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-9/_attention_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-9/_attention_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [12, 64, 768]\n",
+            "encoder/layer_with_weights-9/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [3072]\n",
+            "encoder/layer_with_weights-9/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 3072]\n",
+            "encoder/layer_with_weights-9/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-9/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [3072, 768]\n",
+            "encoder/layer_with_weights-9/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "encoder/layer_with_weights-9/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "next_sentence..pooler_dense/bias/.ATTRIBUTES/VARIABLE_VALUE [768]\n",
+            "next_sentence..pooler_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE [768, 768]\n"
+          ]
+        }
+      ],
+      "source": [
+        "vars = tf.train.list_variables(checkpoint_path)\n",
+        "weights = {}\n",
+        "for name, shape in vars:\n",
+        "    print(name, shape)\n",
+        "    weight = tf.train.load_variable(checkpoint_path, name)\n",
+        "    weights[name] = weight"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "XrwVEkDP5RjE"
+      },
+      "source": [
+        "## Load BertBase model with KerasNLP."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 8,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Y5gJgD2j5BTG",
+        "outputId": "e130445e-1676-4a21-b332-b373814aef14"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Model: \"bert\"\n",
+            "__________________________________________________________________________________________________\n",
+            " Layer (type)                   Output Shape         Param #     Connected to                     \n",
+            "==================================================================================================\n",
+            " input_ids (InputLayer)         [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " token_embedding (Embedding)    (None, None, 768)    23440896    ['input_ids[0][0]']              \n",
+            "                                                                                                  \n",
+            " segment_ids (InputLayer)       [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " position_embedding (PositionEm  (None, None, 768)   393216      ['token_embedding[0][0]']        \n",
+            " bedding)                                                                                         \n",
+            "                                                                                                  \n",
+            " segment_embedding (Embedding)  (None, None, 768)    1536        ['segment_ids[0][0]']            \n",
+            "                                                                                                  \n",
+            " add (Add)                      (None, None, 768)    0           ['token_embedding[0][0]',        \n",
+            "                                                                  'position_embedding[0][0]',     \n",
+            "                                                                  'segment_embedding[0][0]']      \n",
+            "                                                                                                  \n",
+            " embeddings_layer_norm (LayerNo  (None, None, 768)   1536        ['add[0][0]']                    \n",
+            " rmalization)                                                                                     \n",
+            "                                                                                                  \n",
+            " embeddings_dropout (Dropout)   (None, None, 768)    0           ['embeddings_layer_norm[0][0]']  \n",
+            "                                                                                                  \n",
+            " input_mask (InputLayer)        [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " transformer_layer_0 (Transform  (None, None, 768)   7087872     ['embeddings_dropout[0][0]',     \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_1 (Transform  (None, None, 768)   7087872     ['transformer_layer_0[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_2 (Transform  (None, None, 768)   7087872     ['transformer_layer_1[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_3 (Transform  (None, None, 768)   7087872     ['transformer_layer_2[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_4 (Transform  (None, None, 768)   7087872     ['transformer_layer_3[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_5 (Transform  (None, None, 768)   7087872     ['transformer_layer_4[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_6 (Transform  (None, None, 768)   7087872     ['transformer_layer_5[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_7 (Transform  (None, None, 768)   7087872     ['transformer_layer_6[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_8 (Transform  (None, None, 768)   7087872     ['transformer_layer_7[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_9 (Transform  (None, None, 768)   7087872     ['transformer_layer_8[0][0]',    \n",
+            " erEncoder)                                                       'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_10 (Transfor  (None, None, 768)   7087872     ['transformer_layer_9[0][0]',    \n",
+            " merEncoder)                                                      'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " transformer_layer_11 (Transfor  (None, None, 768)   7087872     ['transformer_layer_10[0][0]',   \n",
+            " merEncoder)                                                      'input_mask[0][0]']             \n",
+            "                                                                                                  \n",
+            " tf.__operators__.getitem (Slic  (None, 768)         0           ['transformer_layer_11[0][0]']   \n",
+            " ingOpLambda)                                                                                     \n",
+            "                                                                                                  \n",
+            " pooled_dense (Dense)           (None, 768)          590592      ['tf.__operators__.getitem[0][0]'\n",
+            "                                                                 ]                                \n",
+            "                                                                                                  \n",
+            "==================================================================================================\n",
+            "Total params: 109,482,240\n",
+            "Trainable params: 109,482,240\n",
+            "Non-trainable params: 0\n",
+            "__________________________________________________________________________________________________\n"
+          ]
+        }
+      ],
+      "source": [
+        "model = keras_nlp.models.BertBase(vocabulary_size=VOCAB_SIZE)\n",
+        "model.summary()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tZitZxcFyvlb"
+      },
+      "source": [
+        "## Convert Weights"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 9,
+      "metadata": {
+        "id": "wmuCofTwBgfo"
+      },
+      "outputs": [],
+      "source": [
+        "model.get_layer(\"token_embedding\").embeddings.assign(\n",
+        "    weights[\"encoder/layer_with_weights-0/embeddings/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "model.get_layer(\"position_embedding\").position_embeddings.assign(\n",
+        "    weights[\"encoder/layer_with_weights-1/embeddings/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "model.get_layer(\"segment_embedding\").embeddings.assign(\n",
+        "    weights[\"encoder/layer_with_weights-2/embeddings/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "model.get_layer(\"embeddings_layer_norm\").gamma.assign(\n",
+        "    weights[\"encoder/layer_with_weights-3/gamma/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "model.get_layer(\"embeddings_layer_norm\").beta.assign(\n",
+        "    weights[\"encoder/layer_with_weights-3/beta/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "\n",
+        "for i in range(model.num_layers):\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._key_dense.kernel.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_key_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._key_dense.bias.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_key_dense/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._query_dense.kernel.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_query_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._query_dense.bias.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_query_dense/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._value_dense.kernel.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_value_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._value_dense.bias.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer/_value_dense/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._output_dense.kernel.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._multi_head_attention_layer._output_dense.bias.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._attention_layernorm.gamma.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._attention_layernorm.beta.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_attention_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._intermediate_dense.kernel.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_intermediate_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._intermediate_dense.bias.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_intermediate_dense/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._output_dense.kernel.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_output_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._output_dense.bias.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_output_dense/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._feedforward_layernorm.gamma.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_output_layer_norm/gamma/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "    model.get_layer(f\"transformer_layer_{i}\")._feedforward_layernorm.beta.assign(\n",
+        "        weights[f\"encoder/layer_with_weights-{i + 4}/_output_layer_norm/beta/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        "    )\n",
+        "\n",
+        "model.get_layer(\"pooled_dense\").kernel.assign(\n",
+        "    weights[\"next_sentence..pooler_dense/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "model.get_layer(\"pooled_dense\").bias.assign(\n",
+        "    weights[\"next_sentence..pooler_dense/bias/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "pass"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tNxzCIaF_-IG"
+      },
+      "source": [
+        "## Compare Output"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 10,
+      "metadata": {
+        "id": "plQq3yxI_ry_"
+      },
+      "outputs": [],
+      "source": [
+        "def preprocess(x):\n",
+        "    tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(\n",
+        "        vocabulary=vocab_path,\n",
+        "    )\n",
+        "    packer = keras_nlp.layers.MultiSegmentPacker(\n",
+        "        sequence_length=model.max_sequence_length,\n",
+        "        start_value=tokenizer.token_to_id(\"[CLS]\"),\n",
+        "        end_value=tokenizer.token_to_id(\"[SEP]\"),\n",
+        "    )\n",
+        "    return packer(tokenizer(x))\n",
+        "\n",
+        "token_ids, segment_ids = preprocess([\"the quick brown fox.\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 11,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "chg-o813CJNJ",
+        "outputId": "72dfdb08-dbbb-4f8a-8a0a-2d881659edeb"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7fbd6a23ee50>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 11
+        }
+      ],
+      "source": [
+        "encoder_config = tfm.nlp.encoders.EncoderConfig(\n",
+        "    type=\"bert\",\n",
+        "    bert=json.load(tf.io.gfile.GFile(config_path)),\n",
+        ")\n",
+        "mg_model = tfm.nlp.encoders.build_encoder(encoder_config)\n",
+        "checkpoint = tf.train.Checkpoint(encoder=mg_model)\n",
+        "checkpoint.read(checkpoint_path).assert_consumed()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 12,
+      "metadata": {
+        "id": "oFG-tRIKEzer"
+      },
+      "outputs": [],
+      "source": [
+        "keras_nlp_output = model({\n",
+        "    \"input_ids\": token_ids,\n",
+        "    \"segment_ids\": segment_ids,\n",
+        "    \"input_mask\": token_ids != 0,\n",
+        "})[\"pooled_output\"]\n",
+        "\n",
+        "mg_output = mg_model({\n",
+        "    \"input_word_ids\": token_ids,\n",
+        "    \"input_type_ids\": segment_ids,\n",
+        "    \"input_mask\": token_ids != 0,\n",
+        "})[\"pooled_output\"]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 13,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Swp2Y0OvYIId",
+        "outputId": "8a981dbb-eb5b-432c-ee6d-52419e962cfd"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(10,), dtype=float32, numpy=\n",
+              "array([-0.8692639 , -0.46654877, -0.87829065,  0.81910247,  0.6255867 ,\n",
+              "       -0.09378534,  0.9116511 ,  0.32022637, -0.77492046, -0.9999707 ],\n",
+              "      dtype=float32)>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 13
+        }
+      ],
+      "source": [
+        "keras_nlp_output[0, 0:10]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 14,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "zmhZYi-1YMGH",
+        "outputId": "2cb14435-68bc-48e3-b462-80842686d0e9"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(10,), dtype=float32, numpy=\n",
+              "array([-0.8692642 , -0.46654695, -0.8782886 ,  0.81910086,  0.6255875 ,\n",
+              "       -0.0937808 ,  0.91164833,  0.32022265, -0.77491295, -0.99997073],\n",
+              "      dtype=float32)>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 14
+        }
+      ],
+      "source": [
+        "mg_output[0, 0:10]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 15,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "TvNj8DFBYNPT",
+        "outputId": "c7fc2d2d-919f-443e-85dd-d7139d01b173"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(), dtype=float32, numpy=-1.7592144e-07>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 15
+        }
+      ],
+      "source": [
+        "# Very close! Though not 100% exact.\n",
+        "tf.reduce_mean(keras_nlp_output - mg_output)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 16,
+      "metadata": {
+        "id": "aXu7Y_zue70C"
+      },
+      "outputs": [],
+      "source": [
+        "# Save BertBase checkpoint\n",
+        "model.save_weights(f\"\"\"{MODEL_NAME}.h5\"\"\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 17,
+      "metadata": {
+        "id": "bwlLYTFeg1og"
+      },
+      "outputs": [],
+      "source": [
+        "model2 = keras_nlp.models.BertBase(vocabulary_size=VOCAB_SIZE)\n",
+        "model2.load_weights(f\"\"\"{MODEL_NAME}.h5\"\"\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 18,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "yTqPl39qhPMV",
+        "outputId": "0b8e4fbc-c356-4f66-d01c-d43fc1a326fb"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(), dtype=float32, numpy=0.0>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 18
+        }
+      ],
+      "source": [
+        "# Same output from loaded checkpoint\n",
+        "keras_nlp_output2 = model2({\n",
+        "    \"input_ids\": token_ids,\n",
+        "    \"segment_ids\": segment_ids,\n",
+        "    \"input_mask\": token_ids != 0,\n",
+        "})[\"pooled_output\"]\n",
+        "tf.reduce_mean(keras_nlp_output - keras_nlp_output2)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 19,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "cXiV4NvilLg6",
+        "outputId": "76d51be8-d503-4807-c78e-0b94a64ecb6c"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "228209"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 19
+        }
+      ],
+      "source": [
+        "# Save vocab file as well\n",
+        "vocab_info = tf.io.gfile.GFile(vocab_path).read()\n",
+        "f = open(\"vocab.txt\", \"w\")\n",
+        "f.write(vocab_info)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 20,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "lVutnU-hB6IQ",
+        "outputId": "9312bab4-4889-42b4-a7d6-b7e59c45dfcc"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "94d5ec8bea7816f0c29904c959284343  bert_base_uncased.h5\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Get MD5 of model\n",
+        "!md5sum \"\"\"{MODEL_NAME}.h5\"\"\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 21,
+      "metadata": {
+        "id": "z_0iMTCdFl8t"
+      },
+      "outputs": [],
+      "source": [
+        "# Upload model to drive\n",
+        "#from google.colab import drive\n",
+        "#drive.mount('/content/drive')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 22,
+      "metadata": {
+        "id": "wTd-5vUyVG0Q",
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "outputId": "c92995ea-0e01-4289-f45b-3a52900ac7e7"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Downloading data from https://storage.googleapis.com/keras-nlp/models/bert_base_uncased/model.h5\n",
+            "438162680/438162680 [==============================] - 5s 0us/step\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Check uploaded model once added to repo\n",
+        "model_cloud = keras_nlp.models.BertBase(weights=MODEL_NAME)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 23,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "zs5x_f6GVdNY",
+        "outputId": "fff1c4e5-a41c-42d4-df2d-303d4bc8a0ec"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(), dtype=float32, numpy=0.0>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 23
+        }
+      ],
+      "source": [
+        "# Same output from cloud model\n",
+        "keras_nlp_output_cloud = model_cloud({\n",
+        "    \"input_ids\": token_ids,\n",
+        "    \"segment_ids\": segment_ids,\n",
+        "    \"input_mask\": token_ids != 0,\n",
+        "})[\"pooled_output\"]\n",
+        "tf.reduce_mean(keras_nlp_output - keras_nlp_output_cloud)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "keras_nlp_output_cloud[0, 0:10]"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "RAwrhAcSzHWa",
+        "outputId": "6b0bdfbe-27df-483c-a761-2ef34d8e5417"
+      },
+      "execution_count": 24,
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(10,), dtype=float32, numpy=\n",
+              "array([-0.8692639 , -0.46654877, -0.87829065,  0.81910247,  0.6255867 ,\n",
+              "       -0.09378534,  0.9116511 ,  0.32022637, -0.77492046, -0.9999707 ],\n",
+              "      dtype=float32)>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 24
+        }
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "name": "checkpoint convert model garden -> keras-nlp Bert uncased",
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU",
+    "gpuClass": "standard"
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
\ No newline at end of file


Test patch:
diff --git a/keras_nlp/models/bert_test.py b/keras_nlp/models/bert_test.py
index 1518d2c4..e4fa6b14 100644
--- a/keras_nlp/models/bert_test.py
+++ b/keras_nlp/models/bert_test.py
@@ -34,7 +34,7 @@ class BertTest(tf.test.TestCase):
         )
         input_data = {
             "input_ids": tf.random.uniform(
-                shape=(1, 12), dtype=tf.int64, maxval=30522
+                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)
@@ -88,7 +88,7 @@ class BertTest(tf.test.TestCase):
         classifier(input_data)
 
     def test_valid_call_bert_base(self):
-        model = bert.BertBase(name="encoder")
+        model = bert.BertBase(vocabulary_size=10000, name="encoder")
         input_data = {
             "input_ids": tf.random.uniform(
                 shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
@@ -98,6 +98,26 @@ class BertTest(tf.test.TestCase):
         }
         model(input_data)
 
+    def test_bert_base_vocab_error(self):
+        # Need `vocabulary_size` or `weights`
+        with self.assertRaises(ValueError):
+            bert.BertBase(name="encoder")
+
+        # Only one of `vocabulary_size` or `weights`
+        with self.assertRaises(ValueError):
+            bert.BertBase(
+                weights="bert_base_uncased",
+                vocabulary_size=10000,
+                name="encoder",
+            )
+
+        # Not a checkpoint name
+        with self.assertRaises(ValueError):
+            bert.BertBase(
+                weights="bert_base_clowntown",
+                name="encoder",
+            )
+
     def test_saving_model(self):
         model = bert.BertCustom(
             vocabulary_size=30522,
@@ -110,7 +130,7 @@ class BertTest(tf.test.TestCase):
         )
         input_data = {
             "input_ids": tf.random.uniform(
-                shape=(1, 12), dtype=tf.int64, maxval=30522
+                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)
diff --git a/keras_nlp/models/roberta_test.py b/keras_nlp/models/roberta_test.py
new file mode 100644
index 00000000..cd035df4
--- /dev/null
+++ b/keras_nlp/models/roberta_test.py
@@ -0,0 +1,124 @@
+# 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 Roberta model."""
+
+import os
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.models import roberta
+
+
+class RobertaTest(tf.test.TestCase):
+    def test_valid_call_roberta(self):
+        model = roberta.RobertaCustom(
+            vocabulary_size=50265,
+            num_layers=12,
+            num_heads=12,
+            hidden_dim=768,
+            intermediate_dim=3072,
+            max_sequence_length=12,
+            name="encoder",
+        )
+        input_data = {
+            "input_ids": tf.random.uniform(
+                shape=(1, 12), dtype=tf.int64, maxval=50265
+            ),
+            "input_mask": tf.constant(
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+            ),
+        }
+        model(input_data)
+
+    def test_valid_call_classifier(self):
+        model = roberta.RobertaCustom(
+            vocabulary_size=50265,
+            num_layers=12,
+            num_heads=12,
+            hidden_dim=768,
+            intermediate_dim=3072,
+            max_sequence_length=12,
+            name="encoder",
+        )
+        input_data = {
+            "input_ids": tf.random.uniform(
+                shape=(1, 12), dtype=tf.int64, maxval=50265
+            ),
+            "input_mask": tf.constant(
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+            ),
+        }
+        classifier = roberta.RobertaClassifier(model, 4, 768, name="classifier")
+        classifier(input_data)
+
+    def test_valid_call_roberta_base(self):
+        model = roberta.RobertaBase(name="encoder")
+        input_data = {
+            "input_ids": tf.random.uniform(
+                shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
+            ),
+            "input_mask": tf.constant([1] * 512, shape=(1, 512)),
+        }
+        model(input_data)
+
+    def test_variable_sequence_length_call_roberta(self):
+        model = roberta.RobertaCustom(
+            vocabulary_size=50265,
+            num_layers=12,
+            num_heads=12,
+            hidden_dim=768,
+            intermediate_dim=3072,
+            max_sequence_length=100,
+            name="encoder",
+        )
+        for seq_length in (25, 50, 100):
+            input_data = {
+                "input_ids": tf.ones((8, seq_length), dtype="int32"),
+                "input_mask": tf.ones((8, seq_length), dtype="int32"),
+            }
+            output = model(input_data)
+            self.assertAllEqual(
+                tf.shape(output["sequence_output"]), [8, seq_length, 768]
+            )
+
+    def test_saving_model(self):
+        model = roberta.RobertaCustom(
+            vocabulary_size=50265,
+            num_layers=12,
+            num_heads=12,
+            hidden_dim=768,
+            intermediate_dim=3072,
+            max_sequence_length=12,
+            name="encoder",
+        )
+        input_data = {
+            "input_ids": tf.random.uniform(
+                shape=(1, 12), dtype=tf.int64, maxval=50265
+            ),
+            "input_mask": tf.constant(
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+            ),
+        }
+        model_output = model.predict(input_data)
+
+        save_path = os.path.join(self.get_temp_dir(), "model")
+        model.save(save_path)
+        restored_model = keras.models.load_model(save_path)
+
+        restored_output = restored_model.predict(input_data)
+        self.assertAllClose(
+            model_output["sequence_output"],
+            restored_output["sequence_output"],
+        )
diff --git a/network_tests/test_load_ckpts.py b/network_tests/test_load_ckpts.py
new file mode 100644
index 00000000..2d00ef75
--- /dev/null
+++ b/network_tests/test_load_ckpts.py
@@ -0,0 +1,42 @@
+# 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 loading pretrained model checkpoints."""
+
+import tensorflow as tf
+
+import keras_nlp
+
+
+class BertCkptTest(tf.test.TestCase):
+    def test_load_bert_base_uncased(self):
+        model = keras_nlp.models.BertBase(weights="uncased_en")
+        input_data = {
+            "input_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)),
+            "input_mask": tf.constant([1] * 512, shape=(1, 512)),
+        }
+        model(input_data)
+
+    def test_load_bert_base_cased(self):
+        model = keras_nlp.models.BertBase(weights="cased_en")
+        input_data = {
+            "input_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)),
+            "input_mask": tf.constant([1] * 512, shape=(1, 512)),
+        }
+        model(input_data)