Repository: keras-team/keras-nlp
Pull request #319 (2022-08-27 13:44:25)

`bert_base_zh`, `bert_base_multi_cased`: Add BERT Base Variants
Resolves https://github.com/keras-team/keras-nlp/issues/308

Converted model weights and vocab have been uploaded here: https://drive.google.com/drive/folders/19APUi7fobORdQjoe8YDyk8ou0x_L6hWu?usp=sharing.

Edit: Merging https://github.com/keras-team/keras-nlp/pull/320 with this PR. Might cause conflicts while merging with `master`, otherwise.

mattdangerw: @abheesht17 I think we would want to avoid the add_pooling_layer parameter if possible.

One thing we could check is the bert checkpoints on the original bert repo. Could you see if those include the pooling layer?

https://github.com/google-research/bert
abheesht17: > @abheesht17 I think we would want to avoid the add_pooling_layer parameter if possible.
> 
> One thing we could check is the bert checkpoints on the original bert repo. Could you see if those include the pooling layer?
> 
> https://github.com/google-research/bert

My bad. The pooling layers are present in this ckpt after all: https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/multi_cased_L-12_H-768_A-12.tar.gz. Missed them earlier

Patch:
diff --git a/examples/bert/bert_finetune_glue.py b/examples/bert/bert_finetune_glue.py
index 6b1beec6..afb46672 100644
--- a/examples/bert/bert_finetune_glue.py
+++ b/examples/bert/bert_finetune_glue.py
@@ -140,9 +140,9 @@ def main(_):
         inputs = [tokenizer(x) for x in inputs]
         token_ids, segment_ids = packer(inputs)
         return {
-            "input_ids": token_ids,
-            "input_mask": tf.cast(token_ids != 0, "int32"),
+            "token_ids": token_ids,
             "segment_ids": segment_ids,
+            "padding_mask": tf.cast(token_ids != 0, "int32"),
         }, labels
 
     # Read and preprocess GLUE task data.
