Repository: keras-team/keras-nlp

Issue:
Add an AlbertMaskedLM task model
In https://github.com/keras-team/keras-nlp/pull/653 we added a masked language modeling task for RoBERTa. We can make a similar change for the `Albert` model.

* [ ] Update `AlbertTokenizer` to expect a mask token.
* [ ] Add a `AlbertMaskedLMPreprocessor` preprocessor layer and tests.
* [ ] Add a `AlbertMaskedLM` task model and tests.
* [ ] Update `keras_nlp/models/__init__.py` to export `AlbertMaskedLM` and `AlbertMaskedLMPreprocessor`.

Discussion:
shivance: cc: @jbischof @mattdangerw @chenmoneygithub 
shivance: Ready for review
shivance: CI is green, ready for review @mattdangerw 
shivance: @mattdangerw Please check if the updated test is fine, (passes on local)
shivance: @mattdangerw CI is green in terms of tests.
Not sure why cloud build fails in accelerator test
mattdangerw: @shivance thanks! will take another pass soon. looks like the accelerator test is just a timeout, so probably unrelated to the code changes here.
shivance: GTG @mattdangerw 
shivance: Still WIP
mattdangerw: > Still WIP

No rush! Ping when this is ready for 👀 again!
shivance: The PR is ready for review, waiting for the CI, although tests pass in local.
shivance: CI is green @mattdangerw . Ready for review !

Patch:
diff --git a/keras_nlp/layers/masked_lm_head.py b/keras_nlp/layers/masked_lm_head.py
index 2a31aa7c..2f749eb0 100644
--- a/keras_nlp/layers/masked_lm_head.py
+++ b/keras_nlp/layers/masked_lm_head.py
@@ -140,7 +140,10 @@ class MaskedLMHead(keras.layers.Layer):
             self.vocabulary_size = shape[0]
 
     def build(self, input_shapes):
