Repository: keras-team/keras-nlp
Pull request #751 (2023-02-14 18:35:48)

Adding ragged support to SinePositionEncoding
Solves #109 


Patch:
diff --git a/keras_nlp/layers/sine_position_encoding.py b/keras_nlp/layers/sine_position_encoding.py
index 7ae2ca11..3fabb542 100644
--- a/keras_nlp/layers/sine_position_encoding.py
+++ b/keras_nlp/layers/sine_position_encoding.py
@@ -31,6 +31,11 @@ class SinePositionEncoding(keras.layers.Layer):
     positional encoding the same size as the embedded token tensor, which
     can be added directly to the embedded token tensor.
 
+    This layer optionally accepts `tf.RaggedTensor`s as inputs to process
+    batches of sequences of different lengths. The one ragged dimension must be
+    the dimension that corresponds to the sequence, that is, the penultimate
+    dimension.
+
     Args:
         max_wavelength: The maximum angular wavelength of the sine/cosine
             curves, as described in Attention is All You Need. Defaults to
@@ -65,10 +70,26 @@ class SinePositionEncoding(keras.layers.Layer):
     def call(self, inputs):
         # TODO(jbischof): replace `hidden_size` with`hidden_dim` for consistency
         # with other layers.
-        input_shape = tf.shape(inputs)
-        # length of sequence is the second last dimension of the inputs
-        seq_length = input_shape[-2]
-        hidden_size = input_shape[-1]
+        if isinstance(inputs, tf.RaggedTensor):
+            bounding_shape = inputs.bounding_shape()
+            position_embeddings = (
+                self._compute_trim_and_broadcast_position_embeddings(
+                    bounding_shape,
+                )
+            )
+            # then apply row lengths to recreate the same ragged shape as inputs
+            return tf.RaggedTensor.from_tensor(
+                position_embeddings,
+                inputs.nested_row_lengths(),
+            )
+        else:
+            return self._compute_trim_and_broadcast_position_embeddings(
+                tf.shape(inputs),
+            )
+
+    def _compute_trim_and_broadcast_position_embeddings(self, shape):
+        seq_length = shape[-2]
+        hidden_size = shape[-1]
         position = tf.cast(tf.range(seq_length), self.compute_dtype)
         min_freq = tf.cast(1 / self.max_wavelength, dtype=self.compute_dtype)
         timescales = tf.pow(
@@ -85,7 +106,7 @@ class SinePositionEncoding(keras.layers.Layer):
             tf.sin(angles) * sin_mask + tf.cos(angles) * cos_mask
         )
 
-        return tf.broadcast_to(positional_encodings, input_shape)
+        return tf.broadcast_to(positional_encodings, shape)
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/layers/token_and_position_embedding.py b/keras_nlp/layers/token_and_position_embedding.py
index df46b362..d002d874 100644
--- a/keras_nlp/layers/token_and_position_embedding.py
+++ b/keras_nlp/layers/token_and_position_embedding.py
@@ -23,9 +23,12 @@ from keras_nlp.utils.keras_utils import clone_initializer
 @keras.utils.register_keras_serializable(package="keras_nlp")
 class TokenAndPositionEmbedding(keras.layers.Layer):
     """A layer which sums a token and position embedding.
-
-    This layer assumes that the last dimension in the input corresponds
-    to the sequence dimension.
+    
+    Token and position embeddings are ways of representing words and their order
+    in a sentence. This layer creates a `keras.layers.Embedding` token embedding
+    and a `keras_nlp.layers.PositionEmbedding` position embedding and sums their
+    output when called. This layer assumes that the last dimension in the input
+    corresponds to the sequence dimension.
 
     Args:
         vocabulary_size: The size of the vocabulary.
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/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/layers/sine_position_encoding_test.py b/keras_nlp/layers/sine_position_encoding_test.py
index c61f1c99..05aff73e 100644
--- a/keras_nlp/layers/sine_position_encoding_test.py
+++ b/keras_nlp/layers/sine_position_encoding_test.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 """Tests for Sinusoidal Positional encoding."""
 
-
 import tensorflow as tf
 from tensorflow import keras
 
@@ -92,6 +91,84 @@ class SinePositionEncodingTest(tf.test.TestCase):
         self.assertAllClose(output[0, 0, :], expected_encoding_position_0)
         self.assertAllClose(output[0, 3, :], expected_encoding_position_3)
 
+    def test_ragged_tensor_with_3_dimensions(self):
+        feature_size = 2
+        test_layer = sine_position_encoding.SinePositionEncoding()
+        # Create a 3-dimensional ragged input (the first dimension is implicit).
+        input_tensor = keras.Input(
+            shape=(None, feature_size), dtype=tf.float32, ragged=True
+        )
+        output_tensor = test_layer(input_tensor)
+        model = keras.Model(input_tensor, output_tensor)
+
+        input_data = tf.ragged.constant(
+            [
+                [[1.0, 1.0], [1.0, 1.0]],
+                [],
+                [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]],
+                [[1.0, 1.0]],
+            ],
+            ragged_rank=1,
+            inner_shape=(2,),
+        )
+        expected_output_data = tf.ragged.constant(
+            [
+                [[0.0, 1.0], [0.84147096, 0.5403023]],
+                [],
+                [[0.0, 1.0], [0.84147096, 0.5403023], [0.9092974, -0.41614684]],
+                [[0.0, 1.0]],
+            ],
+            ragged_rank=1,
+            inner_shape=(2,),
+        )
+        output_data = model.predict(input_data)
+        self.assertAllClose(output_data, expected_output_data)
+
+    def test_ragged_tensor_with_4_dimensions(self):
+        feature_size = 2
+        test_layer = sine_position_encoding.SinePositionEncoding()
+        # Create a 4-dimensional ragged input (the first dimension is implicit).
+        input_tensor = keras.Input(
+            shape=(None, None, feature_size), dtype=tf.float32, ragged=True
+        )
+        output_tensor = test_layer(input_tensor)
+        model = keras.Model(input_tensor, output_tensor)
+
+        input_data = tf.ragged.constant(
+            [
+                [
+                    [[1.0, 1.0], [1.0, 1.0]],
+                    [],
+                ],
+                [
+                    [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]],
+                    [[1.0, 1.0]],
+                ],
+            ],
+            ragged_rank=2,
+            inner_shape=(2,),
+        )
+        expected_output_data = tf.ragged.constant(
+            [
+                [
+                    [[0.0, 1.0], [0.84147096, 0.5403023]],
+                    [],
+                ],
+                [
+                    [
+                        [0.0, 1.0],
+                        [0.84147096, 0.5403023],
+                        [0.9092974, -0.41614684],
+                    ],
+                    [[0.0, 1.0]],
+                ],
+            ],
+            ragged_rank=2,
+            inner_shape=(2,),
+        )
+        output_data = model.predict(input_data)
+        self.assertAllClose(output_data, expected_output_data)
+
     def test_get_config_and_from_config(self):
         pos_encoding = sine_position_encoding.SinePositionEncoding(
             max_wavelength=1000,