Repository: keras-team/keras-nlp

Issue:
Add pretrained checkpoints for Bert
# Problem
We exposed an API for fully configurable and fixed Bert architectures in PR #288. Most users will not have the resources to do pretraining on their own, however. We should offer checkpoints for `models.BertBase` similar to those used in the original paper.

# Implementation

## Weights

Model checkpoints from the original paper repo ([gh](https://github.com/google-research/bert)) or similar source can be converted to be compatible with our implementation in PR #288. This code will be a one-off task and not added to the repo.

## Format

`keras-cv` is currently using `h5`, but that format is deprecated. An obvious alternative is a `tf.checkpoint` using `model.save_weights()`. The resulting `data` and `index` files will have to be zipped and uploaded to our cloud storage bucket.

## API

There some prior art in `keras-cv` and our pretrained vocab to use a string to identify a model of interest which is then checked against a dictionary (e.g., [weights.py](https://github.com/keras-team/keras-cv/blob/master/keras_cv/models/weights.py#L58)). However, this requires the user to dive into the code to find possible names and a lot of input validation and error handling in the implementation.

I propose to create an `enum` for each `Model` class (e.g., `models.BertBase`) with possible checkpoints that will expose these options as part of the API and work with users' autocomplete engines:
```
class BertBaseCkpt(Enum):
   uncased = 0
   cased = 1
   zh = 2
```

Under the hood this `enum` would then be the key for a dictionary of metadata for each model including checkpoint location, vocabulary location, and a description string explaining where the model came from and what data was used to train it. We could also write a method to print the description for each/all `enum` value(s) to again protect the user from having to read through utility functions.

```
base_models = {
   BertBaseCkpt.uncased: {
      "ckpt": ...
      "description": ...
      "vocab": ...
   }
}
```

Model checkpoints will be hosted in our Google Cloud storage bucket and loaded behind the scenes given an enum value. The user should never have to interact with this dictionary.

Here is the enhanced API:
```
def BertBase(ckpt=None, name=None, trainable=True):
    """Bi-directional Transformer-based encoder network (Bert) using "Base"
    architecture.

    This network implements a bi-directional Transformer-based encoder as
    described in ["BERT: Pre-training of Deep Bidirectional Transformers for
    Language Understanding"](https://arxiv.org/abs/1810.04805). It includes the
    embedding lookups and transformer layers, but not the masked language model
    or classification task networks.

    Args:
        ckpt: BertBaseCkpt. A checkpoint of pretrained weights to load into the model. 
             If `None`, the model will be randomly initialized.
        name: String, optional. Name of the model.
        trainable: Boolean, optional. If the model's variables should be
            trainable.
```

## Vocabulary
Note that this proposal does not cover the vocabulary file, which will be incorporated in our preprocessing work. One potential proposal is to host the `vocab.txt` file used to train the model in the same folder to be loaded if the user does not supply their own. The rationale is that except for very advanced use cases the vocabulary must match that used in training, and requiring the user to specify their own will expose our Cloud storage implementation to the user for little gain. However, this is beyond the scope of the proposal.

Discussion:

Patch:
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/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/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)