diff --git a/examples/bert/bert_preprocess.py b/examples/bert/bert_preprocess.py
index 00e0e8cb..91113efd 100644
--- a/examples/bert/bert_preprocess.py
+++ b/examples/bert/bert_preprocess.py
@@ -125,18 +125,18 @@ def write_instance_to_example_files(
     total_written = 0
     lookup = dict(zip(vocab, range(len(vocab))))
     for (inst_index, instance) in enumerate(instances):
-        input_ids = [lookup[x] for x in instance.tokens]
-        input_mask = [1] * len(input_ids)
+        token_ids = [lookup[x] for x in instance.tokens]
+        padding_mask = [1] * len(token_ids)
         segment_ids = list(instance.segment_ids)
-        assert len(input_ids) <= max_seq_length
+        assert len(token_ids) <= max_seq_length
 
-        while len(input_ids) < max_seq_length:
-            input_ids.append(0)
-            input_mask.append(0)
+        while len(token_ids) < max_seq_length:
+            token_ids.append(0)
+            padding_mask.append(0)
             segment_ids.append(0)
 
-        assert len(input_ids) == max_seq_length
-        assert len(input_mask) == max_seq_length
+        assert len(token_ids) == max_seq_length
+        assert len(padding_mask) == max_seq_length
         assert len(segment_ids) == max_seq_length
 
         masked_lm_positions = list(instance.masked_lm_positions)
@@ -151,8 +151,8 @@ def write_instance_to_example_files(
         next_sentence_label = 1 if instance.is_random_next else 0
 
         features = collections.OrderedDict()
-        features["input_ids"] = int_feature(input_ids)
-        features["input_mask"] = int_feature(input_mask)
+        features["token_ids"] = int_feature(token_ids)
+        features["padding_mask"] = int_feature(padding_mask)
         features["segment_ids"] = int_feature(segment_ids)
         features["masked_lm_positions"] = int_feature(masked_lm_positions)
         features["masked_lm_ids"] = int_feature(masked_lm_ids)
diff --git a/examples/bert/bert_train.py b/examples/bert/bert_train.py
index 3f30b004..9e870440 100644
--- a/examples/bert/bert_train.py
+++ b/examples/bert/bert_train.py
@@ -220,9 +220,9 @@ class BertPretrainingModel(keras.Model):
     def call(self, data):
         encoder_output = self.encoder(
             {
-                "input_ids": data["input_ids"],
-                "input_mask": data["input_mask"],
+                "token_ids": data["token_ids"],
                 "segment_ids": data["segment_ids"],
+                "padding_mask": data["padding_mask"],
             }
         )
         sequence_output, pooled_output = (
@@ -312,8 +312,8 @@ def decode_record(record):
     seq_length = PREPROCESSING_CONFIG["max_seq_length"]
     lm_length = PREPROCESSING_CONFIG["max_predictions_per_seq"]
     name_to_features = {
-        "input_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
-        "input_mask": tf.io.FixedLenFeature([seq_length], tf.int64),
+        "token_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
+        "padding_mask": tf.io.FixedLenFeature([seq_length], tf.int64),
         "segment_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
         "masked_lm_positions": tf.io.FixedLenFeature([lm_length], tf.int64),
         "masked_lm_ids": tf.io.FixedLenFeature([lm_length], tf.int64),
diff --git a/keras_nlp/models/bert.py b/keras_nlp/models/bert.py
index 25e672e8..41e6720e 100644
--- a/keras_nlp/models/bert.py
+++ b/keras_nlp/models/bert.py
@@ -33,15 +33,26 @@ checkpoints = {
         "uncased_en": {
             "md5": "9b2b2139f221988759ac9cdd17050b31",
             "description": "Base size of Bert where all input is lowercased. "
-            "Trained on English wikipedia + books corpora.",
+            "Trained on English Wikipedia + BooksCorpus.",
             "vocabulary_size": 30522,
         },
         "cased_en": {
             "md5": "f94a6cb012e18f4fb8ec92abb91864e9",
             "description": "Base size of Bert where case is maintained. "
-            "Trained on English wikipedia + books corpora.",
+            "Trained on English Wikipedia + BooksCorpus.",
             "vocabulary_size": 28996,
         },
+        "zh": {
+            "md5": "79afa421e386076e62ab42dad555ab0c",
+            "description": "Base size of BERT. Trained on Chinese Wikipedia.",
+            "vocabulary_size": 21128,
+        },
+        "multi_cased": {
+            "md5": "b0631cec0a1f2513c6cfd75ba29c33aa",
+            "description": "Base size of BERT. Trained on Wikipedias of "
+            "104 languages.",
+            "vocabulary_size": 119547,
+        },
     }
 }
 
diff --git a/keras_nlp/utils/text_generation.py b/keras_nlp/utils/text_generation.py
index afc20b84..b2206c5d 100644
--- a/keras_nlp/utils/text_generation.py
+++ b/keras_nlp/utils/text_generation.py
@@ -18,27 +18,16 @@ import tensorflow as tf
 from absl import logging
 from tensorflow import keras
 
-# TODO (@chenmoneygithub): Refactor code to reuse snippets.
 
-
-def validate_prompt(prompt):
+def _validate_prompt(prompt):
     """Helper function to validate input to text_generation utils."""
-    if isinstance(prompt, tf.RaggedTensor):
-        raise ValueError(
-            "RaggedTensor `prompt` is not supported, please "
-            "provide `prompt` as a list or Tensor."
-        )
-    if not isinstance(prompt, tf.Tensor):
+    if not isinstance(prompt, (tf.Tensor, tf.RaggedTensor)):
         prompt = tf.convert_to_tensor(prompt)
-    if prompt.shape[-1] == 0:
-        raise ValueError(
-            "Length of `prompt` is 0, please provide a non-empty `prompt`."
-        )
     return prompt
 
 
-def validate_token_probability_fn(token_probability_fn, prompt):
-    """Helper function to validate token probability fn output"""
+def _validate_token_probability_fn(token_probability_fn, prompt):
+    """Helper function to validate token probability fn output."""
     test_pred = token_probability_fn(prompt)
     if len(test_pred.shape) != 2:
         raise ValueError(
@@ -48,7 +37,45 @@ def validate_token_probability_fn(token_probability_fn, prompt):
         )
 
 
-def mask_tokens_after_end_token(prompt, max_length, end_token_id, pad_token_id):
+def _get_prompt_shape(prompt):
+    """Helper function to get the batch size and prompt length."""
+    if isinstance(prompt, tf.Tensor):
+        shape = tf.shape(prompt)
+        return (shape[0], shape[1])
+    elif isinstance(prompt, tf.RaggedTensor):
+        batch_size = prompt.nrows()
+        length = tf.math.reduce_min(tf.RaggedTensor.row_lengths(prompt))
+        return (batch_size, length)
+
+
+def _pad_prompt(prompt, max_length):
+    """Pad prompt to `max_length` and compute a mask for controlled updates.
+
+    This utility will pad the (possibly ragged) prompt to `max_length`, and
+    compute a mask where the input was originally set in the prompt, to avoid
+    overwriting the original inputs when generating token(s) for the next
+    timestep.
+    """
+    if isinstance(prompt, tf.Tensor):
+        shape = tf.shape(prompt)
+        pad_shape = (shape[0], max_length - shape[1])
+
+        mask = tf.ones(shape, tf.bool)
+        mask = tf.concat((mask, tf.zeros(pad_shape, tf.bool)), axis=1)
+        prompt = tf.concat((prompt, tf.zeros(pad_shape, prompt.dtype)), axis=1)
+    elif isinstance(prompt, tf.RaggedTensor):
+        # TODO: `to_tensor()` works with `jit_compile = True` in TF 2.8.x but
+        # fails in TF 2.9.x. Fix this. After this issue has been fixed, we can
+        # condense the two branches into one by starting off with a ragged tensor.
+        mask = tf.ones_like(prompt, dtype=tf.bool)
+        mask = mask.to_tensor(shape=(None, max_length))
+        prompt = prompt.to_tensor(shape=(None, max_length))
+    return prompt, mask
+
+
+def _mask_tokens_after_end_token(
+    prompt, max_length, end_token_id, pad_token_id
+):
     """Helper function to mask the tokens after the end token."""
     # Mask out tokens after `end_token_id` is encountered.
     # Find index of first end_token_id.
@@ -134,40 +161,36 @@ def greedy_search(
     ```
 
     """
-    prompt = validate_prompt(prompt)
+    prompt = _validate_prompt(prompt)
 
     input_is_1d = prompt.shape.rank == 1
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
-    validate_token_probability_fn(token_probability_fn, prompt)
 
-    shape = tf.shape(prompt)
-    batch_size = shape[0]
-    length = shape[1]
+    _validate_token_probability_fn(token_probability_fn, prompt)
 
-    # Pad the prompt with `pad_token_id` to `max_length`.
-    padding = tf.fill((batch_size, max_length - length), pad_token_id)
-    prompt = tf.concat((prompt, padding), axis=1)
+    batch_size, length = _get_prompt_shape(prompt)
+    prompt, mask = _pad_prompt(prompt, max_length)
 
     def one_step(length, prompt):
         pred = token_probability_fn(prompt[:, :length])
         next_token = tf.cast(tf.argmax(pred, axis=-1), dtype=prompt.dtype)
+        next_token = tf.where(mask[:, length], prompt[:, length], next_token)
 
         # Append the next token to current sequence.
-        def add_token(args):
-            sequence, token = args
-            return tf.tensor_scatter_nd_update(
-                tensor=sequence, indices=[[length]], updates=[token]
-            )
-
-        prompt = tf.map_fn(
-            fn=add_token,
-            elems=(prompt, next_token),
-            fn_output_signature=tf.TensorSpec(
-                shape=(max_length), dtype=prompt.dtype
+        prompt = tf.tensor_scatter_nd_update(
+            tensor=prompt,
+            indices=tf.stack(
+                (
+                    tf.cast(tf.range(batch_size), dtype=length.dtype),
+                    tf.repeat(length, batch_size),
+                ),
+                axis=1,
             ),
+            updates=next_token,
         )
-        length += 1
+
+        length = tf.add(length, 1)
         return (length, prompt)
 
     # Run a while loop till text of length `max_length` has been generated.
@@ -178,7 +201,7 @@ def greedy_search(
     )
 
     if end_token_id is not None:
-        prompt = mask_tokens_after_end_token(
+        prompt = _mask_tokens_after_end_token(
             prompt, max_length, end_token_id, pad_token_id
         )
 
@@ -277,12 +300,17 @@ def beam_search(
             f"`num_beams` should be strictly positive. Received: `num_beams={num_beams}`."
         )
 
-    prompt = validate_prompt(prompt)
+    prompt = _validate_prompt(prompt)
+
+    if prompt.shape[-1] == 0:
+        raise ValueError(
+            "Length of `prompt` is 0, please provide a non-empty `prompt`."
+        )
 
     input_is_1d = prompt.shape.rank == 1
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
-    validate_token_probability_fn(token_probability_fn, prompt)
+    _validate_token_probability_fn(token_probability_fn, prompt)
 
     batch_size, length = prompt.shape
     if length >= max_length:
@@ -325,7 +353,7 @@ def beam_search(
     prompt = tf.squeeze(max_beams)
 
     if end_token_id is not None:
-        prompt = mask_tokens_after_end_token(
+        prompt = _mask_tokens_after_end_token(
             prompt, max_length, end_token_id, pad_token_id
         )
     return tf.squeeze(prompt) if input_is_1d else prompt
@@ -407,22 +435,19 @@ def random_search(
     ```
 
     """
-    prompt = validate_prompt(prompt)
+    prompt = _validate_prompt(prompt)
+
     input_is_1d = prompt.shape.rank == 1
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
-    validate_token_probability_fn(token_probability_fn, prompt)
 
-    shape = tf.shape(prompt)
-    batch_size = shape[0]
-    length = shape[1]
+    _validate_token_probability_fn(token_probability_fn, prompt)
 
-    # Pad the prompt with `pad_token_id` to `max_length`.
-    padding = tf.fill((batch_size, max_length - length), pad_token_id)
-    prompt = tf.concat((prompt, padding), axis=1)
+    batch_size, length = _get_prompt_shape(prompt)
+    prompt, mask = _pad_prompt(prompt, max_length)
 
     def one_step(length, prompt):
-        pred = token_probability_fn(prompt)
+        pred = token_probability_fn(prompt[:, :length])
         if from_logits:
             pred = keras.activations.softmax(pred, axis=-1)
         next_token = tf.squeeze(
@@ -432,22 +457,22 @@ def random_search(
             ),
             axis=1,
         )
+        next_token = tf.where(mask[:, length], prompt[:, length], next_token)
 
         # Append the next token to current sequence.
-        def add_token(args):
-            sequence, token = args
-            return tf.tensor_scatter_nd_update(
-                tensor=sequence, indices=[[length]], updates=[token]
-            )
-
-        prompt = tf.map_fn(
-            fn=add_token,
-            elems=(prompt, next_token),
-            fn_output_signature=tf.TensorSpec(
-                shape=(max_length), dtype=prompt.dtype
+        prompt = tf.tensor_scatter_nd_update(
+            tensor=prompt,
+            indices=tf.stack(
+                (
+                    tf.cast(tf.range(batch_size), dtype=length.dtype),
+                    tf.repeat(length, batch_size),
+                ),
+                axis=1,
             ),
+            updates=next_token,
         )
-        length += 1
+
+        length = tf.add(length, 1)
         return (length, prompt)
 
     # Run a while loop till text of length `max_length` has been generated.
@@ -458,9 +483,10 @@ def random_search(
     )
 
     if end_token_id is not None:
-        prompt = mask_tokens_after_end_token(
+        prompt = _mask_tokens_after_end_token(
             prompt, max_length, end_token_id, pad_token_id
         )
+
     return tf.squeeze(prompt) if input_is_1d else prompt
 
 
@@ -544,14 +570,17 @@ def top_k_search(
     ```
 
     """
-    prompt = validate_prompt(prompt)
+    if k <= 0:
+        raise ValueError(f"`k` should be strictly positive. Received: `k={k}`.")
+
+    prompt = _validate_prompt(prompt)
+
     input_is_1d = prompt.shape.rank == 1
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
-    validate_token_probability_fn(token_probability_fn, prompt)
 
-    if k <= 0:
-        raise ValueError(f"`k` should be strictly positive. Received: `k={k}`.")
+    _validate_token_probability_fn(token_probability_fn, prompt)
+
     # If k is greater than the vocabulary size, use the entire vocabulary.
     pred = token_probability_fn(prompt)
     if k > pred.shape[1]:
@@ -561,16 +590,11 @@ def top_k_search(
         )
         k = pred.shape[1]
 
-    shape = tf.shape(prompt)
-    batch_size = shape[0]
-    length = shape[1]
-
-    # Pad the prompt with `pad_token_id` to `max_length`.
-    padding = tf.fill((batch_size, max_length - length), pad_token_id)
-    prompt = tf.concat((prompt, padding), axis=1)
+    batch_size, length = _get_prompt_shape(prompt)
+    prompt, mask = _pad_prompt(prompt, max_length)
 
     def one_step(length, prompt):
-        pred = token_probability_fn(prompt)
+        pred = token_probability_fn(prompt[:, :length])
         if from_logits:
             pred = keras.activations.softmax(pred, axis=-1)
 
@@ -584,22 +608,22 @@ def top_k_search(
         # Rearrange to get the next token idx from the original order.
         next_token = tf.gather_nd(top_k_indices, next_token, batch_dims=1)
         next_token = tf.cast(next_token, dtype=prompt.dtype)
+        next_token = tf.where(mask[:, length], prompt[:, length], next_token)
 
         # Append the next token to current sequence.
-        def add_token(args):
-            sequence, token = args
-            return tf.tensor_scatter_nd_update(
-                tensor=sequence, indices=[[length]], updates=[token]
-            )
-
-        prompt = tf.map_fn(
-            fn=add_token,
-            elems=(prompt, next_token),
-            fn_output_signature=tf.TensorSpec(
-                shape=(max_length), dtype=prompt.dtype
+        prompt = tf.tensor_scatter_nd_update(
+            tensor=prompt,
+            indices=tf.stack(
+                (
+                    tf.cast(tf.range(batch_size), dtype=length.dtype),
+                    tf.repeat(length, batch_size),
+                ),
+                axis=1,
             ),
+            updates=next_token,
         )
-        length += 1
+
+        length = tf.add(length, 1)
         return (length, prompt)
 
     # Run a while loop till text of length `max_length` has been generated.
@@ -610,9 +634,10 @@ def top_k_search(
     )
 
     if end_token_id is not None:
-        prompt = mask_tokens_after_end_token(
+        prompt = _mask_tokens_after_end_token(
             prompt, max_length, end_token_id, pad_token_id
         )
+
     return tf.squeeze(prompt) if input_is_1d else prompt
 
 
@@ -703,22 +728,19 @@ def top_p_search(
             f"`p` should be in the range (0, 1). Received: `p={p}`."
         )
 
-    prompt = validate_prompt(prompt)
+    prompt = _validate_prompt(prompt)
+
     input_is_1d = prompt.shape.rank == 1
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
-    validate_token_probability_fn(token_probability_fn, prompt)
 
-    shape = tf.shape(prompt)
-    batch_size = shape[0]
-    length = shape[1]
+    _validate_token_probability_fn(token_probability_fn, prompt)
 
-    # Pad the prompt with `pad_token_id` to `max_length`.
-    padding = tf.fill((batch_size, max_length - length), pad_token_id)
-    prompt = tf.concat((prompt, padding), axis=1)
+    batch_size, length = _get_prompt_shape(prompt)
+    prompt, mask = _pad_prompt(prompt, max_length)
 
     def one_step(length, prompt):
-        pred = token_probability_fn(prompt)
+        pred = token_probability_fn(prompt[:, :length])
         if from_logits:
             pred = keras.activations.softmax(pred, axis=-1)
         # Sort preds in descending order.
@@ -746,22 +768,22 @@ def top_p_search(
             sorted_indices, sorted_next_token, batch_dims=1
         )
         next_token = tf.cast(next_token, dtype=prompt.dtype)
+        next_token = tf.where(mask[:, length], prompt[:, length], next_token)
 
         # Append the next token to current sequence.
-        def add_token(args):
-            sequence, token = args
-            return tf.tensor_scatter_nd_update(
-                tensor=sequence, indices=[[length]], updates=[token]
-            )
-
-        prompt = tf.map_fn(
-            fn=add_token,
-            elems=(prompt, next_token),
-            fn_output_signature=tf.TensorSpec(
-                shape=(max_length), dtype=prompt.dtype
+        prompt = tf.tensor_scatter_nd_update(
+            tensor=prompt,
+            indices=tf.stack(
+                (
+                    tf.cast(tf.range(batch_size), dtype=length.dtype),
+                    tf.repeat(length, batch_size),
+                ),
+                axis=1,
             ),
+            updates=next_token,
         )
-        length += 1
+
+        length = tf.add(length, 1)
         return (length, prompt)
 
     # Run a while loop till text of length `max_length` has been generated.
@@ -772,7 +794,8 @@ def top_p_search(
     )
 
     if end_token_id is not None:
-        prompt = mask_tokens_after_end_token(
+        prompt = _mask_tokens_after_end_token(
             prompt, max_length, end_token_id, pad_token_id
         )
+
     return tf.squeeze(prompt) if input_is_1d else prompt
diff --git a/tools/checkpoint_conversion/bert_base_cased.ipynb b/tools/checkpoint_conversion/bert_base_cased.ipynb
index ce805b77..878d8fc9 100644
--- a/tools/checkpoint_conversion/bert_base_cased.ipynb
+++ b/tools/checkpoint_conversion/bert_base_cased.ipynb
@@ -425,9 +425,9 @@
             "__________________________________________________________________________________________________\n",
             " Layer (type)                   Output Shape         Param #     Connected to                     \n",
             "==================================================================================================\n",
-            " input_ids (InputLayer)         [(None, None)]       0           []                               \n",
+            " token_ids (InputLayer)         [(None, None)]       0           []                               \n",
             "                                                                                                  \n",
-            " token_embedding (Embedding)    (None, None, 768)    22268928    ['input_ids[0][0]']              \n",
+            " token_embedding (Embedding)    (None, None, 768)    22268928    ['token_ids[0][0]']              \n",
             "                                                                                                  \n",
             " segment_ids (InputLayer)       [(None, None)]       0           []                               \n",
             "                                                                                                  \n",
@@ -445,43 +445,43 @@
             "                                                                                                  \n",
             " embeddings_dropout (Dropout)   (None, None, 768)    0           ['embeddings_layer_norm[0][0]']  \n",
             "                                                                                                  \n",
-            " input_mask (InputLayer)        [(None, None)]       0           []                               \n",
+            " padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " merEncoder)                                                      'padding_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",
+            " merEncoder)                                                      'padding_mask[0][0]']             \n",
             "                                                                                                  \n",
             " tf.__operators__.getitem (Slic  (None, 768)         0           ['transformer_layer_11[0][0]']   \n",
             " ingOpLambda)                                                                                     \n",
@@ -666,15 +666,15 @@
       "outputs": [],
       "source": [
         "keras_nlp_output = model({\n",
-        "    \"input_ids\": token_ids,\n",
+        "    \"token_ids\": token_ids,\n",
         "    \"segment_ids\": segment_ids,\n",
-        "    \"input_mask\": token_ids != 0,\n",
+        "    \"padding_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",
+        "    \"padding_mask\": token_ids != 0,\n",
         "})[\"pooled_output\"]"
       ]
     },
@@ -812,9 +812,9 @@
       "source": [
         "# Same output from loaded checkpoint\n",
         "keras_nlp_output2 = model2({\n",
-        "    \"input_ids\": token_ids,\n",
+        "    \"token_ids\": token_ids,\n",
         "    \"segment_ids\": segment_ids,\n",
-        "    \"input_mask\": token_ids != 0,\n",
+        "    \"padding_mask\": token_ids != 0,\n",
         "})[\"pooled_output\"]\n",
         "tf.reduce_mean(keras_nlp_output - keras_nlp_output2)"
       ]
@@ -935,9 +935,9 @@
       "source": [
         "# Same output from cloud model\n",
         "keras_nlp_output_cloud = model_cloud({\n",
-        "    \"input_ids\": token_ids,\n",
+        "    \"token_ids\": token_ids,\n",
         "    \"segment_ids\": segment_ids,\n",
-        "    \"input_mask\": token_ids != 0,\n",
+        "    \"padding_mask\": token_ids != 0,\n",
         "})[\"pooled_output\"]\n",
         "tf.reduce_mean(keras_nlp_output - keras_nlp_output_cloud)"
       ]
diff --git a/tools/checkpoint_conversion/bert_base_multi_cased.ipynb b/tools/checkpoint_conversion/bert_base_multi_cased.ipynb
new file mode 100644
index 00000000..249402df
--- /dev/null
+++ b/tools/checkpoint_conversion/bert_base_multi_cased.ipynb
@@ -0,0 +1,1002 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github"
+      },
+      "source": [
+        "<a href=\"https://github.com/abheesht17/keras-nlp/blob/bert-base-chinese/tools/checkpoint_conversion/bert_base_multi_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": "52ac4b95-b8c9-4c2d-b4ca-eb57505d2cae"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "\u001b[K     |████████████████████████████████| 511.7 MB 6.3 kB/s \n",
+            "\u001b[K     |████████████████████████████████| 2.1 MB 42.7 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 4.6 MB 49.4 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 438 kB 69.8 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 5.8 MB 48.8 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 1.6 MB 57.7 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 1.1 MB 72.3 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 43 kB 2.2 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 99 kB 10.6 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 116 kB 63.0 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 1.3 MB 64.7 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 352 kB 63.4 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 238 kB 100.5 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 636 kB 73.7 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/abheesht17/keras-nlp.git@bert-base-chinese tensorflow tf-models-official --upgrade --quiet"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {
+        "id": "JsbnAdSz5DzZ"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import os\n",
+        "\n",
+        "import keras_nlp\n",
+        "import tensorflow as tf\n",
+        "import tensorflow_models as tfm\n",
+        "from tensorflow import keras"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "id": "DmVlNiSexzR7"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_TYPE = \"bert_base\"\n",
+        "MODEL_SUFFIX = \"multi_cased\"\n",
+        "MODEL_NAME = f\"{MODEL_TYPE}_{MODEL_SUFFIX}\"\n",
+        "VOCAB_SIZE = 119547"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "L3j5PFBt5JeR"
+      },
+      "source": [
+        "## Load the model garden checkpoints and weights"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "JdXFWsMVEf-x",
+        "outputId": "043339cd-3d95-4083-c915-f9632effc319"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Downloading data from https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/multi_cased_L-12_H-768_A-12.tar.gz\n",
+            "660198400/660198400 [==============================] - 5s 0us/step\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Model garden BERT paths.\n",
+        "zip_path = f\"\"\"https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/{MODEL_SUFFIX}_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": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "LHYiSsvYtfEU",
+        "outputId": "c185b640-d358-4d25-c249-7ac5ec4234df"
+      },
+      "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 = os.path.join(extract_dir, \"vocab.txt\")\n",
+        "checkpoint_path = os.path.join(extract_dir, \"bert_model.ckpt\")\n",
+        "config_path = os.path.join(extract_dir, \"bert_config.json\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 7,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "vCUJk59B2rai",
+        "outputId": "d1f58eaa-a601-4bf1-bd1f-0b5ad1a28a28"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "_CHECKPOINTABLE_OBJECT_GRAPH []\n",
+            "encoder/layer_with_weights-0/embeddings/.ATTRIBUTES/VARIABLE_VALUE [119547, 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": "937d8db5-4567-4fdc-e4c8-3a14c1f50b36"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Model: \"bert_custom\"\n",
+            "__________________________________________________________________________________________________\n",
+            " Layer (type)                   Output Shape         Param #     Connected to                     \n",
+            "==================================================================================================\n",
+            " token_ids (InputLayer)         [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " token_embedding (Embedding)    (None, None, 768)    91812096    ['token_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",
+            " padding_mask (InputLayer)      [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " transformer_layer_0 (Transform  (None, None, 768)   7087872     ['embeddings_dropout[0][0]',     \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_1 (Transform  (None, None, 768)   7087872     ['transformer_layer_0[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_2 (Transform  (None, None, 768)   7087872     ['transformer_layer_1[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_3 (Transform  (None, None, 768)   7087872     ['transformer_layer_2[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_4 (Transform  (None, None, 768)   7087872     ['transformer_layer_3[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_5 (Transform  (None, None, 768)   7087872     ['transformer_layer_4[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_6 (Transform  (None, None, 768)   7087872     ['transformer_layer_5[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_7 (Transform  (None, None, 768)   7087872     ['transformer_layer_6[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_8 (Transform  (None, None, 768)   7087872     ['transformer_layer_7[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_9 (Transform  (None, None, 768)   7087872     ['transformer_layer_8[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_10 (Transfor  (None, None, 768)   7087872     ['transformer_layer_9[0][0]',    \n",
+            " merEncoder)                                                      'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_11 (Transfor  (None, None, 768)   7087872     ['transformer_layer_10[0][0]',   \n",
+            " merEncoder)                                                      'padding_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: 177,853,440\n",
+            "Trainable params: 177,853,440\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[f\"encoder/layer_with_weights-{model.num_layers + 4}/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "model.get_layer(\"pooled_dense\").bias.assign(\n",
+        "    weights[f\"encoder/layer_with_weights-{model.num_layers + 4}/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",
+        "        lowercase=False\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 झटपट brown लोमड़ी.\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 11,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "chg-o813CJNJ",
+        "outputId": "9cdd88a8-efc8-43bb-c469-9496578f46d6"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7f8f003c50d0>"
+            ]
+          },
+          "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",
+        "    \"token_ids\": token_ids,\n",
+        "    \"segment_ids\": segment_ids,\n",
+        "    \"padding_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": "9gF51d-CbwUh",
+        "outputId": "39dc40e7-7e9d-4a70-a55e-2ecc8d3e91ce"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(10,), dtype=float32, numpy=\n",
+              "array([ 0.23481919, -0.09908431,  0.191012  , -0.09323437,  0.0487478 ,\n",
+              "        0.2899589 ,  0.08575539,  0.19417824, -0.25255474,  0.22692412],\n",
+              "      dtype=float32)>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 13
+        }
+      ],
+      "source": [
+        "keras_nlp_output[0, :10]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 14,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "zmhZYi-1YMGH",
+        "outputId": "48a0565d-30bb-4db0-c658-a4a4af7dc599"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(10,), dtype=float32, numpy=\n",
+              "array([ 0.23481807, -0.09908444,  0.19100997, -0.09323332,  0.04874839,\n",
+              "        0.28995794,  0.08575849,  0.19417666, -0.2525542 ,  0.22692516],\n",
+              "      dtype=float32)>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 14
+        }
+      ],
+      "source": [
+        "mg_output[0, :10]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 15,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "TvNj8DFBYNPT",
+        "outputId": "12d8cb41-85f9-44a5-e990-9e497ac45fff"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(), dtype=float32, numpy=8.3201634e-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": "32cb37c4-344e-4b34-dab9-907f4f6b7216"
+      },
+      "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",
+        "    \"token_ids\": token_ids,\n",
+        "    \"segment_ids\": segment_ids,\n",
+        "    \"padding_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": "3289713c-5ba2-4c9a-d3fa-fefa52a6223a"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "764415"
+            ]
+          },
+          "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": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "lVutnU-hB6IQ",
+        "outputId": "d394313b-ccf0-4db5-b270-aee6f27d1f06"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "b0631cec0a1f2513c6cfd75ba29c33aa  bert_base_multi_cased.h5\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Get MD5 of model\n",
+        "!md5sum \"\"\"{MODEL_NAME}.h5\"\"\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "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": 20,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "wTd-5vUyVG0Q",
+        "outputId": "d66aca01-c1aa-404a-d2a1-090644f842bd"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Downloading data from https://storage.googleapis.com/keras-nlp/models/bert_base_multi_cased/model.h5\n",
+            "711647480/711647480 [==============================] - 8s 0us/step\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Check uploaded model once added to repo\n",
+        "model_cloud = keras_nlp.models.BertBase(weights=MODEL_SUFFIX)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 21,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "zs5x_f6GVdNY",
+        "outputId": "330f724f-29a4-4e90-d1b8-669501575d80"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(), dtype=float32, numpy=0.0>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 21
+        }
+      ],
+      "source": [
+        "# Same output from cloud model\n",
+        "keras_nlp_output_cloud = model_cloud({\n",
+        "    \"token_ids\": token_ids,\n",
+        "    \"segment_ids\": segment_ids,\n",
+        "    \"padding_mask\": token_ids != 0,\n",
+        "})[\"pooled_output\"]\n",
+        "tf.reduce_mean(keras_nlp_output - keras_nlp_output_cloud)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 22,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "RAwrhAcSzHWa",
+        "outputId": "57dd61f6-5415-47a6-af9b-01d082f6b745"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(10,), dtype=float32, numpy=\n",
+              "array([ 0.23481919, -0.09908431,  0.191012  , -0.09323437,  0.0487478 ,\n",
+              "        0.2899589 ,  0.08575539,  0.19417824, -0.25255474,  0.22692412],\n",
+              "      dtype=float32)>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 22
+        }
+      ],
+      "source": [
+        "keras_nlp_output_cloud[0, :10]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "UqSREee-sLmP"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [],
+      "provenance": []
+    },
+    "gpuClass": "standard",
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "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
index d8bf48b5..07206838 100644
--- a/tools/checkpoint_conversion/bert_base_uncased.ipynb
+++ b/tools/checkpoint_conversion/bert_base_uncased.ipynb
@@ -425,9 +425,9 @@
             "__________________________________________________________________________________________________\n",
             " Layer (type)                   Output Shape         Param #     Connected to                     \n",
             "==================================================================================================\n",
-            " input_ids (InputLayer)         [(None, None)]       0           []                               \n",
+            " token_ids (InputLayer)         [(None, None)]       0           []                               \n",
             "                                                                                                  \n",
-            " token_embedding (Embedding)    (None, None, 768)    23440896    ['input_ids[0][0]']              \n",
+            " token_embedding (Embedding)    (None, None, 768)    23440896    ['token_ids[0][0]']              \n",
             "                                                                                                  \n",
             " segment_ids (InputLayer)       [(None, None)]       0           []                               \n",
             "                                                                                                  \n",
@@ -445,43 +445,43 @@
             "                                                                                                  \n",
             " embeddings_dropout (Dropout)   (None, None, 768)    0           ['embeddings_layer_norm[0][0]']  \n",
             "                                                                                                  \n",
-            " input_mask (InputLayer)        [(None, None)]       0           []                               \n",
+            " padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " erEncoder)                                                       'padding_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",
+            " merEncoder)                                                      'padding_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",
+            " merEncoder)                                                      'padding_mask[0][0]']             \n",
             "                                                                                                  \n",
             " tf.__operators__.getitem (Slic  (None, 768)         0           ['transformer_layer_11[0][0]']   \n",
             " ingOpLambda)                                                                                     \n",
@@ -666,15 +666,15 @@
       "outputs": [],
       "source": [
         "keras_nlp_output = model({\n",
-        "    \"input_ids\": token_ids,\n",
+        "    \"token_ids\": token_ids,\n",
         "    \"segment_ids\": segment_ids,\n",
-        "    \"input_mask\": token_ids != 0,\n",
+        "    \"padding_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",
+        "    \"padding_mask\": token_ids != 0,\n",
         "})[\"pooled_output\"]"
       ]
     },
@@ -812,9 +812,9 @@
       "source": [
         "# Same output from loaded checkpoint\n",
         "keras_nlp_output2 = model2({\n",
-        "    \"input_ids\": token_ids,\n",
+        "    \"token_ids\": token_ids,\n",
         "    \"segment_ids\": segment_ids,\n",
-        "    \"input_mask\": token_ids != 0,\n",
+        "    \"padding_mask\": token_ids != 0,\n",
         "})[\"pooled_output\"]\n",
         "tf.reduce_mean(keras_nlp_output - keras_nlp_output2)"
       ]
@@ -935,9 +935,9 @@
       "source": [
         "# Same output from cloud model\n",
         "keras_nlp_output_cloud = model_cloud({\n",
-        "    \"input_ids\": token_ids,\n",
+        "    \"token_ids\": token_ids,\n",
         "    \"segment_ids\": segment_ids,\n",
-        "    \"input_mask\": token_ids != 0,\n",
+        "    \"padding_mask\": token_ids != 0,\n",
         "})[\"pooled_output\"]\n",
         "tf.reduce_mean(keras_nlp_output - keras_nlp_output_cloud)"
       ]
diff --git a/tools/checkpoint_conversion/bert_base_zh.ipynb b/tools/checkpoint_conversion/bert_base_zh.ipynb
new file mode 100644
index 00000000..94603a4c
--- /dev/null
+++ b/tools/checkpoint_conversion/bert_base_zh.ipynb
@@ -0,0 +1,1001 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github"
+      },
+      "source": [
+        "<a href=\"https://github.com/abheesht17/keras-nlp/blob/bert-base-chinese/tools/checkpoint_conversion/bert_base_zh.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": "0e8b9004-fe3f-4fb7-b905-cebb361887ab"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "\u001b[K     |████████████████████████████████| 511.7 MB 6.2 kB/s \n",
+            "\u001b[K     |████████████████████████████████| 2.1 MB 47.5 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 4.6 MB 43.2 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 5.8 MB 51.1 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 438 kB 72.3 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 1.6 MB 54.9 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 352 kB 71.1 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 43 kB 2.2 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 636 kB 60.0 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 1.1 MB 68.5 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 1.3 MB 67.5 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 238 kB 75.7 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 116 kB 71.3 MB/s \n",
+            "\u001b[K     |████████████████████████████████| 99 kB 12.0 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/abheesht17/keras-nlp.git@bert-base-chinese tensorflow tf-models-official --upgrade --quiet"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {
+        "id": "JsbnAdSz5DzZ"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import os\n",
+        "\n",
+        "import keras_nlp\n",
+        "import tensorflow as tf\n",
+        "import tensorflow_models as tfm\n",
+        "from tensorflow import keras"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "id": "DmVlNiSexzR7"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL_TYPE = \"bert_base\"\n",
+        "MODEL_SUFFIX = \"chinese\"\n",
+        "MODEL_NAME = f\"{MODEL_TYPE}_{MODEL_SUFFIX}\"\n",
+        "VOCAB_SIZE = 21128"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "L3j5PFBt5JeR"
+      },
+      "source": [
+        "## Load the model garden checkpoints and weights"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "JdXFWsMVEf-x",
+        "outputId": "e47ef71a-11c2-4c53-f398-ca3de4da77f9"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Downloading data from https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/chinese_L-12_H-768_A-12.tar.gz\n",
+            "379546341/379546341 [==============================] - 3s 0us/step\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Model garden BERT paths.\n",
+        "zip_path = f\"\"\"https://storage.googleapis.com/tf_model_garden/nlp/bert/v3/{MODEL_SUFFIX}_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": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "LHYiSsvYtfEU",
+        "outputId": "47518515-60ce-49f3-e6ac-39273b1d4616"
+      },
+      "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 = os.path.join(extract_dir, \"vocab.txt\")\n",
+        "checkpoint_path = os.path.join(extract_dir, \"bert_model.ckpt\")\n",
+        "config_path = os.path.join(extract_dir, \"bert_config.json\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 7,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "vCUJk59B2rai",
+        "outputId": "59d95000-83a8-48f8-870a-51c2dbb5db6c"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "_CHECKPOINTABLE_OBJECT_GRAPH []\n",
+            "encoder/layer_with_weights-0/embeddings/.ATTRIBUTES/VARIABLE_VALUE [21128, 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": "6b7cd86d-33af-4864-a3c7-63acf53fecc2"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Model: \"bert_custom\"\n",
+            "__________________________________________________________________________________________________\n",
+            " Layer (type)                   Output Shape         Param #     Connected to                     \n",
+            "==================================================================================================\n",
+            " token_ids (InputLayer)         [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " token_embedding (Embedding)    (None, None, 768)    16226304    ['token_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",
+            " padding_mask (InputLayer)      [(None, None)]       0           []                               \n",
+            "                                                                                                  \n",
+            " transformer_layer_0 (Transform  (None, None, 768)   7087872     ['embeddings_dropout[0][0]',     \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_1 (Transform  (None, None, 768)   7087872     ['transformer_layer_0[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_2 (Transform  (None, None, 768)   7087872     ['transformer_layer_1[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_3 (Transform  (None, None, 768)   7087872     ['transformer_layer_2[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_4 (Transform  (None, None, 768)   7087872     ['transformer_layer_3[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_5 (Transform  (None, None, 768)   7087872     ['transformer_layer_4[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_6 (Transform  (None, None, 768)   7087872     ['transformer_layer_5[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_7 (Transform  (None, None, 768)   7087872     ['transformer_layer_6[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_8 (Transform  (None, None, 768)   7087872     ['transformer_layer_7[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_9 (Transform  (None, None, 768)   7087872     ['transformer_layer_8[0][0]',    \n",
+            " erEncoder)                                                       'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_10 (Transfor  (None, None, 768)   7087872     ['transformer_layer_9[0][0]',    \n",
+            " merEncoder)                                                      'padding_mask[0][0]']           \n",
+            "                                                                                                  \n",
+            " transformer_layer_11 (Transfor  (None, None, 768)   7087872     ['transformer_layer_10[0][0]',   \n",
+            " merEncoder)                                                      'padding_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: 102,267,648\n",
+            "Trainable params: 102,267,648\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[f\"encoder/layer_with_weights-{model.num_layers + 4}/kernel/.ATTRIBUTES/VARIABLE_VALUE\"]\n",
+        ")\n",
+        "model.get_layer(\"pooled_dense\").bias.assign(\n",
+        "    weights[f\"encoder/layer_with_weights-{model.num_layers + 4}/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([\"敏捷的棕色狐狸\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 11,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "chg-o813CJNJ",
+        "outputId": "5c8144f8-4a07-4a6d-8c06-6bca7e8625cf"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7f7b120f1750>"
+            ]
+          },
+          "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",
+        "    \"token_ids\": token_ids,\n",
+        "    \"segment_ids\": segment_ids,\n",
+        "    \"padding_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": "9gF51d-CbwUh",
+        "outputId": "dc0fdef7-1f6c-4e7d-ac84-71066108846d"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(10,), dtype=float32, numpy=\n",
+              "array([ 0.99850804,  0.99867415,  0.9855039 ,  0.6341696 ,  0.87277335,\n",
+              "        0.31392148, -0.7766987 , -0.20329826,  0.93192303, -0.9924547 ],\n",
+              "      dtype=float32)>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 13
+        }
+      ],
+      "source": [
+        "keras_nlp_output[0, :10]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 14,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "zmhZYi-1YMGH",
+        "outputId": "893835bc-6438-45e5-c326-0ba631e1a4d6"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(10,), dtype=float32, numpy=\n",
+              "array([ 0.9985083 ,  0.9986742 ,  0.9855096 ,  0.6341621 ,  0.87276965,\n",
+              "        0.3139575 , -0.77669966, -0.20335098,  0.9319198 , -0.9924562 ],\n",
+              "      dtype=float32)>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 14
+        }
+      ],
+      "source": [
+        "mg_output[0, :10]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 15,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "TvNj8DFBYNPT",
+        "outputId": "e8094ce2-2cf5-4c6a-8dc3-474f20b9b58e"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(), dtype=float32, numpy=-3.5652678e-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": "3282584c-38a9-486c-a3bf-1589dc17083f"
+      },
+      "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",
+        "    \"token_ids\": token_ids,\n",
+        "    \"segment_ids\": segment_ids,\n",
+        "    \"padding_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": "8da86cef-d78f-4979-db6b-5fe8c9428c19"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "75770"
+            ]
+          },
+          "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": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "lVutnU-hB6IQ",
+        "outputId": "438c793e-bbc8-4c46-808c-a6bb69ef68e6"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "79afa421e386076e62ab42dad555ab0c  bert_base_chinese.h5\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Get MD5 of model\n",
+        "!md5sum \"\"\"{MODEL_NAME}.h5\"\"\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "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": 20,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "wTd-5vUyVG0Q",
+        "outputId": "5db37007-d7eb-4431-b182-8304425d8f71"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Downloading data from https://storage.googleapis.com/keras-nlp/models/bert_base_zh/model.h5\n",
+            "409304312/409304312 [==============================] - 3s 0us/step\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Check uploaded model once added to repo\n",
+        "model_cloud = keras_nlp.models.BertBase(weights=\"zh\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 21,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "zs5x_f6GVdNY",
+        "outputId": "ba0b05d4-f1f2-41d0-ab93-3fa5fe7ea2e5"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(), dtype=float32, numpy=0.0>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 21
+        }
+      ],
+      "source": [
+        "# Same output from cloud model\n",
+        "keras_nlp_output_cloud = model_cloud({\n",
+        "    \"token_ids\": token_ids,\n",
+        "    \"segment_ids\": segment_ids,\n",
+        "    \"padding_mask\": token_ids != 0,\n",
+        "})[\"pooled_output\"]\n",
+        "tf.reduce_mean(keras_nlp_output - keras_nlp_output_cloud)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 22,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "RAwrhAcSzHWa",
+        "outputId": "0f3528b5-262f-4736-e377-bd450ee586de"
+      },
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": [
+              "<tf.Tensor: shape=(10,), dtype=float32, numpy=\n",
+              "array([ 0.99850804,  0.99867415,  0.9855039 ,  0.6341696 ,  0.87277335,\n",
+              "        0.31392148, -0.7766987 , -0.20329826,  0.93192303, -0.9924547 ],\n",
+              "      dtype=float32)>"
+            ]
+          },
+          "metadata": {},
+          "execution_count": 22
+        }
+      ],
+      "source": [
+        "keras_nlp_output_cloud[0, :10]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "ae9jwK4xqZQf"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [],
+      "provenance": []
+    },
+    "gpuClass": "standard",
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "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 e86d92af..0fd607e5 100644
--- a/keras_nlp/models/bert_test.py
+++ b/keras_nlp/models/bert_test.py
@@ -16,12 +16,13 @@
 import os
 
 import tensorflow as tf
+from absl.testing import parameterized
 from tensorflow import keras
 
 from keras_nlp.models import bert
 
 
-class BertTest(tf.test.TestCase):
+class BertTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         self.model = bert.BertCustom(
             vocabulary_size=1000,
@@ -33,7 +34,7 @@ class BertTest(tf.test.TestCase):
             name="encoder",
         )
         self.batch_size = 8
-        self.input_data = {
+        self.input_batch = {
             "token_ids": tf.ones(
                 (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
@@ -45,8 +46,12 @@ class BertTest(tf.test.TestCase):
             ),
         }
 
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
     def test_valid_call_bert(self):
-        self.model(self.input_data)
+        self.model(self.input_batch)
 
     def test_variable_sequence_length_call_bert(self):
         for seq_length in (25, 50, 75):
@@ -65,7 +70,7 @@ class BertTest(tf.test.TestCase):
 
     def test_valid_call_classifier(self):
         classifier = bert.BertClassifier(self.model, 4, name="classifier")
-        classifier(self.input_data)
+        classifier(self.input_batch)
 
     def test_valid_call_bert_base(self):
         model = bert.BertBase(vocabulary_size=1000, name="encoder")
@@ -82,6 +87,38 @@ class BertTest(tf.test.TestCase):
         }
         model(input_data)
 
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_bert_base_compile(self, jit_compile):
+        model = bert.BertBase(vocabulary_size=1000, name="encoder")
+        model.compile(jit_compile=jit_compile)
+        model.predict(self.input_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_bert_base_compile_batched_ds(self, jit_compile):
+        model = bert.BertBase(vocabulary_size=1000, name="encoder")
+        model.compile(jit_compile=jit_compile)
+        model.predict(self.input_dataset)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_bert_classifier_compile(self, jit_compile):
+        model = bert.BertClassifier(self.model, 4, name="classifier")
+        model.compile(jit_compile=jit_compile)
+        model.predict(self.input_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_bert_classifier_compile_batched_ds(self, jit_compile):
+        model = bert.BertClassifier(self.model, 4, name="classifier")
+        model.compile(jit_compile=jit_compile)
+        model.predict(self.input_dataset)
+
     def test_bert_base_vocab_error(self):
         # Need `vocabulary_size` or `weights`
         with self.assertRaises(ValueError):
@@ -103,12 +140,12 @@ class BertTest(tf.test.TestCase):
             )
 
     def test_saving_model(self):
-        model_output = self.model(self.input_data)
+        model_output = self.model(self.input_batch)
         save_path = os.path.join(self.get_temp_dir(), "model")
         self.model.save(save_path)
         restored_model = keras.models.load_model(save_path)
 
-        restored_output = restored_model(self.input_data)
+        restored_output = restored_model(self.input_batch)
         self.assertAllClose(
             model_output["pooled_output"], restored_output["pooled_output"]
         )
diff --git a/keras_nlp/models/roberta_test.py b/keras_nlp/models/roberta_test.py
index 26fb65d7..f362de13 100644
--- a/keras_nlp/models/roberta_test.py
+++ b/keras_nlp/models/roberta_test.py
@@ -16,12 +16,13 @@
 import os
 
 import tensorflow as tf
+from absl.testing import parameterized
 from tensorflow import keras
 
 from keras_nlp.models import roberta
 
 
-class RobertaTest(tf.test.TestCase):
+class RobertaTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         self.model = roberta.RobertaCustom(
             vocabulary_size=1000,
@@ -33,7 +34,7 @@ class RobertaTest(tf.test.TestCase):
             name="encoder",
         )
         self.batch_size = 8
-        self.input_data = {
+        self.input_batch = {
             "token_ids": tf.ones(
                 (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
@@ -42,14 +43,18 @@ class RobertaTest(tf.test.TestCase):
             ),
         }
 
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
     def test_valid_call_roberta(self):
-        self.model(self.input_data)
+        self.model(self.input_batch)
 
     def test_valid_call_classifier(self):
         classifier = roberta.RobertaClassifier(
             self.model, 4, 128, name="classifier"
         )
-        classifier(self.input_data)
+        classifier(self.input_batch)
 
     def test_valid_call_roberta_base(self):
         model = roberta.RobertaBase(vocabulary_size=1000, name="encoder")
@@ -63,6 +68,38 @@ class RobertaTest(tf.test.TestCase):
         }
         model(input_data)
 
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_roberta_base_compile(self, jit_compile):
+        model = roberta.RobertaBase(vocabulary_size=1000, name="encoder")
+        model.compile(jit_compile=jit_compile)
+        model.predict(self.input_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_roberta_base_compile_batched_ds(self, jit_compile):
+        model = roberta.RobertaBase(vocabulary_size=1000, name="encoder")
+        model.compile(jit_compile=jit_compile)
+        model.predict(self.input_dataset)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_roberta_classifier_compile(self, jit_compile):
+        model = roberta.RobertaClassifier(self.model, 4, 128, name="classifier")
+        model.compile(jit_compile=jit_compile)
+        model.predict(self.input_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_roberta_classifier_compile_batched_ds(self, jit_compile):
+        model = roberta.RobertaClassifier(self.model, 4, 128, name="classifier")
+        model.compile(jit_compile=jit_compile)
+        model.predict(self.input_dataset)
+
     def test_variable_sequence_length_call_roberta(self):
         for seq_length in (25, 50, 75):
             input_data = {
@@ -80,13 +117,13 @@ class RobertaTest(tf.test.TestCase):
             )
 
     def test_saving_model(self):
-        model_output = self.model.predict(self.input_data)
+        model_output = self.model.predict(self.input_batch)
 
         save_path = os.path.join(self.get_temp_dir(), "model")
         self.model.save(save_path)
         restored_model = keras.models.load_model(save_path)
 
-        restored_output = restored_model.predict(self.input_data)
+        restored_output = restored_model.predict(self.input_batch)
         self.assertAllClose(
             model_output,
             restored_output,
diff --git a/keras_nlp/utils/text_generation_test.py b/keras_nlp/utils/text_generation_test.py
index fab6ed8f..2c96f281 100644
--- a/keras_nlp/utils/text_generation_test.py
+++ b/keras_nlp/utils/text_generation_test.py
@@ -51,14 +51,6 @@ class GreedySearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
         self.token_probability_fn = token_probability_fn
 
-    def test_generate_with_empty_prompt(self):
-        inputs = tf.constant([])
-        with self.assertRaises(ValueError):
-            greedy_search(self.token_probability_fn, inputs, max_length=5)
-        inputs = tf.constant([[]])
-        with self.assertRaises(ValueError):
-            greedy_search(self.token_probability_fn, inputs, max_length=5)
-
     def test_generate_with_1d_prompt(self):
         inputs = tf.constant([1])
         outputs = greedy_search(self.token_probability_fn, inputs, max_length=5)
@@ -75,9 +67,13 @@ class GreedySearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
         self.assertEquals(outputs.shape, [2, 5])
 
     def test_generate_with_ragged_prompt(self):
-        inputs = tf.ragged.constant([[1], [2, 3]])
-        with self.assertRaises(ValueError):
-            greedy_search(self.token_probability_fn, inputs, max_length=5)
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        inputs = tf.ragged.constant([[1], [2, 1, 2]])
+        outputs = greedy_search(token_probability_fn, inputs, max_length=5)
+        self.assertEquals(outputs.shape, [2, 5])
 
     def test_assert_generation_is_correct(self):
         def token_probability_fn(inputs):
@@ -115,9 +111,11 @@ class GreedySearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(outputs, expected_outputs)
 
     @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
+        ("dense_jit_compile_false", False, False),
+        ("dense_jit_compile_true", True, False),
+        ("ragged_jit_compile_false", False, True),
     )
-    def test_model_compile(self, jit_compile):
+    def test_model_compile(self, jit_compile, ragged):
         def token_probability_fn(inputs):
             prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, 2, axis=0)
@@ -135,9 +133,12 @@ class GreedySearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
                 )
                 return generated
 
-        inputs = tf.constant([[0, 1], [1, 2]])
-        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
-        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+        if ragged:
+            inputs = tf.ragged.constant([[0, 1], [1, 1, 2]])
+            expected_outputs = tf.constant([[0, 1, 3, 3, 3], [1, 1, 2, 0, 0]])
+        else:
+            inputs = tf.constant([[0, 1], [1, 2]])
+            expected_outputs = [[0, 1, 3, 3, 3], [1, 2, 0, 0, 0]]
 
         model = TestModel()
         model.compile(jit_compile=jit_compile)
@@ -146,9 +147,11 @@ class GreedySearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(outputs, expected_outputs)
 
     @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
+        ("dense_jit_compile_false", False, False),
+        ("dense_jit_compile_true", True, False),
+        ("ragged_jit_compile_false", False, True),
     )
-    def test_model_compile_batched_ds(self, jit_compile):
+    def test_model_compile_batched_ds(self, jit_compile, ragged):
         def token_probability_fn(inputs):
             prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, 2, axis=0)
@@ -166,10 +169,14 @@ class GreedySearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
                 )
                 return generated
 
-        inputs = tf.constant([[0, 1], [1, 2]])
+        if ragged:
+            inputs = tf.ragged.constant([[0, 1], [1, 1, 2]])
+            expected_outputs = tf.constant([[0, 1, 3, 3, 3], [1, 1, 2, 0, 0]])
+        else:
+            inputs = tf.constant([[0, 1], [1, 2]])
+            expected_outputs = [[0, 1, 3, 3, 3], [1, 2, 0, 0, 0]]
+
         ds = tf.data.Dataset.from_tensor_slices(inputs).batch(2)
-        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
-        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
 
         model = TestModel()
         model.compile(jit_compile=jit_compile)
@@ -375,14 +382,6 @@ class RandomSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
         self.token_probability_fn = token_probability_fn
 
-    def test_generate_with_empty_prompt(self):
-        inputs = tf.constant([])
-        with self.assertRaises(ValueError):
-            random_search(self.token_probability_fn, inputs, max_length=5)
-        inputs = tf.constant([[]])
-        with self.assertRaises(ValueError):
-            random_search(self.token_probability_fn, inputs, max_length=5)
-
     def test_generate_with_1d_prompt(self):
         inputs = tf.constant([1])
         outputs = random_search(self.token_probability_fn, inputs, max_length=5)
@@ -399,9 +398,13 @@ class RandomSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
         self.assertEquals(outputs.shape, [2, 5])
 
     def test_generate_with_ragged_prompt(self):
-        inputs = tf.ragged.constant([[1], [2, 3]])
-        with self.assertRaises(ValueError):
-            random_search(self.token_probability_fn, inputs, max_length=5)
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        inputs = tf.ragged.constant([[1], [2, 1, 2]])
+        outputs = random_search(token_probability_fn, inputs, max_length=5)
+        self.assertEquals(outputs.shape, [2, 5])
 
     def test_assert_probability_distribution_generation_is_correct(self):
         def token_probability_fn(inputs):
@@ -451,9 +454,11 @@ class RandomSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(outputs, expected_outputs)
 
     @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
+        ("dense_jit_compile_false", False, False),
+        ("dense_jit_compile_true", True, False),
+        ("ragged_jit_compile_false", False, True),
     )
-    def test_model_compile(self, jit_compile):
+    def test_model_compile(self, jit_compile, ragged):
         def token_probability_fn(inputs):
             prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, 2, axis=0)
@@ -472,9 +477,12 @@ class RandomSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
                 )
                 return generated
 