-        feature_size = input_shapes[-1]
+        if self.embedding_weights is not None:
+            feature_size = self.embedding_weights.shape[-1]
+        else:
+            feature_size = input_shapes[-1]
 
         self._dense = keras.layers.Dense(
             feature_size,
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..16cd757b 100644
--- a/keras_nlp/layers/token_and_position_embedding.py
+++ b/keras_nlp/layers/token_and_position_embedding.py
@@ -24,8 +24,11 @@ from keras_nlp.utils.keras_utils import clone_initializer
 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/__init__.py b/keras_nlp/models/__init__.py
index 893d9300..a490a113 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -13,6 +13,10 @@
 # limitations under the License.
 
 from keras_nlp.models.albert.albert_backbone import AlbertBackbone
+from keras_nlp.models.albert.albert_masked_lm import AlbertMaskedLM
+from keras_nlp.models.albert.albert_masked_lm_preprocessor import (
+    AlbertMaskedLMPreprocessor,
+)
 from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
 from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
 from keras_nlp.models.bart.bart_backbone import BartBackbone
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/albert/albert_masked_lm.py b/keras_nlp/models/albert/albert_masked_lm.py
new file mode 100644
index 00000000..2fe24885
--- /dev/null
+++ b/keras_nlp/models/albert/albert_masked_lm.py
@@ -0,0 +1,154 @@
+# 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.
+
+"""ALBERT masked LM model."""
+
+import copy
+
+from tensorflow import keras
+
+from keras_nlp.layers.masked_lm_head import MaskedLMHead
+from keras_nlp.models.albert.albert_backbone import AlbertBackbone
+from keras_nlp.models.albert.albert_backbone import albert_kernel_initializer
+from keras_nlp.models.albert.albert_masked_lm_preprocessor import (
+    AlbertMaskedLMPreprocessor,
+)
+from keras_nlp.models.albert.albert_presets import backbone_presets
+from keras_nlp.models.task import Task
+from keras_nlp.utils.python_utils import classproperty
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class AlbertMaskedLM(Task):
+    """An end-to-end ALBERT model for the masked language modeling task.
+
+    This model will train ALBERT on a masked language modeling task.
+    The model will predict labels for a number of masked tokens in the
+    input data. For usage of this model with pre-trained weights, see the
+    `from_preset()` method.
+
+    This model can optionally be configured with a `preprocessor` layer, in
+    which case inputs can be raw string features during `fit()`, `predict()`,
+    and `evaluate()`. Inputs will be tokenized and dynamically masked during
+    training and evaluation. This is done by default when creating the model
+    with `from_preset()`.
+
+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind.
+
+    Args:
+        backbone: A `keras_nlp.models.AlbertBackbone` instance.
+        preprocessor: A `keras_nlp.models.AlbertMaskedLMPreprocessor` or
+            `None`. If `None`, this model will not apply preprocessing, and
+            inputs should be preprocessed before calling the model.
+
+    Example usage:
+
+    Raw string inputs and pretrained backbone.
+    ```python
+    # Create a dataset with raw string features. Labels are inferred.
+    features = ["The quick brown fox jumped.", "I forgot my homework."]
+
+    # Create a AlbertMaskedLM with a pretrained backbone and further train
+    # on an MLM task.
+    masked_lm = keras_nlp.models.AlbertMaskedLM.from_preset(
+        "albert_base_en_uncased",
+    )
+    masked_lm.compile(
+        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+    )
+    masked_lm.fit(x=features, batch_size=2)
+    ```
+
+    Preprocessed inputs and custom backbone.
+    ```python
+    # Create a preprocessed dataset where 0 is the mask token.
+    preprocessed_features = {
+        "segment_ids": tf.constant(
+            [[1, 0, 0, 4, 0, 6, 7, 8]] * 2, shape=(2, 8)
+        ),
+        "token_ids": tf.constant(
+            [[1, 2, 0, 4, 0, 6, 7, 8]] * 2, shape=(2, 8)
+        ),
+        "padding_mask": tf.constant(
+            [[1, 1, 1, 1, 1, 1, 1, 1]] * 2, shape=(2, 8)
+        ),
+        "mask_positions": tf.constant([[2, 4]] * 2, shape=(2, 2))
+    }
+    # Labels are the original masked values.
+    labels = [[3, 5]] * 2
+
+    # Randomly initialize a ALBERT encoder
+    backbone = keras_nlp.models.AlbertBackbone(
+        vocabulary_size=1000,
+        num_layers=2,
+        num_heads=2,
+        embedding_dim=64,
+        hidden_dim=64,
+        intermediate_dim=128,
+        max_sequence_length=128)
+
+    # Create a ALBERT masked LM and fit the data.
+    masked_lm = keras_nlp.models.AlbertMaskedLM(
+        backbone,
+        preprocessor=None,
+    )
+    masked_lm.compile(
+        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        jit_compile=True
+    )
+    masked_lm.fit(x=preprocessed_features, y=labels, batch_size=2)
+    ```
+    """
+
+    def __init__(self, backbone, preprocessor=None, **kwargs):
+        inputs = {
+            **backbone.input,
+            "mask_positions": keras.Input(
+                shape=(None,), dtype="int32", name="mask_positions"
+            ),
+        }
+
+        backbone_outputs = backbone(backbone.input)
+        outputs = MaskedLMHead(
+            vocabulary_size=backbone.vocabulary_size,
+            embedding_weights=backbone.token_embedding.embeddings,
+            intermediate_activation=lambda x: keras.activations.gelu(
+                x, approximate=True
+            ),
+            kernel_initializer=albert_kernel_initializer(),
+            name="mlm_head",
+        )(backbone_outputs["sequence_output"], inputs["mask_positions"])
+
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            include_preprocessing=preprocessor is not None,
+            **kwargs
+        )
+
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+
+    @classproperty
+    def backbone_cls(cls):
+        return AlbertBackbone
+
+    @classproperty
+    def preprocessor_cls(cls):
+        return AlbertMaskedLMPreprocessor
+
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy(backbone_presets)
diff --git a/keras_nlp/models/albert/albert_masked_lm_preprocessor.py b/keras_nlp/models/albert/albert_masked_lm_preprocessor.py
new file mode 100644
index 00000000..1c874a0b
--- /dev/null
+++ b/keras_nlp/models/albert/albert_masked_lm_preprocessor.py
@@ -0,0 +1,198 @@
+# 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.
+
+"""ALBERT masked language model preprocessor layer."""
+
+from absl import logging
+from tensorflow import keras
+
+from keras_nlp.layers.masked_lm_mask_generator import MaskedLMMaskGenerator
+from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
+from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class AlbertMaskedLMPreprocessor(AlbertPreprocessor):
+    """ALBERT preprocessing for the masked language modeling task.
+
+    This preprocessing layer will prepare inputs for a masked language modeling
+    task. It is primarily intended for use with the
+    `keras_nlp.models.AlbertMaskedLM` task model. Preprocessing will occur in
+    multiple steps.
+
+    - Tokenize any number of input segments using the `tokenizer`.
+    - Pack the inputs together with the appropriate `"
"`, `"
"` and
+      `"<pad>"` tokens, i.e., adding a single `"
"` at the start of the
+      entire sequence, `"

"` between each segment,
+      and a `"
"` at the end of the entire sequence.
+    - Randomly select non-special tokens to mask, controlled by
+      `mask_selection_rate`.
+    - Construct a `(x, y, sample_weight)` tuple suitable for training with a
+      `keras_nlp.models.AlbertMaskedLM` task model.
+
+    Args:
+        tokenizer: A `keras_nlp.models.AlbertTokenizer` instance.
+        sequence_length: The length of the packed inputs.
+        mask_selection_rate: The probability an input token will be dynamically
+            masked.
+        mask_selection_length: The maximum number of masked tokens supported
+            by the layer.
+        mask_token_rate: float, defaults to 0.8. `mask_token_rate` must be
+            between 0 and 1 which indicates how often the mask_token is
+            substituted for tokens selected for masking.
+        random_token_rate: float, defaults to 0.1. `random_token_rate` must be
+            between 0 and 1 which indicates how often a random token is
+            substituted for tokens selected for masking. Default is 0.1.
+            Note: mask_token_rate + random_token_rate <= 1,  and for
+            (1 - mask_token_rate - random_token_rate), the token will not be
+            changed.
+        truncate: string. The algorithm to truncate a list of batched segments
+            to fit within `sequence_length`. The value can be either
+            `round_robin` or `waterfall`:
+                - `"round_robin"`: Available space is assigned one token at a
+                    time in a round-robin fashion to the inputs that still need
+                    some, until the limit is reached.
+                - `"waterfall"`: The allocation of the budget is done using a
+                    "waterfall" algorithm that allocates quota in a
+                    left-to-right manner and fills up the buckets until we run
+                    out of budget. It supports an arbitrary number of segments.
+
+    Examples:
+    ```python
+    # Load the preprocessor from a preset.
+    preprocessor = keras_nlp.models.AlbertMaskedLMPreprocessor.from_preset(
+        "albert_base_en_uncased"
+    )
+
+    # Tokenize and mask a single sentence.
+    sentence = tf.constant("The quick brown fox jumped.")
+    preprocessor(sentence)
+
+    # Tokenize and mask a batch of sentences.
+    sentences = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    preprocessor(sentences)
+
+    # Tokenize and mask a dataset of sentences.
+    features = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    ds = tf.data.Dataset.from_tensor_slices((features))
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Alternatively, you can create a preprocessor from your own vocabulary.
+    vocab_data = tf.data.Dataset.from_tensor_slices(
+        ["the quick brown fox", "the earth is round"]
+    )
+
+    # Creating sentencepiece tokenizer for ALBERT LM preprocessor
+    bytes_io = io.BytesIO()
+
+    sentencepiece.SentencePieceTrainer.train(
+        sentence_iterator=vocab_data.as_numpy_iterator(),
+        model_writer=bytes_io,
+        vocab_size=12,
+        model_type="WORD",
+        pad_id=0,
+        unk_id=1,
+        bos_id=2,
+        eos_id=3,
+        pad_piece="<pad>",
+        unk_piece="<unk>",
+        bos_piece="[CLS]",
+        eos_piece="[SEP]",
+        user_defined_symbols="[MASK]"
+    )
+
+    proto = bytes_io.getvalue()
+
+    tokenizer = AlbertTokenizer(proto=proto)
+
+    preprocessor = AlbertMaskedLMPreprocessor(
+        tokenizer=tokenizer
+    )
+
+    ```
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        sequence_length=512,
+        truncate="round_robin",
+        mask_selection_rate=0.15,
+        mask_selection_length=96,
+        mask_token_rate=0.8,
+        random_token_rate=0.1,
+        **kwargs,
+    ):
+        super().__init__(
+            tokenizer,
+            sequence_length=sequence_length,
+            truncate=truncate,
+            **kwargs,
+        )
+
+        self.masker = MaskedLMMaskGenerator(
+            mask_selection_rate=mask_selection_rate,
+            mask_selection_length=mask_selection_length,
+            mask_token_rate=mask_token_rate,
+            random_token_rate=random_token_rate,
+            vocabulary_size=tokenizer.vocabulary_size(),
+            mask_token_id=tokenizer.mask_token_id,
+            unselectable_token_ids=[
+                tokenizer.cls_token_id,
+                tokenizer.sep_token_id,
+                tokenizer.pad_token_id,
+            ],
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "mask_selection_rate": self.masker.mask_selection_rate,
+                "mask_selection_length": self.masker.mask_selection_length,
+                "mask_token_rate": self.masker.mask_token_rate,
+                "random_token_rate": self.masker.random_token_rate,
+            }
+        )
+        return config
+
+    def call(self, x, y=None, sample_weight=None):
+        if y is not None or sample_weight is not None:
+            logging.warning(
+                f"{self.__class__.__name__} generates `y` and `sample_weight` "
+                "based on your input data, but your data already contains `y` "
+                "or `sample_weight`. Your `y` and `sample_weight` will be "
+                "ignored."
+            )
+
+        x = super().call(x)
+        token_ids, segment_ids, padding_mask = (
+            x["token_ids"],
+            x["segment_ids"],
+            x["padding_mask"],
+        )
+        masker_outputs = self.masker(token_ids)
+        x = {
+            "token_ids": masker_outputs["token_ids"],
+            "segment_ids": segment_ids,
+            "padding_mask": padding_mask,
+            "mask_positions": masker_outputs["mask_positions"],
+        }
+        y = masker_outputs["mask_ids"]
+        sample_weight = masker_outputs["mask_weights"]
+        return pack_x_y_sample_weight(x, y, sample_weight)
diff --git a/keras_nlp/models/albert/albert_tokenizer.py b/keras_nlp/models/albert/albert_tokenizer.py
index 35dc7ec8..5ba47af5 100644
--- a/keras_nlp/models/albert/albert_tokenizer.py
+++ b/keras_nlp/models/albert/albert_tokenizer.py
@@ -72,7 +72,8 @@ class AlbertTokenizer(SentencePieceTokenizer):
         cls_token = "[CLS]"
         sep_token = "[SEP]"
         pad_token = "<pad>"
-        for token in [cls_token, sep_token, pad_token]:
+        mask_token = "[MASK]"
+        for token in [cls_token, sep_token, pad_token, mask_token]:
             if token not in self.get_vocabulary():
                 raise ValueError(
                     f"Cannot find token `'{token}'` in the provided "
@@ -83,6 +84,7 @@ class AlbertTokenizer(SentencePieceTokenizer):
         self.cls_token_id = self.token_to_id(cls_token)
         self.sep_token_id = self.token_to_id(sep_token)
         self.pad_token_id = self.token_to_id(pad_token)
+        self.mask_token_id = self.token_to_id(mask_token)
 
     @classproperty
     def presets(cls):
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/masked_lm_head_test.py b/keras_nlp/layers/masked_lm_head_test.py
index 15734c88..f66f0b19 100644
--- a/keras_nlp/layers/masked_lm_head_test.py
+++ b/keras_nlp/layers/masked_lm_head_test.py
@@ -46,15 +46,17 @@ class MaskedLMHeadTest(tf.test.TestCase):
             embedding_weights=embedding.embeddings,
             activation="softmax",
         )
-        encoded_tokens = keras.Input(shape=(10, 16))
+        # Use a difference "hidden dim" for the model than "embedding dim", we
+        # need to support this in the layer.
+        sequence = keras.Input(shape=(10, 32))
         positions = keras.Input(shape=(5,), dtype="int32")
-        outputs = head(encoded_tokens, mask_positions=positions)
-        model = keras.Model((encoded_tokens, positions), outputs)
-        token_data = tf.random.uniform(shape=(4, 10, 16))
+        outputs = head(sequence, mask_positions=positions)
+        model = keras.Model((sequence, positions), outputs)
+        sequence_data = tf.random.uniform(shape=(4, 10, 32))
         position_data = tf.random.uniform(
             shape=(4, 5), maxval=10, dtype="int32"
         )
-        model((token_data, position_data))
+        model((sequence_data, position_data))
 
     def test_get_config_and_from_config(self):
         head = masked_lm_head.MaskedLMHead(
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,
diff --git a/keras_nlp/models/albert/albert_classifier_test.py b/keras_nlp/models/albert/albert_classifier_test.py
index 622e8bae..40fec534 100644
--- a/keras_nlp/models/albert/albert_classifier_test.py
+++ b/keras_nlp/models/albert/albert_classifier_test.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for BERT classification model."""
+"""Tests for ALBERT classification model."""
 
 import io
 import os
