Repository: keras-team/keras-nlp

Issue:
Incorrect `epsilon` (LayerNorm) for the Transformer Layer in All Backbones
Similar to https://github.com/keras-team/keras-nlp/issues/642. Please check whether the correct epsilon value has been passed. Do this for all backbones.

Discussion:
abheesht17: LGTM!
TheAthleticCoder: @mattdangerw A gentle reminder regarding the commit
mattdangerw: @TheAthleticCoder thanks! This looks good, but would it be possible to drop a comment here with links to where you found these constants? Would be a useful reference anyway.
TheAthleticCoder: @mattdangerw The values of the constants were already there in the model files (model_backbone.py). I used those pre-existing constants and passed them into the main model architecture.  

Patch:
diff --git a/keras_nlp/models/albert/albert_backbone.py b/keras_nlp/models/albert/albert_backbone.py
index c3609f91..14840987 100644
--- a/keras_nlp/models/albert/albert_backbone.py
+++ b/keras_nlp/models/albert/albert_backbone.py
@@ -190,6 +190,7 @@ class AlbertBackbone(Backbone):
                         x, approximate=True
                     ),
                     dropout=dropout,
+                    layer_norm_epsilon=1e-12,
                     kernel_initializer=albert_kernel_initializer(),
                     name=f"group_{group_idx}_inner_layer_{inner_idx}",
                 )
diff --git a/keras_nlp/models/bart/bart_backbone.py b/keras_nlp/models/bart/bart_backbone.py
index a66a6210..9256c11c 100644
--- a/keras_nlp/models/bart/bart_backbone.py
+++ b/keras_nlp/models/bart/bart_backbone.py
@@ -156,8 +156,8 @@ class BartBackbone(Backbone):
                 activation=lambda x: keras.activations.gelu(
                     x, approximate=False
                 ),
-                layer_norm_epsilon=1e-5,
                 dropout=dropout,
+                layer_norm_epsilon=1e-5,
                 kernel_initializer=bart_kernel_initializer(),
                 name=f"transformer_encoder_layer_{i}",
             )(x, padding_mask=encoder_padding_mask)
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_backbone.py b/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
index eaac7dd9..ab8ebd21 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
@@ -155,6 +155,7 @@ class DebertaV3Backbone(Backbone):
                 activation=lambda x: keras.activations.gelu(
                     x, approximate=False
                 ),
+                layer_norm_epsilon=1e-7,
                 kernel_initializer=deberta_kernel_initializer(),
                 name=f"disentangled_attention_encoder_layer_{i}",
             )(
diff --git a/keras_nlp/models/distil_bert/distil_bert_backbone.py b/keras_nlp/models/distil_bert/distil_bert_backbone.py
index 20ef77f6..5c614fb9 100644
--- a/keras_nlp/models/distil_bert/distil_bert_backbone.py
+++ b/keras_nlp/models/distil_bert/distil_bert_backbone.py
@@ -141,6 +141,7 @@ class DistilBertBackbone(Backbone):
                 intermediate_dim=intermediate_dim,
                 activation="gelu",
                 dropout=dropout,
+                layer_norm_epsilon=1e-12,
                 kernel_initializer=distilbert_kernel_initializer(),
                 name=f"transformer_layer_{i}",
             )(x, padding_mask=padding_mask)
diff --git a/keras_nlp/models/f_net/f_net_backbone.py b/keras_nlp/models/f_net/f_net_backbone.py
index 403912c2..4a0eeb17 100644
--- a/keras_nlp/models/f_net/f_net_backbone.py
+++ b/keras_nlp/models/f_net/f_net_backbone.py
@@ -159,8 +159,8 @@ class FNetBackbone(Backbone):
                 activation=lambda x: keras.activations.gelu(
                     x, approximate=True
                 ),
-                layer_norm_epsilon=1e-12,
                 dropout=dropout,
+                layer_norm_epsilon=1e-12,
                 kernel_initializer=f_net_kernel_initializer(),
                 bias_initializer=f_net_bias_initializer(),
                 name=f"f_net_layer_{i}",
diff --git a/keras_nlp/models/gpt2/gpt2_backbone.py b/keras_nlp/models/gpt2/gpt2_backbone.py
index f47a14af..29959478 100644
--- a/keras_nlp/models/gpt2/gpt2_backbone.py
+++ b/keras_nlp/models/gpt2/gpt2_backbone.py
@@ -137,10 +137,10 @@ class GPT2Backbone(Backbone):
                 intermediate_dim=intermediate_dim,
                 num_heads=num_heads,
                 dropout=dropout,
+                layer_norm_epsilon=1e-05,
                 activation=lambda x: keras.activations.gelu(
                     x, approximate=True
                 ),
-                layer_norm_epsilon=1e-05,
                 kernel_initializer=_gpt_2_kernel_initializer(stddev=0.02),
                 normalize_first=True,
                 name=f"transformer_layer_{i}",
diff --git a/keras_nlp/models/roberta/roberta_backbone.py b/keras_nlp/models/roberta/roberta_backbone.py
index 926f3157..cf702354 100644
--- a/keras_nlp/models/roberta/roberta_backbone.py
+++ b/keras_nlp/models/roberta/roberta_backbone.py
@@ -135,6 +135,7 @@ class RobertaBackbone(Backbone):
                 intermediate_dim=intermediate_dim,
                 activation="gelu",
                 dropout=dropout,
+                layer_norm_epsilon=1e-5,
                 kernel_initializer=roberta_kernel_initializer(),
                 name=f"transformer_layer_{i}",
             )(x, padding_mask=padding_mask)
diff --git a/keras_nlp/models/roberta/roberta_tokenizer.py b/keras_nlp/models/roberta/roberta_tokenizer.py
index c5ba13b5..0e3284fa 100644
--- a/keras_nlp/models/roberta/roberta_tokenizer.py
+++ b/keras_nlp/models/roberta/roberta_tokenizer.py
@@ -105,7 +105,7 @@ class RobertaTokenizer(BytePairTokenizer):
         pad_token = "<pad>"
         end_token = "
"
         mask_token = "<mask>"