-        inputs = tf.constant([[0, 1], [1, 2]])
-        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
-        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+        if ragged:
+            inputs = tf.ragged.constant([[0, 1], [1, 1, 2]])
+            expected_outputs = tf.constant([[0, 1, 3, 3, 3], [1, 1, 2, 0, 0]])
+        else:
+            inputs = tf.constant([[0, 1], [1, 2]])
+            expected_outputs = [[0, 1, 3, 3, 3], [1, 2, 0, 0, 0]]
 
         model = TestModel()
         model.compile(jit_compile=jit_compile)
@@ -483,9 +491,11 @@ class RandomSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(outputs, expected_outputs)
 
     @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
+        ("dense_jit_compile_false", False, False),
+        ("dense_jit_compile_true", True, False),
+        ("ragged_jit_compile_false", False, True),
     )
-    def test_model_compile_batched_ds(self, jit_compile):
+    def test_model_compile_batched_ds(self, jit_compile, ragged):
         def token_probability_fn(inputs):
             prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, 2, axis=0)
@@ -504,10 +514,14 @@ class RandomSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
                 )
                 return generated
 
-        inputs = tf.constant([[0, 1], [1, 2]])
+        if ragged:
+            inputs = tf.ragged.constant([[0, 1], [1, 1, 2]])
+            expected_outputs = tf.constant([[0, 1, 3, 3, 3], [1, 1, 2, 0, 0]])
+        else:
+            inputs = tf.constant([[0, 1], [1, 2]])
+            expected_outputs = [[0, 1, 3, 3, 3], [1, 2, 0, 0, 0]]
+
         ds = tf.data.Dataset.from_tensor_slices(inputs).batch(2)