@@ -57,6 +57,7 @@ class AlbertClassifierTest(tf.test.TestCase, parameterized.TestCase):
             unk_piece="<unk>",
             bos_piece="[CLS]",
             eos_piece="[SEP]",
+            user_defined_symbols="[MASK]",
         )
         self.proto = bytes_io.getvalue()
 
diff --git a/keras_nlp/models/albert/albert_masked_lm_preprocessor_test.py b/keras_nlp/models/albert/albert_masked_lm_preprocessor_test.py
new file mode 100644
index 00000000..c77f24c2
--- /dev/null
+++ b/keras_nlp/models/albert/albert_masked_lm_preprocessor_test.py
@@ -0,0 +1,162 @@
+# 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.
+"""Tests for ALBERT masked language model preprocessor layer."""
+
+import io
+import os
+
+import sentencepiece
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.albert.albert_masked_lm_preprocessor import (
+    AlbertMaskedLMPreprocessor,
+)
+from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
+
+
+class AlbertMaskedLMPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        vocab_data = tf.data.Dataset.from_tensor_slices(
+            ["the quick brown fox", "the earth is round"]
+        )
+
+        bytes_io = io.BytesIO()
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=vocab_data.as_numpy_iterator(),
+            model_writer=bytes_io,
+            vocab_size=12,
+            model_type="WORD",
+            pad_id=0,
+            unk_id=1,
+            bos_id=2,
+            eos_id=3,
+            pad_piece="<pad>",
+            unk_piece="<unk>",
+            bos_piece="[CLS]",
+            eos_piece="[SEP]",
+            user_defined_symbols="[MASK]",
+        )
+
+        proto = bytes_io.getvalue()
+
+        tokenizer = AlbertTokenizer(proto=proto)
+
+        self.preprocessor = AlbertMaskedLMPreprocessor(
+            tokenizer=tokenizer,
+            # Simplify out testing by masking every available token.
+            mask_selection_rate=1.0,
+            mask_token_rate=1.0,
+            random_token_rate=0.0,
+            mask_selection_length=4,
+            sequence_length=12,
+        )
+
+    def test_preprocess_strings(self):
+        input_data = "the quick brown fox"
+
+        x, y, sw = self.preprocessor(input_data)
+        self.assertAllEqual(
+            x["token_ids"], [2, 4, 4, 4, 4, 3, 0, 0, 0, 0, 0, 0]
+        )
+        self.assertAllEqual(
+            x["padding_mask"], [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]
+        )
+        self.assertAllEqual(x["mask_positions"], [1, 2, 3, 4])
+        self.assertAllEqual(y, [5, 10, 6, 8])
+        self.assertAllEqual(sw, [1.0, 1.0, 1.0, 1.0])
+
+    def test_preprocess_list_of_strings(self):
+        input_data = ["the quick brown fox"] * 4
+
+        x, y, sw = self.preprocessor(input_data)
+        self.assertAllEqual(
+            x["token_ids"], [[2, 4, 4, 4, 4, 3, 0, 0, 0, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(
+            x["padding_mask"], [[1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(x["mask_positions"], [[1, 2, 3, 4]] * 4)
+        self.assertAllEqual(y, [[5, 10, 6, 8]] * 4)
+        self.assertAllEqual(sw, [[1.0, 1.0, 1.0, 1.0]] * 4)
+
+    def test_preprocess_dataset(self):
+        sentences = tf.constant(["the quick brown fox"] * 4)
+        ds = tf.data.Dataset.from_tensor_slices(sentences)
+        ds = ds.map(self.preprocessor)
+        x, y, sw = ds.batch(4).take(1).get_single_element()
+        self.assertAllEqual(
+            x["token_ids"], [[2, 4, 4, 4, 4, 3, 0, 0, 0, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(
+            x["padding_mask"], [[1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(x["mask_positions"], [[1, 2, 3, 4]] * 4)
+        self.assertAllEqual(y, [[5, 10, 6, 8]] * 4)
+        self.assertAllEqual(sw, [[1.0, 1.0, 1.0, 1.0]] * 4)
+
+    def test_mask_multiple_sentences(self):
+        sentence_one = tf.constant("the quick")
+        sentence_two = tf.constant("brown fox")
+
+        x, y, sw = self.preprocessor((sentence_one, sentence_two))
+        self.assertAllEqual(
+            x["token_ids"], [2, 4, 4, 3, 4, 4, 3, 0, 0, 0, 0, 0]
+        )
+        self.assertAllEqual(
+            x["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0]
+        )
+        self.assertAllEqual(x["mask_positions"], [1, 2, 4, 5])
+        self.assertAllEqual(y, [5, 10, 6, 8])
+        self.assertAllEqual(sw, [1.0, 1.0, 1.0, 1.0])
+
+    def test_no_masking_zero_rate(self):
+        no_mask_preprocessor = AlbertMaskedLMPreprocessor(
+            self.preprocessor.tokenizer,
+            mask_selection_rate=0.0,
+            mask_selection_length=4,
+            sequence_length=12,
+        )
+        input_data = "the quick brown fox"
+
+        x, y, sw = no_mask_preprocessor(input_data)
+        self.assertAllEqual(
+            x["token_ids"], [2, 5, 10, 6, 8, 3, 0, 0, 0, 0, 0, 0]
+        )
+        self.assertAllEqual(
+            x["padding_mask"], [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]
+        )
+        self.assertAllEqual(x["mask_positions"], [0, 0, 0, 0])
+        self.assertAllEqual(y, [0, 0, 0, 0])
+        self.assertAllEqual(sw, [0.0, 0.0, 0.0, 0.0])
+
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
+        input_data = tf.constant(["the quick brown fox"])
+
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = self.preprocessor(inputs)
+        model = keras.Model(inputs, outputs)
+
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
+
+        restored_model = keras.models.load_model(path)
+        outputs = model(input_data)[0]["token_ids"]
+        restored_outputs = restored_model(input_data)[0]["token_ids"]
+        self.assertAllEqual(outputs, restored_outputs)
diff --git a/keras_nlp/models/albert/albert_masked_lm_test.py b/keras_nlp/models/albert/albert_masked_lm_test.py
new file mode 100644
index 00000000..557b1bd4
--- /dev/null
+++ b/keras_nlp/models/albert/albert_masked_lm_test.py
@@ -0,0 +1,152 @@
+# 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.
+"""Tests for ALBERT masked language model."""
+
+import io
+import os
+
+import sentencepiece
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.albert.albert_backbone import AlbertBackbone
+from keras_nlp.models.albert.albert_masked_lm import AlbertMaskedLM
+from keras_nlp.models.albert.albert_masked_lm_preprocessor import (
+    AlbertMaskedLMPreprocessor,
+)
+from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
+
+
+class AlbertMaskedLMTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        self.backbone = AlbertBackbone(
+            vocabulary_size=1000,
+            num_layers=2,
+            num_heads=2,
+            embedding_dim=128,
+            hidden_dim=64,
+            intermediate_dim=128,
+            max_sequence_length=128,
+        )
+        vocab_data = tf.data.Dataset.from_tensor_slices(
+            ["the quick brown fox", "the earth is round", "an eagle flew"]
+        )
+
+        bytes_io = io.BytesIO()
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=vocab_data.as_numpy_iterator(),
+            model_writer=bytes_io,
+            vocab_size=15,
+            model_type="WORD",
+            pad_id=0,
+            unk_id=1,
+            bos_id=2,
+            eos_id=3,
+            pad_piece="<pad>",
+            unk_piece="<unk>",
+            bos_piece="[CLS]",
+            eos_piece="[SEP]",
+            user_defined_symbols="[MASK]",
+        )
+
+        proto = bytes_io.getvalue()
+
+        tokenizer = AlbertTokenizer(proto=proto)
+
+        self.preprocessor = AlbertMaskedLMPreprocessor(
+            tokenizer=tokenizer,
+            # Simplify out testing by masking every available token.
+            mask_selection_rate=1.0,
+            mask_token_rate=1.0,
+            random_token_rate=0.0,
+            mask_selection_length=5,
+            sequence_length=12,
+        )
+        self.masked_lm = AlbertMaskedLM(
+            self.backbone,
+            preprocessor=self.preprocessor,
+        )
+        self.masked_lm_no_preprocessing = AlbertMaskedLM(
+            self.backbone,
+            preprocessor=None,
+        )
+
+        self.raw_batch = tf.constant(
+            [
+                "quick brown fox",
+                "eagle flew over fox",
+                "the eagle flew quick",
+                "a brown eagle",
+            ]
+        )
+        self.preprocessed_batch = self.preprocessor(self.raw_batch)[0]
+        self.raw_dataset = tf.data.Dataset.from_tensor_slices(
+            self.raw_batch
+        ).batch(2)
+        self.preprocessed_dataset = self.raw_dataset.map(self.preprocessor)
+
+    def test_valid_call_masked_lm(self):
+        self.masked_lm(self.preprocessed_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_albert_masked_lm_predict(self, jit_compile):
+        self.masked_lm.compile(jit_compile=jit_compile)
+        self.masked_lm.predict(self.raw_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_albert_masked_lm_predict_no_preprocessing(self, jit_compile):
+        self.masked_lm_no_preprocessing.compile(jit_compile=jit_compile)
+        self.masked_lm_no_preprocessing.predict(self.preprocessed_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_albert_masked_lm_fit(self, jit_compile):
+        self.masked_lm.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            jit_compile=jit_compile,
+        )
+        self.masked_lm.fit(self.raw_dataset)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_albert_masked_lm_fit_no_preprocessing(self, jit_compile):
+        self.masked_lm_no_preprocessing.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            jit_compile=jit_compile,
+        )
+        self.masked_lm_no_preprocessing.fit(self.preprocessed_dataset)
+
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        self.masked_lm.save(save_path, save_format=save_format)
+        restored_model = keras.models.load_model(save_path)
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, AlbertMaskedLM)
+
+        model_output = self.masked_lm(self.preprocessed_batch)
+        restored_output = restored_model(self.preprocessed_batch)
+
+        self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/models/albert/albert_preprocessor_test.py b/keras_nlp/models/albert/albert_preprocessor_test.py
index ee6038f8..53639517 100644
--- a/keras_nlp/models/albert/albert_preprocessor_test.py
+++ b/keras_nlp/models/albert/albert_preprocessor_test.py
@@ -35,7 +35,7 @@ class AlbertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         sentencepiece.SentencePieceTrainer.train(
             sentence_iterator=vocab_data.as_numpy_iterator(),
             model_writer=bytes_io,
-            vocab_size=10,
+            vocab_size=12,
             model_type="WORD",
             pad_id=0,
             unk_id=1,
@@ -45,6 +45,7 @@ class AlbertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
             unk_piece="<unk>",
             bos_piece="[CLS]",
             eos_piece="[SEP]",
+            user_defined_symbols="[MASK]",
         )
         self.proto = bytes_io.getvalue()
 
@@ -57,7 +58,7 @@ class AlbertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         input_data = "the quick brown fox"
         output = self.preprocessor(input_data)
         self.assertAllEqual(
-            output["token_ids"], [2, 4, 9, 5, 7, 3, 0, 0, 0, 0, 0, 0]
+            output["token_ids"], [2, 5, 10, 6, 8, 3, 0, 0, 0, 0, 0, 0]
         )
         self.assertAllEqual(
             output["segment_ids"], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
@@ -72,7 +73,7 @@ class AlbertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         output = self.preprocessor(input_data)
         self.assertAllEqual(
             output["token_ids"],
-            [[2, 4, 9, 5, 7, 3, 0, 0, 0, 0, 0, 0]] * 4,
+            [[2, 5, 10, 6, 8, 3, 0, 0, 0, 0, 0, 0]] * 4,
         )
         self.assertAllEqual(
             output["segment_ids"], [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]] * 4
@@ -88,7 +89,7 @@ class AlbertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         x_out, y_out, sw_out = self.preprocessor(x, y, sw)
         self.assertAllEqual(
             x_out["token_ids"],
-            [[2, 4, 9, 5, 7, 3, 0, 0, 0, 0, 0, 0]] * 4,
+            [[2, 5, 10, 6, 8, 3, 0, 0, 0, 0, 0, 0]] * 4,
         )
         self.assertAllEqual(
             x_out["segment_ids"], [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]] * 4
@@ -108,7 +109,7 @@ class AlbertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         x_out, y_out, sw_out = ds.batch(4).take(1).get_single_element()
         self.assertAllEqual(
             x_out["token_ids"],
-            [[2, 4, 9, 5, 7, 3, 0, 0, 0, 0, 0, 0]] * 4,
+            [[2, 5, 10, 6, 8, 3, 0, 0, 0, 0, 0, 0]] * 4,
         )
         self.assertAllEqual(
             x_out["segment_ids"], [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]] * 4
@@ -125,7 +126,7 @@ class AlbertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         output = self.preprocessor((sentence_one, sentence_two))
         self.assertAllEqual(
             output["token_ids"],
-            [2, 4, 9, 5, 7, 3, 4, 6, 3, 0, 0, 0],
+            [2, 5, 10, 6, 8, 3, 5, 7, 3, 0, 0, 0],
         )
         self.assertAllEqual(
             output["segment_ids"], [0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0]
@@ -142,7 +143,7 @@ class AlbertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
         output = self.preprocessor((sentence_one, sentence_two))
         self.assertAllEqual(
             output["token_ids"],
-            [[2, 4, 9, 5, 7, 3, 4, 6, 3, 0, 0, 0]] * 4,
+            [[2, 5, 10, 6, 8, 3, 5, 7, 3, 0, 0, 0]] * 4,
         )
         self.assertAllEqual(
             output["segment_ids"], [[0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0]] * 4
diff --git a/keras_nlp/models/albert/albert_tokenizer_test.py b/keras_nlp/models/albert/albert_tokenizer_test.py
index d80657db..e97af4ee 100644
--- a/keras_nlp/models/albert/albert_tokenizer_test.py
+++ b/keras_nlp/models/albert/albert_tokenizer_test.py
@@ -34,7 +34,7 @@ class AlbertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
         sentencepiece.SentencePieceTrainer.train(
             sentence_iterator=vocab_data.as_numpy_iterator(),
             model_writer=bytes_io,
-            vocab_size=10,
+            vocab_size=12,
             model_type="WORD",
             pad_id=0,
             unk_id=1,
@@ -44,6 +44,7 @@ class AlbertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
             unk_piece="<unk>",
             bos_piece="[CLS]",
             eos_piece="[SEP]",
+            user_defined_symbols="[MASK]",
         )
         self.proto = bytes_io.getvalue()
 
@@ -52,20 +53,21 @@ class AlbertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
     def test_tokenize(self):
         input_data = "the quick brown fox"
         output = self.tokenizer(input_data)
-        self.assertAllEqual(output, [4, 9, 5, 7])
+        self.assertAllEqual(output, [5, 10, 6, 8])
 
     def test_tokenize_batch(self):
         input_data = tf.constant(["the quick brown fox", "the earth is round"])
         output = self.tokenizer(input_data)
-        self.assertAllEqual(output, [[4, 9, 5, 7], [4, 6, 8, 1]])
+        self.assertAllEqual(output, [[5, 10, 6, 8], [5, 7, 9, 11]])
 
     def test_detokenize(self):
-        input_data = tf.constant([[4, 9, 5, 7]])
+        input_data = tf.constant([[5, 10, 6, 8]])
         output = self.tokenizer.detokenize(input_data)
         self.assertEqual(output, tf.constant(["the quick brown fox"]))
 
     def test_vocabulary_size(self):
-        self.assertEqual(self.tokenizer.vocabulary_size(), 10)
+        tokenizer = AlbertTokenizer(proto=self.proto)
+        self.assertEqual(tokenizer.vocabulary_size(), 12)
 
     def test_errors_missing_special_tokens(self):
         bytes_io = io.BytesIO()