-        for token in [start_token, pad_token, end_token]:
+        for token in [start_token, pad_token, end_token, mask_token]:
             if token not in self.get_vocabulary():
                 raise ValueError(
                     f"Cannot find token `'{token}'` in the provided "
diff --git a/keras_nlp/samplers/sampler.py b/keras_nlp/samplers/sampler.py
index dde87a66..b9a8dbaa 100644
--- a/keras_nlp/samplers/sampler.py
+++ b/keras_nlp/samplers/sampler.py
@@ -189,14 +189,18 @@ class Sampler:
 
     def _mask_tokens_after_end_token(
         self,
-        prompt,
+        generated_result,
+        original_padding_mask,
         max_length,
         end_token_id,
     ):
         """Helper function to truncate the tokens after the end token."""
-        # Mask out tokens after `end_token_id` is encountered.
+        # Create a tensor with True for each end_token_id.
+        end_tokens = generated_result == end_token_id
+        # Remove all end_token_ids in the original input.
+        end_tokens = end_tokens & (original_padding_mask == tf.constant(False))
         # Find index of first end_token_id.
-        end_indices = tf.math.argmax(prompt == end_token_id, -1)
+        end_indices = tf.math.argmax(end_tokens, -1)
         # Use max_length if no `end_token_id` is found.
         end_indices = tf.where(
             end_indices == 0,
@@ -205,7 +209,7 @@ class Sampler:
         )
         # Truncate out tokens after (including) the end token.
         mask_indices = tf.sequence_mask(end_indices, maxlen=max_length)
-        return tf.ragged.boolean_mask(prompt, mask_indices)
+        return tf.ragged.boolean_mask(generated_result, mask_indices)
 
     def __call__(
         self,
@@ -217,7 +221,6 @@ class Sampler:
         from_logits=True,
     ):
         prompt, mask = self._validate_prompt_and_mask(prompt, mask)
-
         input_is_1d = prompt.shape.rank == 1
         if input_is_1d:
             prompt = tf.RaggedTensor.from_tensor(prompt[tf.newaxis, :])
@@ -228,6 +231,7 @@ class Sampler:
         # static shape, which means we cannot concatenate generated token to
         # current prompt.
         prompt, mask = self._pad_prompt(prompt, max_length)
+        original_padding_mask = tf.identity(mask)
         self._validate_token_probability_fn(token_probability_fn, prompt, mask)
 
         # Convert `sample` method to a `tf.function` if `self.run_eagerly=False`
@@ -247,6 +251,7 @@ class Sampler:
         if end_token_id is not None:
             prompt = self._mask_tokens_after_end_token(
                 prompt,
+                original_padding_mask,
                 max_length,
                 end_token_id,
             )
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index ed593a32..c5e0e578 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -331,7 +331,7 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
         self.oov_token = oov_token
 
         if oov_token not in self.vocabulary:
-            raise RuntimeError(
+            raise ValueError(
                 f'Cannot find `oov_token="{self.oov_token}"` in the '
                 "vocabulary.\n"
                 "You can either update the vocabulary to include "
diff --git a/tools/checkpoint_conversion/convert_albert_checkpoints.py b/tools/checkpoint_conversion/convert_albert_checkpoints.py
new file mode 100644
index 00000000..180478b0
--- /dev/null
+++ b/tools/checkpoint_conversion/convert_albert_checkpoints.py
@@ -0,0 +1,319 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import shutil
+
+import numpy as np
+import tensorflow as tf
+import transformers
+from absl import app
+from absl import flags
+
+import keras_nlp
+from tools.checkpoint_conversion.checkpoint_conversion_utils import (
+    get_md5_checksum,
+)
+
+PRESET_MAP = {
+    "albert_base_en_uncased": "albert-base-v2",
+    "albert_large_en_uncased": "albert-large-v2",
+    "albert_extra_large_en_uncased": "albert-xlarge-v2",
+    "albert_extra_extra_large_en_uncased": "albert-xxlarge-v2",
+}
+
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string(
+    "preset", None, f'Must be one of {",".join(PRESET_MAP.keys())}'
+)
+
+
+def convert_checkpoints(hf_model):
+    print("\n-> Convert original weights to KerasNLP format.")
+
+    print("\n-> Load KerasNLP model.")
+    keras_nlp_model = keras_nlp.models.AlbertBackbone.from_preset(
+        FLAGS.preset, load_weights=False
+    )
+
+    hf_wts = hf_model.state_dict()
+    print("Original weights:")
+    print(list(hf_wts.keys()))
+
+    num_heads = keras_nlp_model.num_heads
+    hidden_dim = keras_nlp_model.hidden_dim
+
+    keras_nlp_model.get_layer("token_embedding").embeddings.assign(
+        hf_wts["embeddings.word_embeddings.weight"]
+    )
+    keras_nlp_model.get_layer("position_embedding").position_embeddings.assign(
+        hf_wts["embeddings.position_embeddings.weight"]
+    )
+    keras_nlp_model.get_layer("segment_embedding").embeddings.assign(
+        hf_wts["embeddings.token_type_embeddings.weight"]
+    )
+
+    keras_nlp_model.get_layer("embeddings_layer_norm").gamma.assign(
+        hf_wts["embeddings.LayerNorm.weight"]
+    )
+    keras_nlp_model.get_layer("embeddings_layer_norm").beta.assign(
+        hf_wts["embeddings.LayerNorm.bias"]
+    )
+
+    keras_nlp_model.get_layer("embedding_projection").kernel.assign(
+        hf_wts["encoder.embedding_hidden_mapping_in.weight"].T
+    )
+    keras_nlp_model.get_layer("embedding_projection").bias.assign(
+        hf_wts["encoder.embedding_hidden_mapping_in.bias"]
+    )
+
+    for i in range(keras_nlp_model.num_groups):
+        for j in range(keras_nlp_model.num_inner_repetitions):
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._self_attention_layer._query_dense.kernel.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.attention.query.weight"
+                ]
+                .transpose(1, 0)
+                .reshape((hidden_dim, num_heads, -1))
+                .numpy()
+            )
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._self_attention_layer._query_dense.bias.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.attention.query.bias"
+                ]
+                .reshape((num_heads, -1))
+                .numpy()
+            )
+
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._self_attention_layer._key_dense.kernel.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.attention.key.weight"
+                ]
+                .transpose(1, 0)
+                .reshape((hidden_dim, num_heads, -1))
+                .numpy()
+            )
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._self_attention_layer._key_dense.bias.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.attention.key.bias"
+                ]
+                .reshape((num_heads, -1))
+                .numpy()
+            )
+
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._self_attention_layer._value_dense.kernel.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.attention.value.weight"
+                ]
+                .transpose(1, 0)
+                .reshape((hidden_dim, num_heads, -1))
+                .numpy()
+            )
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._self_attention_layer._value_dense.bias.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.attention.value.bias"
+                ]
+                .reshape((num_heads, -1))
+                .numpy()
+            )
+
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._self_attention_layer._output_dense.kernel.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.attention.dense.weight"
+                ]
+                .transpose(1, 0)
+                .reshape((num_heads, -1, hidden_dim))
+                .numpy()
+            )
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._self_attention_layer._output_dense.bias.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.attention.dense.bias"
+                ].numpy()
+            )
+
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._self_attention_layernorm.gamma.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.attention.LayerNorm.weight"
+                ].numpy()
+            )
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._self_attention_layernorm.beta.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.attention.LayerNorm.bias"
+                ].numpy()
+            )
+
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._feedforward_intermediate_dense.kernel.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.ffn.weight"
+                ]
+                .transpose(1, 0)
+                .numpy()
+            )
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._feedforward_intermediate_dense.bias.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.ffn.bias"
+                ].numpy()
+            )
+
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._feedforward_output_dense.kernel.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.ffn_output.weight"
+                ]
+                .transpose(1, 0)
+                .numpy()
+            )
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._feedforward_output_dense.bias.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.ffn_output.bias"
+                ].numpy()
+            )
+
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._feedforward_layernorm.gamma.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.full_layer_layer_norm.weight"
+                ].numpy()
+            )
+            keras_nlp_model.get_layer(
+                f"group_{i}_inner_layer_{j}"
+            )._feedforward_layernorm.beta.assign(
+                hf_wts[
+                    f"encoder.albert_layer_groups.{i}.albert_layers.{j}.full_layer_layer_norm.bias"
+                ].numpy()
+            )
+
+    keras_nlp_model.get_layer("pooled_dense").kernel.assign(
+        hf_wts["pooler.weight"].transpose(1, 0).numpy()
+    )
+    keras_nlp_model.get_layer("pooled_dense").bias.assign(
+        hf_wts["pooler.bias"].numpy()
+    )
+
+    # Save the model.
+    print("\n-> Save KerasNLP model weights.")
+    keras_nlp_model.save_weights(os.path.join(FLAGS.preset, "model.h5"))
+
+    return keras_nlp_model
+
+
+def extract_vocab(hf_tokenizer):
+    spm_path = os.path.join(FLAGS.preset, "spiece.model")
+    print(f"\n-> Save KerasNLP SPM vocabulary file to `{spm_path}`.")
+
+    shutil.copyfile(
+        transformers.utils.hub.get_file_from_repo(
+            hf_tokenizer.name_or_path, "spiece.model"
+        ),
+        spm_path,
+    )
+
+    keras_nlp_tokenizer = keras_nlp.models.AlbertTokenizer(
+        proto=spm_path,
+    )
+    keras_nlp_preprocessor = keras_nlp.models.AlbertPreprocessor(
+        keras_nlp_tokenizer
+    )
+
+    print("-> Print MD5 checksum of the vocab files.")
+    print(f"`{spm_path}` md5sum: ", get_md5_checksum(spm_path))
+
+    return keras_nlp_preprocessor
+
+
+def check_output(
+    keras_nlp_preprocessor,
+    keras_nlp_model,
+    hf_tokenizer,
+    hf_model,
+):
+    print("\n-> Check the outputs.")
+    sample_text = ["cricket is awesome, easily the best sport in the world!"]
+
+    # KerasNLP
+    keras_nlp_inputs = keras_nlp_preprocessor(tf.constant(sample_text))
+    keras_nlp_output = keras_nlp_model.predict(keras_nlp_inputs)[
+        "sequence_output"
+    ]
+
+    # HF
+    hf_inputs = hf_tokenizer(
+        sample_text, padding="max_length", return_tensors="pt"
+    )
+    hf_output = hf_model(**hf_inputs).last_hidden_state
+
+    print("KerasNLP output:", keras_nlp_output[0, 0, :10])
+    print("HF output:", hf_output[0, 0, :10])
+    print("Difference:", np.mean(keras_nlp_output - hf_output.detach().numpy()))
+
+    # Show the MD5 checksum of the model weights.
+    print(
+        "Model md5sum: ",
+        get_md5_checksum(os.path.join(FLAGS.preset, "model.h5")),
+    )
+
+
+def main(_):
+    os.makedirs(FLAGS.preset)
+
+    hf_model_name = PRESET_MAP[FLAGS.preset]
+
+    print("\n-> Load HF model and HF tokenizer.")
+    hf_model = transformers.AutoModel.from_pretrained(hf_model_name)
+    hf_model.eval()
+    hf_tokenizer = transformers.AutoTokenizer.from_pretrained(hf_model_name)
+
+    keras_nlp_model = convert_checkpoints(hf_model)
+    print("\n -> Load KerasNLP preprocessor.")
+    keras_nlp_preprocessor = extract_vocab(hf_tokenizer)
+
+    check_output(
+        keras_nlp_preprocessor,
+        keras_nlp_model,
+        hf_tokenizer,
+        hf_model,
+    )
+
+
+if __name__ == "__main__":
+    flags.mark_flag_as_required("preset")
+    app.run(main)
diff --git a/tools/checkpoint_conversion/convert_bart_checkpoints.py b/tools/checkpoint_conversion/convert_bart_checkpoints.py
new file mode 100644
index 00000000..54d54b96
--- /dev/null
+++ b/tools/checkpoint_conversion/convert_bart_checkpoints.py
@@ -0,0 +1,531 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import shutil
+
+import numpy as np
+import tensorflow as tf
+import transformers
+from absl import app
+from absl import flags
+
+import keras_nlp
+from tools.checkpoint_conversion.checkpoint_conversion_utils import (
+    get_md5_checksum,
+)
+
+PRESET_MAP = {
+    "bart_base_en": "facebook/bart-base",
+    "bart_large_en": "facebook/bart-large",
+}
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string(
+    "preset", None, f'Must be one of {",".join(PRESET_MAP.keys())}'
+)
+
+
+def convert_checkpoints(hf_model):
+    print("\n-> Convert original weights to KerasNLP format.")
+
+    print("\n-> Load KerasNLP model.")
+    keras_nlp_model = keras_nlp.models.BartBackbone.from_preset(
+        FLAGS.preset, load_weights=False
+    )
+
+    hf_wts = hf_model.state_dict()
+    print("Original weights:")
+    print(list(hf_wts.keys()))
+
+    hidden_dim = keras_nlp_model.hidden_dim
+    num_heads = keras_nlp_model.num_heads
+
+    # Token embedding weights shared by encoder and decoder.
+    keras_nlp_model.get_layer("token_embedding").embeddings.assign(
+        hf_wts["shared.weight"]
+    )
+
+    # Encoder weights.
+    keras_nlp_model.get_layer(
+        "encoder_position_embedding"
+    ).position_embeddings.assign(hf_wts["encoder.embed_positions.weight"][2:])
+
+    keras_nlp_model.get_layer("encoder_embeddings_layer_norm").gamma.assign(
+        hf_wts["encoder.layernorm_embedding.weight"]
+    )
+    keras_nlp_model.get_layer("encoder_embeddings_layer_norm").beta.assign(
+        hf_wts["encoder.layernorm_embedding.bias"]
+    )
+
+    for i in range(keras_nlp_model.num_layers):
+        # Self-attention.
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._self_attention_layer._query_dense.kernel.assign(
+            hf_wts[f"encoder.layers.{i}.self_attn.q_proj.weight"]
+            .transpose(1, 0)
+            .reshape((hidden_dim, num_heads, -1))
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._self_attention_layer._query_dense.bias.assign(
+            hf_wts[f"encoder.layers.{i}.self_attn.q_proj.bias"]
+            .reshape((num_heads, -1))
+            .numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._self_attention_layer._key_dense.kernel.assign(
+            hf_wts[f"encoder.layers.{i}.self_attn.k_proj.weight"]
+            .transpose(1, 0)
+            .reshape((hidden_dim, num_heads, -1))
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._self_attention_layer._key_dense.bias.assign(
+            hf_wts[f"encoder.layers.{i}.self_attn.k_proj.bias"]
+            .reshape((num_heads, -1))
+            .numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._self_attention_layer._value_dense.kernel.assign(
+            hf_wts[f"encoder.layers.{i}.self_attn.v_proj.weight"]
+            .transpose(1, 0)
+            .reshape((hidden_dim, num_heads, -1))
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._self_attention_layer._value_dense.bias.assign(
+            hf_wts[f"encoder.layers.{i}.self_attn.v_proj.bias"]
+            .reshape((num_heads, -1))
+            .numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._self_attention_layer._output_dense.kernel.assign(
+            hf_wts[f"encoder.layers.{i}.self_attn.out_proj.weight"]
+            .transpose(1, 0)
+            .reshape((num_heads, -1, hidden_dim))
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._self_attention_layer._output_dense.bias.assign(
+            hf_wts[f"encoder.layers.{i}.self_attn.out_proj.bias"].numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._self_attention_layernorm.gamma.assign(
+            hf_wts[f"encoder.layers.{i}.self_attn_layer_norm.weight"].numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._self_attention_layernorm.beta.assign(
+            hf_wts[f"encoder.layers.{i}.self_attn_layer_norm.bias"].numpy()
+        )
+
+        # Post self-attention layers.
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._feedforward_intermediate_dense.kernel.assign(
+            hf_wts[f"encoder.layers.{i}.fc1.weight"].transpose(1, 0).numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._feedforward_intermediate_dense.bias.assign(
+            hf_wts[f"encoder.layers.{i}.fc1.bias"].numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._feedforward_output_dense.kernel.assign(
+            hf_wts[f"encoder.layers.{i}.fc2.weight"].transpose(1, 0).numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._feedforward_output_dense.bias.assign(
+            hf_wts[f"encoder.layers.{i}.fc2.bias"].numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._feedforward_layernorm.gamma.assign(
+            hf_wts[f"encoder.layers.{i}.final_layer_norm.weight"].numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_encoder_layer_{i}"
+        )._feedforward_layernorm.beta.assign(
+            hf_wts[f"encoder.layers.{i}.final_layer_norm.bias"].numpy()
+        )
+
+    # Decoder weights.
+
+    keras_nlp_model.get_layer(
+        "decoder_position_embedding"
+    ).position_embeddings.assign(hf_wts["decoder.embed_positions.weight"][2:])
+
+    keras_nlp_model.get_layer("decoder_embeddings_layer_norm").gamma.assign(
+        hf_wts["decoder.layernorm_embedding.weight"]
+    )
+    keras_nlp_model.get_layer("decoder_embeddings_layer_norm").beta.assign(
+        hf_wts["decoder.layernorm_embedding.bias"]
+    )
+
+    for i in range(keras_nlp_model.num_layers):
+        # Self-attention.
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._self_attention_layer._query_dense.kernel.assign(
+            hf_wts[f"decoder.layers.{i}.self_attn.q_proj.weight"]
+            .transpose(1, 0)
+            .reshape((hidden_dim, num_heads, -1))
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._self_attention_layer._query_dense.bias.assign(
+            hf_wts[f"decoder.layers.{i}.self_attn.q_proj.bias"]
+            .reshape((num_heads, -1))
+            .numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._self_attention_layer._key_dense.kernel.assign(
+            hf_wts[f"decoder.layers.{i}.self_attn.k_proj.weight"]
+            .transpose(1, 0)
+            .reshape((hidden_dim, num_heads, -1))
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._self_attention_layer._key_dense.bias.assign(
+            hf_wts[f"decoder.layers.{i}.self_attn.k_proj.bias"]
+            .reshape((num_heads, -1))
+            .numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._self_attention_layer._value_dense.kernel.assign(
+            hf_wts[f"decoder.layers.{i}.self_attn.v_proj.weight"]
+            .transpose(1, 0)
+            .reshape((hidden_dim, num_heads, -1))
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._self_attention_layer._value_dense.bias.assign(
+            hf_wts[f"decoder.layers.{i}.self_attn.v_proj.bias"]
+            .reshape((num_heads, -1))
+            .numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._self_attention_layer._output_dense.kernel.assign(
+            hf_wts[f"decoder.layers.{i}.self_attn.out_proj.weight"]
+            .transpose(1, 0)
+            .reshape((num_heads, -1, hidden_dim))
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._self_attention_layer._output_dense.bias.assign(
+            hf_wts[f"decoder.layers.{i}.self_attn.out_proj.bias"].numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._self_attention_layernorm.gamma.assign(
+            hf_wts[f"decoder.layers.{i}.self_attn_layer_norm.weight"].numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._self_attention_layernorm.beta.assign(
+            hf_wts[f"decoder.layers.{i}.self_attn_layer_norm.bias"].numpy()
+        )
+
+        # Cross-attention.
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._cross_attention_layer._query_dense.kernel.assign(
+            hf_wts[f"decoder.layers.{i}.encoder_attn.q_proj.weight"]
+            .transpose(1, 0)
+            .reshape((hidden_dim, num_heads, -1))
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._cross_attention_layer._query_dense.bias.assign(
+            hf_wts[f"decoder.layers.{i}.encoder_attn.q_proj.bias"]
+            .reshape((num_heads, -1))
+            .numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._cross_attention_layer._key_dense.kernel.assign(
+            hf_wts[f"decoder.layers.{i}.encoder_attn.k_proj.weight"]
+            .transpose(1, 0)
+            .reshape((hidden_dim, num_heads, -1))
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._cross_attention_layer._key_dense.bias.assign(
+            hf_wts[f"decoder.layers.{i}.encoder_attn.k_proj.bias"]
+            .reshape((num_heads, -1))
+            .numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._cross_attention_layer._value_dense.kernel.assign(
+            hf_wts[f"decoder.layers.{i}.encoder_attn.v_proj.weight"]
+            .transpose(1, 0)
+            .reshape((hidden_dim, num_heads, -1))
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._cross_attention_layer._value_dense.bias.assign(
+            hf_wts[f"decoder.layers.{i}.encoder_attn.v_proj.bias"]
+            .reshape((num_heads, -1))
+            .numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._cross_attention_layer._output_dense.kernel.assign(
+            hf_wts[f"decoder.layers.{i}.encoder_attn.out_proj.weight"]
+            .transpose(1, 0)
+            .reshape((num_heads, -1, hidden_dim))
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._cross_attention_layer._output_dense.bias.assign(
+            hf_wts[f"decoder.layers.{i}.encoder_attn.out_proj.bias"].numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._cross_attention_layernorm.gamma.assign(
+            hf_wts[f"decoder.layers.{i}.encoder_attn_layer_norm.weight"].numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._cross_attention_layernorm.beta.assign(
+            hf_wts[f"decoder.layers.{i}.encoder_attn_layer_norm.bias"].numpy()
+        )
+
+        # Post self-attention and cross-attention layers.
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._feedforward_intermediate_dense.kernel.assign(
+            hf_wts[f"decoder.layers.{i}.fc1.weight"].transpose(1, 0).numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._feedforward_intermediate_dense.bias.assign(
+            hf_wts[f"decoder.layers.{i}.fc1.bias"].numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._feedforward_output_dense.kernel.assign(
+            hf_wts[f"decoder.layers.{i}.fc2.weight"].transpose(1, 0).numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._feedforward_output_dense.bias.assign(
+            hf_wts[f"decoder.layers.{i}.fc2.bias"].numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._feedforward_layernorm.gamma.assign(
+            hf_wts[f"decoder.layers.{i}.final_layer_norm.weight"].numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"transformer_decoder_layer_{i}"
+        )._feedforward_layernorm.beta.assign(
+            hf_wts[f"decoder.layers.{i}.final_layer_norm.bias"].numpy()
+        )
+
+    # Save the model.
+    print("\n-> Save KerasNLP model weights.")
+    keras_nlp_model.save_weights(os.path.join(FLAGS.preset, "model.h5"))
+
+    return keras_nlp_model
+
+
+def extract_vocab(hf_tokenizer):
+    vocabulary_path = os.path.join(FLAGS.preset, "vocab.json")
+    merges_path = os.path.join(FLAGS.preset, "merges.txt")
+    print(f"\n-> Save KerasNLP vocab to `{vocabulary_path}`.")
+    print(f"-> Save KerasNLP merges to `{merges_path}`.")
+
+    # Huggingface has a save_vocabulary function but it's not byte-for-byte
+    # with the source. Instead copy the original downloaded file directly.
+    shutil.copyfile(
+        transformers.utils.hub.get_file_from_repo(
+            hf_tokenizer.name_or_path, "vocab.json"
+        ),
+        vocabulary_path,
+    )
+    shutil.copyfile(
+        transformers.utils.hub.get_file_from_repo(
+            hf_tokenizer.name_or_path, "merges.txt"
+        ),
+        merges_path,
+    )
+
+    keras_nlp_tokenizer = keras_nlp.models.BartTokenizer(
+        vocabulary=vocabulary_path, merges=merges_path
+    )
+
+    print("-> Print MD5 checksum of the vocab files.")
+    print(f"`{vocabulary_path}` md5sum: ", get_md5_checksum(vocabulary_path))
+    print(f"`{merges_path}` md5sum: ", get_md5_checksum(merges_path))
+
+    return keras_nlp_tokenizer
+
+
+def check_output(
+    keras_nlp_tokenizer,
+    keras_nlp_model,
+    hf_tokenizer,
+    hf_model,
+):
+    print("\n-> Check the outputs.")
+    enc_sample_text = [
+        "cricket is awesome, easily the best sport in the world!"
+    ]
+    dec_sample_text = [
+        "football is good too, but nowhere near as good as cricket."
+    ]
+
+    # KerasNLP
+    keras_nlp_enc_token_ids = keras_nlp_tokenizer(
+        tf.constant(enc_sample_text)
+    ).to_tensor()
+    keras_nlp_enc_token_ids = tf.concat(
+        [
+            tf.constant([[keras_nlp_tokenizer.start_token_id]]),
+            keras_nlp_enc_token_ids,
+            tf.constant([[keras_nlp_tokenizer.end_token_id]]),
+        ],
+        axis=-1,
+    )
+    keras_nlp_dec_token_ids = keras_nlp_tokenizer(
+        tf.constant(dec_sample_text)
+    ).to_tensor()
+    keras_nlp_dec_token_ids = tf.concat(
+        [
+            tf.constant([[keras_nlp_tokenizer.start_token_id]]),
+            keras_nlp_dec_token_ids,
+            tf.constant([[keras_nlp_tokenizer.end_token_id]]),
+        ],
+        axis=-1,
+    )
+    keras_nlp_inputs = {
+        "encoder_token_ids": keras_nlp_enc_token_ids,
+        "encoder_padding_mask": keras_nlp_enc_token_ids
+        != keras_nlp_tokenizer.pad_token_id,
+        "decoder_token_ids": keras_nlp_dec_token_ids,
+        "decoder_padding_mask": keras_nlp_dec_token_ids
+        != keras_nlp_tokenizer.pad_token_id,
+    }
+    keras_nlp_output = keras_nlp_model.predict(keras_nlp_inputs)
+
+    # HF
+    hf_enc_inputs = hf_tokenizer(enc_sample_text, return_tensors="pt")
+    hf_dec_inputs = hf_tokenizer(dec_sample_text, return_tensors="pt")
+
+    hf_output = hf_model(
+        **hf_enc_inputs,
+        decoder_input_ids=hf_dec_inputs["input_ids"],
+        decoder_attention_mask=hf_dec_inputs["attention_mask"],
+    )
+
+    print("Encoder Outputs:")
+    print(
+        "KerasNLP output:",
+        keras_nlp_output["encoder_sequence_output"][0, 0, :10],
+    )
+    print("HF output:", hf_output.encoder_last_hidden_state[0, 0, :10])
+    print(
+        "Difference:",
+        np.mean(
+            keras_nlp_output["encoder_sequence_output"]
+            - hf_output.encoder_last_hidden_state.detach().numpy()
+        ),
+    )
+
+    print("Decoder Outputs:")
+    print(
+        "KerasNLP output:",
+        keras_nlp_output["decoder_sequence_output"][0, 0, :10],
+    )
+    print("HF output:", hf_output.last_hidden_state[0, 0, :10])
+    print(
+        "Difference:",
+        np.mean(
+            keras_nlp_output["decoder_sequence_output"]
+            - hf_output.last_hidden_state.detach().numpy()
+        ),
+    )
+
+    # Show the MD5 checksum of the model weights.
+    print(
+        "Model md5sum: ",
+        get_md5_checksum(os.path.join(FLAGS.preset, "model.h5")),
+    )
+
+
+def main(_):
+    os.makedirs(FLAGS.preset)
+
+    hf_model_name = PRESET_MAP[FLAGS.preset]
+
+    print("\n-> Load HF model and HF tokenizer.")
+    hf_model = transformers.AutoModel.from_pretrained(hf_model_name)
+    hf_model.eval()
+    hf_tokenizer = transformers.AutoTokenizer.from_pretrained(hf_model_name)
+
+    keras_nlp_model = convert_checkpoints(hf_model)
+    print("\n -> Load KerasNLP tokenizer.")
+    keras_nlp_tokenizer = extract_vocab(hf_tokenizer)
+
+    check_output(
+        keras_nlp_tokenizer,
+        keras_nlp_model,
+        hf_tokenizer,
+        hf_model,
+    )
+
+
+if __name__ == "__main__":
+    flags.mark_flag_as_required("preset")
+    app.run(main)
diff --git a/tools/checkpoint_conversion/convert_f_net_checkpoints.py b/tools/checkpoint_conversion/convert_f_net_checkpoints.py
new file mode 100644
index 00000000..68e52057
--- /dev/null
+++ b/tools/checkpoint_conversion/convert_f_net_checkpoints.py
@@ -0,0 +1,216 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import shutil
+
+import numpy as np
+import tensorflow as tf
+import transformers
+from absl import app
+from absl import flags
+
+import keras_nlp
+from tools.checkpoint_conversion.checkpoint_conversion_utils import (
+    get_md5_checksum,
+)
+
+PRESET_MAP = {
+    "f_net_base_en": "google/fnet-base",
+    "f_net_large_en": "google/fnet-large",
+}
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string(
+    "preset", None, f'Must be one of {",".join(PRESET_MAP.keys())}'
+)
+
+
+def convert_checkpoints(hf_model):
+    print("\n-> Convert original weights to KerasNLP format.")
+
+    print("\n-> Load KerasNLP model.")
+    keras_nlp_model = keras_nlp.models.FNetBackbone.from_preset(
+        FLAGS.preset, load_weights=False
+    )
+
+    hf_wts = hf_model.state_dict()
+    print("Original weights:")
+    print(list(hf_wts.keys()))
+
+    keras_nlp_model.get_layer("token_embedding").embeddings.assign(
+        hf_wts["embeddings.word_embeddings.weight"]
+    )
+    keras_nlp_model.get_layer("position_embedding").position_embeddings.assign(
+        hf_wts["embeddings.position_embeddings.weight"]
+    )
+    keras_nlp_model.get_layer("segment_embedding").embeddings.assign(
+        hf_wts["embeddings.token_type_embeddings.weight"]
+    )
+
+    keras_nlp_model.get_layer("embeddings_layer_norm").gamma.assign(
+        hf_wts["embeddings.LayerNorm.weight"]
+    )
+    keras_nlp_model.get_layer("embeddings_layer_norm").beta.assign(
+        hf_wts["embeddings.LayerNorm.bias"]
+    )
+
+    keras_nlp_model.get_layer("embedding_projection").kernel.assign(
+        hf_wts["embeddings.projection.weight"].T
+    )
+    keras_nlp_model.get_layer("embedding_projection").bias.assign(
+        hf_wts["embeddings.projection.bias"]
+    )
+
+    for i in range(keras_nlp_model.num_layers):
+        keras_nlp_model.get_layer(
+            f"f_net_layer_{i}"
+        )._mixing_layer_norm.gamma.assign(
+            hf_wts[f"encoder.layer.{i}.fourier.output.LayerNorm.weight"].numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"f_net_layer_{i}"
+        )._mixing_layer_norm.beta.assign(
+            hf_wts[f"encoder.layer.{i}.fourier.output.LayerNorm.bias"].numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"f_net_layer_{i}"
+        )._intermediate_dense.kernel.assign(
+            hf_wts[f"encoder.layer.{i}.intermediate.dense.weight"]
+            .transpose(1, 0)
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"f_net_layer_{i}"
+        )._intermediate_dense.bias.assign(
+            hf_wts[f"encoder.layer.{i}.intermediate.dense.bias"].numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"f_net_layer_{i}"
+        )._output_dense.kernel.assign(
+            hf_wts[f"encoder.layer.{i}.output.dense.weight"]
+            .transpose(1, 0)
+            .numpy()
+        )
+        keras_nlp_model.get_layer(f"f_net_layer_{i}")._output_dense.bias.assign(
+            hf_wts[f"encoder.layer.{i}.output.dense.bias"].numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"f_net_layer_{i}"
+        )._output_layer_norm.gamma.assign(
+            hf_wts[f"encoder.layer.{i}.output.LayerNorm.weight"].numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"f_net_layer_{i}"
+        )._output_layer_norm.beta.assign(
+            hf_wts[f"encoder.layer.{i}.output.LayerNorm.bias"].numpy()
+        )
+
+    keras_nlp_model.get_layer("pooled_dense").kernel.assign(
+        hf_wts["pooler.dense.weight"].transpose(1, 0).numpy()
+    )
+    keras_nlp_model.get_layer("pooled_dense").bias.assign(
+        hf_wts["pooler.dense.bias"].numpy()
+    )
+
+    # Save the model.
+    print("\n-> Save KerasNLP model weights.")
+    keras_nlp_model.save_weights(os.path.join(FLAGS.preset, "model.h5"))
+
+    return keras_nlp_model
+
+
+def extract_vocab(hf_tokenizer):
+    spm_path = os.path.join(FLAGS.preset, "spiece.model")
+    print(f"\n-> Save KerasNLP SPM vocabulary file to `{spm_path}`.")
+
+    shutil.copyfile(
+        transformers.utils.hub.get_file_from_repo(
+            hf_tokenizer.name_or_path, "spiece.model"
+        ),
+        spm_path,
+    )
+
+    keras_nlp_tokenizer = keras_nlp.models.FNetTokenizer(
+        proto=spm_path,
+    )
+    keras_nlp_preprocessor = keras_nlp.models.FNetPreprocessor(
+        keras_nlp_tokenizer
+    )
+
+    print("-> Print MD5 checksum of the vocab files.")
+    print(f"`{spm_path}` md5sum: ", get_md5_checksum(spm_path))
+
+    return keras_nlp_preprocessor
+
+
+def check_output(
+    keras_nlp_preprocessor,
+    keras_nlp_model,
+    hf_tokenizer,
+    hf_model,
+):
+    print("\n-> Check the outputs.")
+    sample_text = ["cricket is awesome, easily the best sport in the world!"]
+
+    # KerasNLP
+    keras_nlp_inputs = keras_nlp_preprocessor(tf.constant(sample_text))
+    keras_nlp_output = keras_nlp_model.predict(keras_nlp_inputs)[
+        "sequence_output"
+    ]
+
+    # HF
+    hf_inputs = hf_tokenizer(
+        sample_text, padding="max_length", return_tensors="pt"
+    )
+    hf_output = hf_model(**hf_inputs).last_hidden_state
+
+    print("KerasNLP output:", keras_nlp_output[0, 0, :10])
+    print("HF output:", hf_output[0, 0, :10])
+    print("Difference:", np.mean(keras_nlp_output - hf_output.detach().numpy()))
+
+    # Show the MD5 checksum of the model weights.
+    print(
+        "Model md5sum: ",
+        get_md5_checksum(os.path.join(FLAGS.preset, "model.h5")),
+    )
+
+
+def main(_):
+    os.makedirs(FLAGS.preset)
+
+    hf_model_name = PRESET_MAP[FLAGS.preset]
+
+    print("\n-> Load HF model and HF tokenizer.")
+    hf_model = transformers.AutoModel.from_pretrained(hf_model_name)
+    hf_model.eval()
+    hf_tokenizer = transformers.AutoTokenizer.from_pretrained(hf_model_name)
+
+    keras_nlp_model = convert_checkpoints(hf_model)
+    print("\n -> Load KerasNLP preprocessor.")
+    keras_nlp_preprocessor = extract_vocab(hf_tokenizer)
+
+    check_output(
+        keras_nlp_preprocessor,
+        keras_nlp_model,
+        hf_tokenizer,
+        hf_model,
+    )
+
+
+if __name__ == "__main__":
+    flags.mark_flag_as_required("preset")
+    app.run(main)


Test patch:
diff --git a/keras_nlp/models/albert/albert_tokenizer_test.py b/keras_nlp/models/albert/albert_tokenizer_test.py
index b7099b8d..d80657db 100644
--- a/keras_nlp/models/albert/albert_tokenizer_test.py
+++ b/keras_nlp/models/albert/albert_tokenizer_test.py
@@ -65,8 +65,20 @@ class AlbertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
         self.assertEqual(output, tf.constant(["the quick brown fox"]))
 
     def test_vocabulary_size(self):
-        tokenizer = AlbertTokenizer(proto=self.proto)
-        self.assertEqual(tokenizer.vocabulary_size(), 10)
+        self.assertEqual(self.tokenizer.vocabulary_size(), 10)
+
+    def test_errors_missing_special_tokens(self):
+        bytes_io = io.BytesIO()
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=iter(["abc"]),
+            model_writer=bytes_io,
+            vocab_size=5,
+            pad_id=-1,
+            eos_id=-1,
+            bos_id=-1,
+        )
+        with self.assertRaises(ValueError):
+            AlbertTokenizer(proto=bytes_io.getvalue())
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
diff --git a/keras_nlp/models/bart/bart_tokenizer_test.py b/keras_nlp/models/bart/bart_tokenizer_test.py
index cd8fd624..1ec896eb 100644
--- a/keras_nlp/models/bart/bart_tokenizer_test.py
+++ b/keras_nlp/models/bart/bart_tokenizer_test.py
@@ -65,6 +65,10 @@ class BartTokenizerTest(tf.test.TestCase, parameterized.TestCase):
     def test_vocabulary_size(self):
         self.assertEqual(self.tokenizer.vocabulary_size(), 12)
 
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            BartTokenizer(vocabulary=["a", "b", "c"], merges=[])
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
diff --git a/keras_nlp/models/bert/bert_tokenizer_test.py b/keras_nlp/models/bert/bert_tokenizer_test.py
index 5f72a59f..1d17d9aa 100644
--- a/keras_nlp/models/bert/bert_tokenizer_test.py
+++ b/keras_nlp/models/bert/bert_tokenizer_test.py
@@ -27,17 +27,16 @@ class BertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
         self.vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
         self.vocab += ["THE", "QUICK", "BROWN", "FOX"]
         self.vocab += ["the", "quick", "brown", "fox"]
+        self.tokenizer = BertTokenizer(vocabulary=self.vocab)
 
     def test_tokenize(self):
         input_data = "THE QUICK BROWN FOX."
-        tokenizer = BertTokenizer(vocabulary=self.vocab)
-        output = tokenizer(input_data)
+        output = self.tokenizer(input_data)
         self.assertAllEqual(output, [5, 6, 7, 8, 1])
 
     def test_tokenize_batch(self):
         input_data = tf.constant(["THE QUICK BROWN FOX.", "THE FOX."])
-        tokenizer = BertTokenizer(vocabulary=self.vocab)
-        output = tokenizer(input_data)
+        output = self.tokenizer(input_data)
         self.assertAllEqual(output, [[5, 6, 7, 8, 1], [5, 8, 1]])
 
     def test_lowercase(self):
@@ -48,13 +47,15 @@ class BertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
 
     def test_detokenize(self):
         input_tokens = [[5, 6, 7, 8]]
-        tokenizer = BertTokenizer(vocabulary=self.vocab)
-        output = tokenizer.detokenize(input_tokens)
+        output = self.tokenizer.detokenize(input_tokens)
         self.assertAllEqual(output, ["THE QUICK BROWN FOX"])
 
     def test_vocabulary_size(self):
-        tokenizer = BertTokenizer(vocabulary=self.vocab)
-        self.assertEqual(tokenizer.vocabulary_size(), 13)
+        self.assertEqual(self.tokenizer.vocabulary_size(), 13)
+
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            BertTokenizer(vocabulary=["a", "b", "c"])
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py b/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py
index 1be123c9..b85f3491 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py
@@ -65,8 +65,20 @@ class DebertaV3TokenizerTest(tf.test.TestCase, parameterized.TestCase):
         self.assertEqual(output, tf.constant(["the quick brown fox"]))
 
     def test_vocabulary_size(self):
-        tokenizer = DebertaV3Tokenizer(proto=self.proto)
-        self.assertEqual(tokenizer.vocabulary_size(), 10)
+        self.assertEqual(self.tokenizer.vocabulary_size(), 10)
+
+    def test_errors_missing_special_tokens(self):
+        bytes_io = io.BytesIO()
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=iter(["abc"]),
+            model_writer=bytes_io,
+            vocab_size=5,
+            pad_id=-1,
+            eos_id=-1,
+            bos_id=-1,
+        )
+        with self.assertRaises(ValueError):
+            DebertaV3Tokenizer(proto=bytes_io.getvalue())
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
diff --git a/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py b/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py
index c3fb810a..b51c2247 100644
--- a/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py
@@ -29,17 +29,16 @@ class DistilBertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
         self.vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
         self.vocab += ["THE", "QUICK", "BROWN", "FOX"]
         self.vocab += ["the", "quick", "brown", "fox"]
+        self.tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
 
     def test_tokenize(self):
         input_data = "THE QUICK BROWN FOX."
-        tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
-        output = tokenizer(input_data)
+        output = self.tokenizer(input_data)
         self.assertAllEqual(output, [5, 6, 7, 8, 1])
 
     def test_tokenize_batch(self):
         input_data = tf.constant(["THE QUICK BROWN FOX.", "THE FOX."])
-        tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
-        output = tokenizer(input_data)
+        output = self.tokenizer(input_data)
         self.assertAllEqual(output, [[5, 6, 7, 8, 1], [5, 8, 1]])
 
     def test_lowercase(self):
@@ -50,13 +49,15 @@ class DistilBertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
 
     def test_detokenize(self):
         input_tokens = [[5, 6, 7, 8]]
-        tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
-        output = tokenizer.detokenize(input_tokens)
+        output = self.tokenizer.detokenize(input_tokens)
         self.assertAllEqual(output, ["THE QUICK BROWN FOX"])
 
     def test_vocabulary_size(self):
-        tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
-        self.assertEqual(tokenizer.vocabulary_size(), 13)
+        self.assertEqual(self.tokenizer.vocabulary_size(), 13)
+
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            DistilBertTokenizer(vocabulary=["a", "b", "c"])
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
diff --git a/keras_nlp/models/f_net/f_net_tokenizer_test.py b/keras_nlp/models/f_net/f_net_tokenizer_test.py
index a4257ffc..b5c186a2 100644
--- a/keras_nlp/models/f_net/f_net_tokenizer_test.py
+++ b/keras_nlp/models/f_net/f_net_tokenizer_test.py
@@ -68,6 +68,19 @@ class FNetTokenizerTest(tf.test.TestCase, parameterized.TestCase):
         tokenizer = FNetTokenizer(proto=self.proto)
         self.assertEqual(tokenizer.vocabulary_size(), 10)
 
+    def test_errors_missing_special_tokens(self):
+        bytes_io = io.BytesIO()
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=iter(["abc"]),
+            model_writer=bytes_io,
+            vocab_size=5,
+            pad_id=-1,
+            eos_id=-1,
+            bos_id=-1,
+        )
+        with self.assertRaises(ValueError):
+            FNetTokenizer(proto=bytes_io.getvalue())
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
diff --git a/keras_nlp/models/gpt2/gpt2_tokenizer_test.py b/keras_nlp/models/gpt2/gpt2_tokenizer_test.py
index 7476bbc2..bbc11cfc 100644
--- a/keras_nlp/models/gpt2/gpt2_tokenizer_test.py
+++ b/keras_nlp/models/gpt2/gpt2_tokenizer_test.py
@@ -66,6 +66,10 @@ class GPT2TokenizerTest(tf.test.TestCase, parameterized.TestCase):
     def test_vocabulary_size(self):
         self.assertEqual(self.tokenizer.vocabulary_size(), 10)
 
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            GPT2Tokenizer(vocabulary=["a", "b", "c"], merges=[])
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
diff --git a/keras_nlp/models/opt/opt_tokenizer_test.py b/keras_nlp/models/opt/opt_tokenizer_test.py
index 2e0eb461..b63493a3 100644
--- a/keras_nlp/models/opt/opt_tokenizer_test.py
+++ b/keras_nlp/models/opt/opt_tokenizer_test.py
@@ -66,6 +66,10 @@ class OPTTokenizerTest(tf.test.TestCase, parameterized.TestCase):
     def test_vocabulary_size(self):
         self.assertEqual(self.tokenizer.vocabulary_size(), 12)
 
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            OPTTokenizer(vocabulary=["a", "b", "c"], merges=[])
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
diff --git a/keras_nlp/models/roberta/roberta_tokenizer_test.py b/keras_nlp/models/roberta/roberta_tokenizer_test.py
index 4f96d455..5da219b9 100644
--- a/keras_nlp/models/roberta/roberta_tokenizer_test.py
+++ b/keras_nlp/models/roberta/roberta_tokenizer_test.py
@@ -66,6 +66,10 @@ class RobertaTokenizerTest(tf.test.TestCase, parameterized.TestCase):
     def test_vocabulary_size(self):
         self.assertEqual(self.tokenizer.vocabulary_size(), 13)
 
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            RobertaTokenizer(vocabulary=["a", "b", "c"], merges=[])
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
diff --git a/keras_nlp/samplers/beam_sampler_test.py b/keras_nlp/samplers/beam_sampler_test.py
index 82abc889..b92a460b 100644
--- a/keras_nlp/samplers/beam_sampler_test.py
+++ b/keras_nlp/samplers/beam_sampler_test.py
@@ -125,7 +125,7 @@ class BeamSamplerTest(tf.test.TestCase, parameterized.TestCase):
             prob = tf.constant([[[0.0, 0.0, 0.0, 1.0]]])
             return tf.tile(prob, [batch_size, seq_length, 1])
 
-        max_length = 5
+        max_length = 4
         inputs = tf.constant([[0, 1], [1, 2]])
         outputs = self.sampler(
             inputs,
@@ -133,5 +133,17 @@ class BeamSamplerTest(tf.test.TestCase, parameterized.TestCase):
             max_length=max_length,
             end_token_id=2,
         )
-        expected_outputs = tf.ragged.constant([[0, 1, 3, 3, 3], [1]])
+        # end_token in prompt does not trigger truncation.
+        expected_outputs = tf.ragged.constant([[0, 1, 3, 3], [1, 2, 3, 3]])
+        self.assertAllEqual(outputs, expected_outputs)
+
+        max_length = 4
+        inputs = tf.constant([[0, 1], [1, 3]])
+        outputs = self.sampler(
+            inputs,
+            token_probability_fn,
+            max_length=max_length,
+            end_token_id=3,
+        )
+        expected_outputs = tf.ragged.constant([[0, 1], [1, 3]])
         self.assertAllEqual(outputs, expected_outputs)
diff --git a/keras_nlp/samplers/greedy_sampler_test.py b/keras_nlp/samplers/greedy_sampler_test.py
index 63c0cb39..1d1eb6ae 100644
--- a/keras_nlp/samplers/greedy_sampler_test.py
+++ b/keras_nlp/samplers/greedy_sampler_test.py
@@ -93,8 +93,6 @@ class GreedySamplerTest(tf.test.TestCase, parameterized.TestCase):
         )
 
     def test_end_token_id(self):
-        max_length = 5
-
         def token_probability_fn(inputs, mask):
             batch_size = inputs.shape[0]
             prob = tf.constant([[[0.0, 0.0, 0.0, 1.0]]])
@@ -102,6 +100,7 @@ class GreedySamplerTest(tf.test.TestCase, parameterized.TestCase):
                 tf.repeat(prob, batch_size, axis=0), max_length, axis=1
             )
 
+        max_length = 4
         sampler = GreedySampler()
         inputs = tf.constant([[0, 1], [1, 2]])
         outputs = sampler(
@@ -110,7 +109,18 @@ class GreedySamplerTest(tf.test.TestCase, parameterized.TestCase):
             max_length=max_length,
             end_token_id=2,
         )
-        expected_outputs = tf.ragged.constant([[0, 1, 3, 3, 3], [1]])
+        # end_token in prompt does not trigger truncation.
+        expected_outputs = tf.ragged.constant([[0, 1, 3, 3], [1, 2, 3, 3]])
+        self.assertAllEqual(outputs, expected_outputs)
+
+        outputs = sampler(
+            inputs,
+            token_probability_fn,
+            max_length=max_length,
+            end_token_id=3,
+        )
+        # Generated end_token will be truncated.
+        expected_outputs = tf.ragged.constant([[0, 1], [1, 2]])
         self.assertAllEqual(outputs, expected_outputs)
 
     def test_compare_xla_noxla_results(self):
diff --git a/keras_nlp/samplers/top_k_sampler_test.py b/keras_nlp/samplers/top_k_sampler_test.py
index 4cca77f6..062f7af7 100644
--- a/keras_nlp/samplers/top_k_sampler_test.py
+++ b/keras_nlp/samplers/top_k_sampler_test.py
@@ -144,10 +144,10 @@ class TopKSamplerTest(tf.test.TestCase, parameterized.TestCase):
             prob = tf.constant([[[0.0, 0.0, 0.0, 1.0]]])
             return tf.tile(prob, [batch_size, seq_length, 1])
 
-        max_length = 5
-        inputs = tf.constant([[0, 1], [1, 2]])
         tf.random.set_seed(42)
         sampler = TopKSampler(k=4, seed=42)
+        max_length = 4
+        inputs = tf.constant([[0, 1], [1, 2]])
         outputs = sampler(
             inputs,
             token_probability_fn,
@@ -155,6 +155,17 @@ class TopKSamplerTest(tf.test.TestCase, parameterized.TestCase):
             end_token_id=2,
             from_logits=False,
         )
-        # Top-k sampling result with seed 42.
-        expected_outputs = tf.ragged.constant([[0, 1, 3, 3, 3], [1]])
+        # end_token in prompt does not trigger truncation.
+        expected_outputs = tf.ragged.constant([[0, 1, 3, 3], [1, 2, 3, 3]])
+        self.assertAllEqual(outputs, expected_outputs)
+
+        outputs = sampler(
+            inputs,
+            token_probability_fn,
+            max_length=max_length,
+            end_token_id=3,
+            from_logits=False,
+        )
+        # Generated end_token will be truncated.
+        expected_outputs = tf.ragged.constant([[0, 1], [1, 2]])
         self.assertAllEqual(outputs, expected_outputs)
diff --git a/keras_nlp/samplers/top_p_sampler_test.py b/keras_nlp/samplers/top_p_sampler_test.py
index 1291de13..6cb56cf4 100644
--- a/keras_nlp/samplers/top_p_sampler_test.py
+++ b/keras_nlp/samplers/top_p_sampler_test.py
@@ -144,10 +144,10 @@ class TopPSamplerTest(tf.test.TestCase, parameterized.TestCase):
             prob = tf.constant([[[0.0, 0.0, 0.0, 1.0]]])
             return tf.tile(prob, [batch_size, seq_length, 1])
 
-        max_length = 5
-        inputs = tf.constant([[0, 1], [1, 2]])
         tf.random.set_seed(42)
         sampler = TopPSampler(p=0.1, seed=42)
+        max_length = 4
+        inputs = tf.constant([[0, 1], [1, 2]])
         outputs = sampler(
             inputs,
             token_probability_fn,
@@ -155,6 +155,17 @@ class TopPSamplerTest(tf.test.TestCase, parameterized.TestCase):
             end_token_id=2,
             from_logits=False,
         )
-        # Top-p sampling result with seed 42.
-        expected_outputs = tf.ragged.constant([[0, 1, 3, 3, 3], [1]])
+        # end_token in prompt does not trigger truncation.
+        expected_outputs = tf.ragged.constant([[0, 1, 3, 3], [1, 2, 3, 3]])
+        self.assertAllEqual(outputs, expected_outputs)
+
+        outputs = sampler(
+            inputs,
+            token_probability_fn,
+            max_length=max_length,
+            end_token_id=3,
+            from_logits=False,
+        )
+        # Generated end_token will be truncated.
+        expected_outputs = tf.ragged.constant([[0, 1], [1, 2]])
         self.assertAllEqual(outputs, expected_outputs)
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
index 86cb3817..fa1b8bd9 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer_test.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
@@ -226,19 +226,19 @@ class WordPieceTokenizerTest(tf.test.TestCase, parameterized.TestCase):
 
     def test_no_oov_token_in_vocabulary(self):
         vocab_data = ["qu", "@@ick", "br", "@@OWN", "fox"]
-        with self.assertRaises(RuntimeError):
+        with self.assertRaises(ValueError):
             WordPieceTokenizer(
                 vocabulary=vocab_data,
             )
 
         vocab_data = ["@UNK@", "qu", "@@ick", "br", "@@OWN", "fox"]
-        with self.assertRaises(RuntimeError):
+        with self.assertRaises(ValueError):
             WordPieceTokenizer(
                 vocabulary=vocab_data,
             )
 
         vocab_data = ["UNK", "qu", "@@ick", "br", "@@OWN", "fox"]
-        with self.assertRaises(RuntimeError):
+        with self.assertRaises(ValueError):
             WordPieceTokenizer(
                 vocabulary=vocab_data,
             )