-        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
-        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
 
         model = TestModel()
         model.compile(jit_compile=jit_compile)
@@ -573,14 +587,6 @@ class TopKSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
         self.token_probability_fn = token_probability_fn
 
-    def test_generate_with_empty_prompt(self):
-        inputs = tf.constant([])
-        with self.assertRaises(ValueError):
-            top_k_search(self.token_probability_fn, inputs, max_length=5, k=2)
-        inputs = tf.constant([[]])
-        with self.assertRaises(ValueError):
-            top_k_search(self.token_probability_fn, inputs, max_length=5, k=2)
-
     def test_generate_with_1d_prompt(self):
         inputs = tf.constant([1])
         outputs = top_k_search(
@@ -613,9 +619,13 @@ class TopKSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
         self.assertEquals(outputs.shape, [2, 5])
 
     def test_generate_with_ragged_prompt(self):
-        inputs = tf.ragged.constant([[1], [2, 3]])
-        with self.assertRaises(ValueError):
-            top_k_search(self.token_probability_fn, inputs, max_length=5, k=2)
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        inputs = tf.ragged.constant([[1], [2, 1, 2]])
+        outputs = top_k_search(token_probability_fn, inputs, max_length=5, k=2)
+        self.assertEquals(outputs.shape, [2, 5])
 
     def test_assert_probability_distribution_generation_is_correct(self):
         def token_probability_fn(inputs):
@@ -689,9 +699,11 @@ class TopKSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(outputs, expected_outputs)
 
     @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
+        ("dense_jit_compile_false", False, False),
+        ("dense_jit_compile_true", True, False),
+        ("ragged_jit_compile_false", False, True),
     )
-    def test_model_compile(self, jit_compile):
+    def test_model_compile(self, jit_compile, ragged):
         def token_probability_fn(inputs):
             prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, 2, axis=0)
@@ -700,19 +712,23 @@ class TopKSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
         class TestModel(tf.keras.Model):
             def call(self, inputs):
-                generated = random_search(
+                generated = top_k_search(
                     token_probability_fn,
                     inputs,
                     max_length=max_length,
+                    k=2,
                     seed=42,
                     end_token_id=2,
                     pad_token_id=0,
                 )
                 return generated
 
-        inputs = tf.constant([[0, 1], [1, 2]])
-        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
-        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+        if ragged:
+            inputs = tf.ragged.constant([[0, 1], [1, 1, 2]])
+            expected_outputs = tf.constant([[0, 1, 3, 3, 3], [1, 1, 2, 0, 0]])
+        else:
+            inputs = tf.constant([[0, 1], [1, 2]])
+            expected_outputs = [[0, 1, 3, 3, 3], [1, 2, 0, 0, 0]]
 
         model = TestModel()
         model.compile(jit_compile=jit_compile)
@@ -721,9 +737,11 @@ class TopKSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(outputs, expected_outputs)
 
     @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
+        ("dense_jit_compile_false", False, False),
+        ("dense_jit_compile_true", True, False),
+        ("ragged_jit_compile_false", False, True),
     )
-    def test_model_compile_batched_ds(self, jit_compile):
+    def test_model_compile_batched_ds(self, jit_compile, ragged):
         def token_probability_fn(inputs):
             prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, 2, axis=0)
@@ -732,20 +750,25 @@ class TopKSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
         class TestModel(tf.keras.Model):
             def call(self, inputs):
-                generated = random_search(
+                generated = top_k_search(
                     token_probability_fn,
                     inputs,
                     max_length=max_length,
+                    k=2,
                     seed=42,
                     end_token_id=2,
                     pad_token_id=0,
                 )
                 return generated
 
-        inputs = tf.constant([[0, 1], [1, 2]])
+        if ragged:
+            inputs = tf.ragged.constant([[0, 1], [1, 1, 2]])
+            expected_outputs = tf.constant([[0, 1, 3, 3, 3], [1, 1, 2, 0, 0]])
+        else:
+            inputs = tf.constant([[0, 1], [1, 2]])
+            expected_outputs = [[0, 1, 3, 3, 3], [1, 2, 0, 0, 0]]
+
         ds = tf.data.Dataset.from_tensor_slices(inputs).batch(2)
-        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
-        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
 
         model = TestModel()
         model.compile(jit_compile=jit_compile)
@@ -813,14 +836,6 @@ class TopPSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
         self.token_probability_fn = token_probability_fn
 
-    def test_generate_with_empty_prompt(self):
-        inputs = tf.constant([])
-        with self.assertRaises(ValueError):
-            top_p_search(self.token_probability_fn, inputs, max_length=5, p=0.8)
-        inputs = tf.constant([[]])
-        with self.assertRaises(ValueError):
-            top_p_search(self.token_probability_fn, inputs, max_length=5, p=0.8)
-
     def test_generate_with_1d_prompt(self):
         inputs = tf.constant([1])
         outputs = top_p_search(
@@ -843,9 +858,15 @@ class TopPSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
         self.assertEquals(outputs.shape, [2, 5])
 
     def test_generate_with_ragged_prompt(self):
-        inputs = tf.ragged.constant([[1], [2, 3]])
-        with self.assertRaises(ValueError):
-            top_p_search(self.token_probability_fn, inputs, max_length=5, p=0.8)
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        inputs = tf.ragged.constant([[1], [2, 1, 2]])
+        outputs = top_p_search(
+            token_probability_fn, inputs, max_length=5, p=0.8
+        )
+        self.assertEquals(outputs.shape, [2, 5])
 
     def test_assert_probability_distribution_generation_is_correct(self):
         def token_probability_fn(inputs):
@@ -915,9 +936,11 @@ class TopPSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(outputs, expected_outputs)
 
     @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
+        ("dense_jit_compile_false", False, False),
+        ("dense_jit_compile_true", True, False),
+        ("ragged_jit_compile_false", False, True),
     )
-    def test_model_compile(self, jit_compile):
+    def test_model_compile(self, jit_compile, ragged):
         def token_probability_fn(inputs):
             prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, 2, axis=0)
@@ -937,9 +960,12 @@ class TopPSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
                 )
                 return generated
 
-        inputs = tf.constant([[0, 1], [1, 2]])
-        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
-        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+        if ragged:
+            inputs = tf.ragged.constant([[0, 1], [1, 1, 2]])
+            expected_outputs = tf.constant([[0, 1, 3, 3, 3], [1, 1, 2, 0, 0]])
+        else:
+            inputs = tf.constant([[0, 1], [1, 2]])
+            expected_outputs = [[0, 1, 3, 3, 3], [1, 2, 0, 0, 0]]
 
         model = TestModel()
         model.compile(jit_compile=jit_compile)
@@ -948,9 +974,11 @@ class TopPSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
         self.assertAllEqual(outputs, expected_outputs)
 
     @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
+        ("dense_jit_compile_false", False, False),
+        ("dense_jit_compile_true", True, False),
+        ("ragged_jit_compile_false", False, True),
     )
-    def test_model_compile_batched_ds(self, jit_compile):
+    def test_model_compile_batched_ds(self, jit_compile, ragged):
         def token_probability_fn(inputs):
             prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, 2, axis=0)
@@ -970,10 +998,14 @@ class TopPSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
                 )
                 return generated
 
-        inputs = tf.constant([[0, 1], [1, 2]])
+        if ragged:
+            inputs = tf.ragged.constant([[0, 1], [1, 1, 2]])
+            expected_outputs = tf.constant([[0, 1, 3, 3, 3], [1, 1, 2, 0, 0]])
+        else:
+            inputs = tf.constant([[0, 1], [1, 2]])
+            expected_outputs = [[0, 1, 3, 3, 3], [1, 2, 0, 0, 0]]
+
         ds = tf.data.Dataset.from_tensor_slices(inputs).batch(2)
-        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
-        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
 
         model = TestModel()
         model.compile(jit_compile=jit_compile)
diff --git a/network_tests/test_load_ckpts.py b/network_tests/test_load_ckpts.py
index ddb251f9..9af16c0b 100644
--- a/network_tests/test_load_ckpts.py
+++ b/network_tests/test_load_ckpts.py
@@ -14,24 +14,20 @@
 """Tests for loading pretrained model checkpoints."""
 
 import tensorflow as tf
+from absl.testing import parameterized
 
 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 = {
-            "token_ids": tf.random.uniform(
-                shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
-            ),
-            "segment_ids": tf.constant([0] * 200 + [1] * 312, shape=(1, 512)),
-            "padding_mask": tf.constant([1] * 512, shape=(1, 512)),
-        }
-        model(input_data)
-
-    def test_load_bert_base_cased(self):
-        model = keras_nlp.models.BertBase(weights="cased_en")
+class BertCkptTest(tf.test.TestCase, parameterized.TestCase):
+    @parameterized.named_parameters(
+        ("uncased_en", "uncased_en"),
+        ("cased_en", "cased_en"),
+        ("zh", "zh"),
+        ("multi_cased", "multi_cased"),
+    )
+    def test_load_bert_base(self, weights):
+        model = keras_nlp.models.BertBase(weights=weights)
         input_data = {
             "token_ids": tf.random.uniform(
                 shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size