Repository: keras-team/keras-nlp
Pull request #621 (2022-12-28 23:51:47)

Create `Backbone` base class
Pilot for how base classes could improve the maintainability of our library. The `Backbone` class can be easily extended to other model types and `Task`, `Preprocessor`, and `Tokenizer` base classes can also be added.

I've also simplified the `from_preset` docstring so we don't have to write a lot of custom copy for each model class. We already have all the information we need in the class docstring.

Part of #530


Patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index bdd60425..21d83940 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
+from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
 from keras_nlp.models.bert.bert_backbone import BertBackbone
 from keras_nlp.models.bert.bert_classifier import BertClassifier
 from keras_nlp.models.bert.bert_preprocessor import BertPreprocessor
diff --git a/keras_nlp/models/albert/__init__.py b/keras_nlp/models/albert/__init__.py
new file mode 100644
index 00000000..6e4df4e7
--- /dev/null
+++ b/keras_nlp/models/albert/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/keras_nlp/models/albert/albert_preprocessor.py b/keras_nlp/models/albert/albert_preprocessor.py
new file mode 100644
index 00000000..61d37c34
--- /dev/null
+++ b/keras_nlp/models/albert/albert_preprocessor.py
@@ -0,0 +1,201 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""ALBERT preprocessor layer."""
+
+from tensorflow import keras
+
+from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
+from keras_nlp.utils.keras_utils import (
+    convert_inputs_to_list_of_tensor_segments,
+)
+from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
+from keras_nlp.utils.python_utils import classproperty
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class AlbertPreprocessor(keras.layers.Layer):
+    """An ALBERT preprocessing layer which tokenizes and packs inputs.
+
+    This preprocessing layer will do three things:
+
+     - Tokenize any number of input segments using the `tokenizer`.
+     - Pack the inputs together using a `keras_nlp.layers.MultiSegmentPacker`.
+       with the appropriate `"[CLS]"`, `"[SEP]"` and `"<pad>"` tokens.
+     - Construct a dictionary with keys `"token_ids"`, `"segment_ids"` and
+       `"padding_mask"`, that can be passed directly to
+       `keras_nlp.models.AlbertBackbone`.
+
+    This layer can be used directly with `tf.data.Dataset.map` to preprocess
+    string data in the `(x, y, sample_weight)` format used by
+    `keras.Model.fit`.
+
+    The call method of this layer accepts three arguments, `x`, `y`, and
+    `sample_weight`. `x` can be a python string or tensor representing a single
+    segment, a list of python strings representing a batch of single segments,
+    or a list of tensors representing multiple segments to be packed together.
+    `y` and `sample_weight` are both optional, can have any format, and will be
+    passed through unaltered.
+
+    Special care should be taken when using `tf.data` to map over an unlabeled
+    tuple of string segments. `tf.data.Dataset.map` will unpack this tuple
+    directly into the call arguments of this layer, rather than forward all
+    argument to `x`. To handle this case, it is recommended to  explicitly call
+    the layer, e.g. `ds.map(lambda seg1, seg2: preprocessor(x=(seg1, seg2)))`.
+
+    Args:
+        tokenizer: A `keras_nlp.models.AlbertTokenizer` instance.
+        sequence_length: The length of the packed inputs.
+        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
+    tokenizer = keras_nlp.models.AlbertTokenizer(proto="model.spm")
+    preprocessor = keras_nlp.models.AlbertPreprocessor(
+        tokenizer=tokenizer,
+        sequence_length=10,
+    )
+
+    # Tokenize and pack a single sentence.
+    sentence = tf.constant("The quick brown fox jumped.")
+    preprocessor(sentence)
+    # Same output.
+    preprocessor("The quick brown fox jumped.")
+
+    # Tokenize and a batch of single sentences.
+    sentences = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    preprocessor(sentences)
+    # Same output.
+    preprocessor(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+
+    # Tokenize and pack a sentence pair.
+    first_sentence = tf.constant("The quick brown fox jumped.")
+    second_sentence = tf.constant("The fox tripped.")
+    preprocessor((first_sentence, second_sentence))
+
+    # Map a dataset to preprocess a single sentence.
+    features = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    labels = tf.constant([0, 1])
+    ds = tf.data.Dataset.from_tensor_slices((features, labels))
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Map a dataset to preprocess sentence pairs.
+    first_sentences = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    second_sentences = tf.constant(
+        ["The fox tripped.", "Oh look, a whale."]
+    )
+    labels = tf.constant([1, 1])
+    ds = tf.data.Dataset.from_tensor_slices(
+        (
+            (first_sentences, second_sentences), labels
+        )
+    )
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Map a dataset to preprocess unlabeled sentence pairs.
+    first_sentences = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    second_sentences = tf.constant(
+        ["The fox tripped.", "Oh look, a whale."]
+    )
+    ds = tf.data.Dataset.from_tensor_slices((first_sentences, second_sentences))
+    # Watch out for tf.data's default unpacking of tuples here!
+    # Best to invoke the `preprocessor` directly in this case.
+    ds = ds.map(
+        lambda s1, s2: preprocessor(x=(s1, s2)),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+    ```
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        sequence_length=512,
+        truncate="round_robin",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self._tokenizer = tokenizer
+        self.packer = MultiSegmentPacker(
+            start_value=self.tokenizer.cls_token_id,
+            end_value=self.tokenizer.sep_token_id,
+            pad_value=self.tokenizer.pad_token_id,
+            truncate=truncate,
+            sequence_length=sequence_length,
+        )
+
+    @property
+    def tokenizer(self):
+        """The `keras_nlp.models.AlbertTokenizer` used to tokenize strings."""
+        return self._tokenizer
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "tokenizer": keras.layers.serialize(self.tokenizer),
+                "sequence_length": self.packer.sequence_length,
+                "truncate": self.packer.truncate,
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        if "tokenizer" in config and isinstance(config["tokenizer"], dict):
+            config["tokenizer"] = keras.layers.deserialize(config["tokenizer"])
+        return cls(**config)
+
+    def call(self, x, y=None, sample_weight=None):
+        x = convert_inputs_to_list_of_tensor_segments(x)
+        x = [self.tokenizer(segment) for segment in x]
+        token_ids, segment_ids = self.packer(x)
+        x = {
+            "token_ids": token_ids,
+            "segment_ids": segment_ids,
+            "padding_mask": token_ids != self.tokenizer.pad_token_id,
+        }
+        return pack_x_y_sample_weight(x, y, sample_weight)
+
+    @classproperty
+    def presets(cls):
+        return {}
+
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        sequence_length=None,
+        truncate="round_robin",
+        **kwargs,
+    ):
+        raise NotImplementedError
diff --git a/keras_nlp/models/albert/albert_tokenizer.py b/keras_nlp/models/albert/albert_tokenizer.py
new file mode 100644
index 00000000..e9e5ac19
--- /dev/null
+++ b/keras_nlp/models/albert/albert_tokenizer.py
@@ -0,0 +1,95 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""ALBERT tokenizer."""
+
+
+from tensorflow import keras
+
+from keras_nlp.tokenizers.sentence_piece_tokenizer import SentencePieceTokenizer
+from keras_nlp.utils.python_utils import classproperty
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class AlbertTokenizer(SentencePieceTokenizer):
+    """ALBERT tokenizer layer based on SentencePiece.
+
+    This tokenizer class will tokenize raw strings into integer sequences and
+    is based on `keras_nlp.tokenizers.SentencePieceTokenizer`. Unlike the
+    underlying tokenizer, it will check for all special tokens needed by
+    ALBERT models and provides a `from_preset()` method to automatically
+    download a matching vocabulary for a ALBERT preset.
+
+    This tokenizer does not provide truncation or padding of inputs. It can be
+    combined with a `keras_nlp.models.AlbertPreprocessor` layer for input
+    packing.
+
+    If input is a batch of strings (rank > 0), the layer will output a
+    `tf.RaggedTensor` where the last dimension of the output is ragged.
+
+    If input is a scalar string (rank == 0), the layer will output a dense
+    `tf.Tensor` with static shape `[None]`.
+
+    Args:
+        proto: Either a `string` path to a SentencePiece proto file, or a
+            `bytes` object with a serialized SentencePiece proto. See the
+            [SentencePiece repository](https://github.com/google/sentencepiece)
+            for more details on the format.
+
+    Examples:
+
+    ```python
+    tokenizer = keras_nlp.models.AlbertTokenizer(proto="model.spm")
+
+    # Batched inputs.
+    tokenizer(["the quick brown fox", "the earth is round"])
+
+    # Unbatched inputs.
+    tokenizer("the quick brown fox")
+
+    # Detokenization.
+    tokenizer.detokenize(tf.constant([[[2, 14, 2231, 886, 2385, 3]]))
+    ```
+    """
+
+    def __init__(self, proto, **kwargs):
+        super().__init__(proto=proto, **kwargs)
+
+        # Check for necessary special tokens.
+        cls_token = "[CLS]"
+        sep_token = "[SEP]"
+        pad_token = "<pad>"
+        for token in [cls_token, sep_token, pad_token]:
+            if token not in self.get_vocabulary():
+                raise ValueError(
+                    f"Cannot find token `'{token}'` in the provided "
+                    f"`vocabulary`. Please provide `'{token}'` in your "
+                    "`vocabulary` or use a pretrained `vocabulary` name."
+                )
+
+        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)
+
+    @classproperty
+    def presets(cls):
+        return {}
+
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        **kwargs,
+    ):
+        raise NotImplementedError
diff --git a/keras_nlp/models/backbone.py b/keras_nlp/models/backbone.py
new file mode 100644
index 00000000..510235a9
--- /dev/null
+++ b/keras_nlp/models/backbone.py
@@ -0,0 +1,89 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Base class for Backbone models."""
+
+import os
+
+from tensorflow import keras
+
+from keras_nlp.utils.python_utils import classproperty
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class Backbone(keras.Model):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
+
+    @classproperty
+    def presets(cls):
+        return {}
+
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        load_weights=True,
+        **kwargs,
+    ):
+        """Instantiate {{model_name}} model from preset architecture and weights.
+
+        Args:
+            preset: string. Must be one of "{{preset_names}}".
+            load_weights: Whether to load pre-trained weights into model.
+                Defaults to `True`.
+
+        Examples:
+        ```python
+        # Load architecture and weights from preset
+        model = {{model_name}}.from_preset("{{example_preset_name}}")
+
+        # Load randomly initialized model from preset architecture
+        model = {{model_name}}.from_preset(
+            "{{example_preset_name}}",
+            load_weights=False
+        )
+        ```
+        """
+
+        if not cls.presets:
+            raise NotImplementedError(
+                "No presets have been created for this class."
+            )
+
+        if preset not in cls.presets:
+            raise ValueError(
+                "`preset` must be one of "
+                f"""{", ".join(cls.presets)}. Received: {preset}."""
+            )
+        metadata = cls.presets[preset]
+        config = metadata["config"]
+        model = cls.from_config({**config, **kwargs})
+
+        if not load_weights:
+            return model
+
+        weights = keras.utils.get_file(
+            "model.h5",
+            metadata["weights_url"],
+            cache_subdir=os.path.join("models", preset),
+            file_hash=metadata["weights_hash"],
+        )
+
+        model.load_weights(weights)
+        return model
diff --git a/keras_nlp/models/bert/bert_backbone.py b/keras_nlp/models/bert/bert_backbone.py
index 2659b7cb..8339d0f8 100644
--- a/keras_nlp/models/bert/bert_backbone.py
+++ b/keras_nlp/models/bert/bert_backbone.py
@@ -12,16 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""BERT backbone models."""
+"""BERT backbone model."""
 
 import copy
-import os
 
 import tensorflow as tf
 from tensorflow import keras
 
 from keras_nlp.layers.position_embedding import PositionEmbedding
 from keras_nlp.layers.transformer_encoder import TransformerEncoder
+from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.bert.bert_presets import backbone_presets
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
@@ -32,7 +32,7 @@ def bert_kernel_initializer(stddev=0.02):
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class BertBackbone(keras.Model):
+class BertBackbone(Backbone):
     """BERT encoder network.
 
     This class implements a bi-directional Transformer-based encoder as
@@ -76,7 +76,11 @@ class BertBackbone(keras.Model):
         ),
     }
 
-    # Randomly initialized BERT encoder
+    # Pretrained BERT encoder
+    model = keras_nlp.models.BertBackbone.from_preset("base_base_en_uncased")
+    output = model(input_data)
+
+    # Randomly initialized BERT encoder with a custom config
     model = keras_nlp.models.BertBackbone(
         vocabulary_size=30552,
         num_layers=12,
@@ -212,71 +216,18 @@ class BertBackbone(keras.Model):
             "trainable": self.trainable,
         }
 
-    @classmethod
-    def from_config(cls, config):
-        return cls(**config)
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
 
     @classmethod
-    @format_docstring(names=", ".join(backbone_presets))
-    def from_preset(
-        cls,
-        preset,
-        load_weights=True,
-        **kwargs,
-    ):
-        """Instantiate BERT model from preset architecture and weights.
-
-        Args:
-            preset: string. Must be one of {{names}}.
-            load_weights: Whether to load pre-trained weights into model.
-                Defaults to `True`.
+    def from_preset(cls, preset, load_weights=True, **kwargs):
+        return super().from_preset(preset, load_weights, **kwargs)
 
-        Examples:
-        ```python
-        input_data = {
-            "token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
-            "segment_ids": tf.constant(
-                [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-            "padding_mask": tf.constant(
-                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-        }
-
-        # Load architecture and weights from preset
-        model = BertBackbone.from_preset("bert_base_en_uncased")
-        output = model(input_data)
-
-        # Load randomly initialized model from preset architecture
-        model = BertBackbone.from_preset(
-            "bert_base_en_uncased",
-            load_weights=False
-        )
-        output = model(input_data)
-        ```
-        """
-        if preset not in cls.presets:
-            raise ValueError(
-                "`preset` must be one of "
-                f"""{", ".join(cls.presets)}. Received: {preset}."""
-            )
-        metadata = cls.presets[preset]
-        config = metadata["config"]
-        model = cls.from_config({**config, **kwargs})
-
-        if not load_weights:
-            return model
-
-        weights = keras.utils.get_file(
-            "model.h5",
-            metadata["weights_url"],
-            cache_subdir=os.path.join("models", preset),
-            file_hash=metadata["weights_hash"],
-        )
 
-        model.load_weights(weights)
-        return model
+BertBackbone.from_preset.__func__.__doc__ = Backbone.from_preset.__doc__
+format_docstring(
+    model_name=BertBackbone.__name__,
+    example_preset_name="bert_base_en_uncased",
+    preset_names='", "'.join(BertBackbone.presets),
+)(BertBackbone.from_preset.__func__)
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_backbone.py b/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
index 6438691c..585a15a5 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
@@ -15,11 +15,11 @@
 """DeBERTa backbone model."""
 
 import copy
-import os
 
 import tensorflow as tf
 from tensorflow import keras
 
+from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.deberta_v3.deberta_v3_presets import backbone_presets
 from keras_nlp.models.deberta_v3.disentangled_attention_encoder import (
     DisentangledAttentionEncoder,
@@ -34,7 +34,7 @@ def deberta_kernel_initializer(stddev=0.02):
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class DebertaV3Backbone(keras.Model):
+class DebertaV3Backbone(Backbone):
     """DeBERTa encoder network.
 
     This network implements a bi-directional Transformer-based encoder as
@@ -45,7 +45,7 @@ class DebertaV3Backbone(keras.Model):
 
     The default constructor gives a fully customizable, randomly initialized
     DeBERTa encoder with any number of layers, heads, and embedding
-    dimensions. To load preset architectures and weights, use the `from_presets`
+    dimensions. To load preset architectures and weights, use the `from_preset`
     constructor.
 
     Disclaimer: Pre-trained models are provided on an "as is" basis, without
@@ -76,7 +76,13 @@ class DebertaV3Backbone(keras.Model):
             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
     }
 
-    # Randomly initialized DeBERTa model
+    # Pretrained DeBERTa encoder
+    model = keras_nlp.models.DebertaV3Backbone.from_preset(
+        "deberta_base_en",
+    )
+    output = model(input_data)
+
+    # Randomly initialized DeBERTa encoder with custom config
     model = keras_nlp.models.DebertaV3Backbone(
         vocabulary_size=128100,
         num_layers=12,
@@ -86,7 +92,6 @@ class DebertaV3Backbone(keras.Model):
         max_sequence_length=512,
         bucket_size=256,
     )
-
     # Call the model on the input data.
     output = model(input_data)
     ```
@@ -194,69 +199,18 @@ class DebertaV3Backbone(keras.Model):
             "trainable": self.trainable,
         }
 
-    @classmethod
-    def from_config(cls, config):
-        return cls(**config)
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
 
     @classmethod
-    @format_docstring(names=", ".join(backbone_presets))
-    def from_preset(
-        cls,
-        preset,
-        load_weights=True,
-        **kwargs,
-    ):
-        """Instantiate DeBERTa model from preset architecture and weights.
-
-        Args:
-            preset: string. Must be one of {{names}}.
-            load_weights: Whether to load pre-trained weights into model.
-                Defaults to `True`.
+    def from_preset(cls, preset, load_weights=True, **kwargs):
+        return super().from_preset(preset, load_weights, **kwargs)
 
-        Examples:
-        ```python
-        input_data = {
-            "token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
-            "padding_mask": tf.constant(
-                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-        }
-
-        # Load architecture and weights from preset
-        model = keras_nlp.models.DebertaV3Backbone.from_preset(
-            "deberta_base_en",
-        )
-        output = model(input_data)
-
-        # Load randomly initialized model from preset architecture
-        model = keras_nlp.models.DebertaV3Backbone.from_preset(
-            "deberta_base_en", load_weights=False
-        )
-        output = model(input_data)
-        ```
-        """
-        if preset not in cls.presets:
-            raise ValueError(
-                "`preset` must be one of "
-                f"""{", ".join(cls.presets)}. Received: {preset}."""
-            )
-        metadata = cls.presets[preset]
-        config = metadata["config"]
-        model = cls.from_config({**config, **kwargs})
-
-        if not load_weights:
-            return model
-
-        weights = keras.utils.get_file(
-            "model.h5",
-            metadata["weights_url"],
-            cache_subdir=os.path.join("models", preset),
-            file_hash=metadata["weights_hash"],
-        )
 
-        model.load_weights(weights)
-        return model
+DebertaV3Backbone.from_preset.__func__.__doc__ = Backbone.from_preset.__doc__
+format_docstring(
+    model_name=DebertaV3Backbone.__name__,
+    example_preset_name="deberta_base_en",
+    preset_names='", "'.join(DebertaV3Backbone.presets),
+)(DebertaV3Backbone.from_preset.__func__)
diff --git a/keras_nlp/models/distil_bert/distil_bert_backbone.py b/keras_nlp/models/distil_bert/distil_bert_backbone.py
index 3c32589d..320010fd 100644
--- a/keras_nlp/models/distil_bert/distil_bert_backbone.py
+++ b/keras_nlp/models/distil_bert/distil_bert_backbone.py
@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""DistilBERT backbone models."""
+"""DistilBERT backbone model."""
 
 import copy
-import os
 
 import tensorflow as tf
 from tensorflow import keras
@@ -24,6 +23,7 @@ from keras_nlp.layers.token_and_position_embedding import (
     TokenAndPositionEmbedding,
 )
 from keras_nlp.layers.transformer_encoder import TransformerEncoder
+from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.distil_bert.distil_bert_presets import backbone_presets
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
@@ -34,7 +34,7 @@ def distilbert_kernel_initializer(stddev=0.02):
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class DistilBertBackbone(keras.Model):
+class DistilBertBackbone(Backbone):
     """DistilBERT encoder network.
 
     This network implements a bi-directional Transformer-based encoder as
@@ -45,7 +45,7 @@ class DistilBertBackbone(keras.Model):
 
     The default constructor gives a fully customizable, randomly initialized
     DistilBERT encoder with any number of layers, heads, and embedding
-    dimensions. To load preset architectures and weights, use the `from_presets`
+    dimensions. To load preset architectures and weights, use the `from_preset`
     constructor.
 
     Disclaimer: Pre-trained models are provided on an "as is" basis, without
@@ -76,7 +76,13 @@ class DistilBertBackbone(keras.Model):
         ),
     }
 
-    # Randomly initialized DistilBERT encoder
+    # Pretrained DistilBERT encoder
+    model = keras_nlp.models.DistilBertBackbone.from_preset(
+        "distil_bert_base_en_uncased"
+    )
+    output = model(input_data)
+
+    # Randomly initialized DistilBERT encoder with custom config
     model = keras_nlp.models.DistilBertBackbone(
         vocabulary_size=30552,
         num_layers=6,
@@ -173,69 +179,18 @@ class DistilBertBackbone(keras.Model):
             "trainable": self.trainable,
         }
 
-    @classmethod
-    def from_config(cls, config):
-        return cls(**config)
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
 
     @classmethod
-    @format_docstring(names=", ".join(backbone_presets))
-    def from_preset(
-        cls,
-        preset,
-        load_weights=True,
-        **kwargs,
-    ):
-        """Instantiate a DistilBERT model from preset architecture and weights.
-
-        Args:
-            preset: string. Must be one of {{names}}.
-            load_weights: Whether to load pre-trained weights into model.
-                Defaults to `True`.
-
-        Examples:
-        ```python
-        input_data = {
-            "token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
-            "padding_mask": tf.constant(
-                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-        }
-
-        # Load architecture and weights from preset
-        model = keras_nlp.models.DistilBertBackbone.from_preset(
-            "distil_bert_base_en_uncased"
-        )
-        output = model(input_data)
+    def from_preset(cls, preset, load_weights=True, **kwargs):
+        return super().from_preset(preset, load_weights, **kwargs)
 
-        # Load randomly initialized model from preset architecture
-        model = keras_nlp.models.DistilBertBackbone.from_preset(
-            "distil_bert_base_en_uncased", load_weights=False
-        )
-        output = model(input_data)
-        ```
-        """
-        if preset not in cls.presets:
-            raise ValueError(
-                "`preset` must be one of "
-                f"""{", ".join(cls.presets)}. Received: {preset}."""
-            )
-        metadata = cls.presets[preset]
-        config = metadata["config"]
-        model = cls.from_config({**config, **kwargs})
-
-        if not load_weights:
-            return model
-
-        weights = keras.utils.get_file(
-            "model.h5",
-            metadata["weights_url"],
-            cache_subdir=os.path.join("models", preset),
-            file_hash=metadata["weights_hash"],
-        )
 
-        model.load_weights(weights)
-        return model
+DistilBertBackbone.from_preset.__func__.__doc__ = Backbone.from_preset.__doc__
+format_docstring(
+    model_name=DistilBertBackbone.__name__,
+    example_preset_name="distil_bert_base_en_uncased",
+    preset_names='", "'.join(DistilBertBackbone.presets),
+)(DistilBertBackbone.from_preset.__func__)
diff --git a/keras_nlp/models/gpt2/gpt2_backbone.py b/keras_nlp/models/gpt2/gpt2_backbone.py
index 87383643..36ec03c5 100644
--- a/keras_nlp/models/gpt2/gpt2_backbone.py
+++ b/keras_nlp/models/gpt2/gpt2_backbone.py
@@ -12,16 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""GPT-2 backbone models."""
+"""GPT-2 backbone model."""
 
 import copy
-import os
 
 import tensorflow as tf
 from tensorflow import keras
 
 from keras_nlp.layers import PositionEmbedding
 from keras_nlp.layers import TransformerDecoder
+from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.gpt2.gpt2_presets import backbone_presets
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
@@ -32,7 +32,7 @@ def _gpt_2_kernel_initializer(stddev=0.02):
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class GPT2Backbone(keras.Model):
+class GPT2Backbone(Backbone):
     """GPT-2 core network with hyperparameters.
 
     This network implements a Transformer-based decoder network,
@@ -42,7 +42,7 @@ class GPT2Backbone(keras.Model):
 
     The default constructor gives a fully customizable, randomly initialized
     GPT-2 model with any number of layers, heads, and embedding
-    dimensions. To load preset architectures and weights, use the `from_presets`
+    dimensions. To load preset architectures and weights, use the `from_preset`
     constructor.
 
     Disclaimer: Pre-trained models are provided on an "as is" basis, without
@@ -73,7 +73,11 @@ class GPT2Backbone(keras.Model):
         ),
     }
 
-    # Randomly initialized GPT-2 decoder
+    # Pretrained GPT-2 decoder
+    model = GPT2Backbone.from_preset("gpt2_base_en")
+    output = model(input_data)
+
+    # Randomly initialized GPT-2 decoder with custom config
     model = keras_nlp.models.GPT2Backbone(
         vocabulary_size=50257,
         num_layers=12,
@@ -182,66 +186,18 @@ class GPT2Backbone(keras.Model):
             "trainable": self.trainable,
         }
 
-    @classmethod
-    def from_config(cls, config):
-        return cls(**config)
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
 
     @classmethod
-    @format_docstring(names=", ".join(backbone_presets))
-    def from_preset(
-        cls,
-        preset,
-        load_weights=True,
-        **kwargs,
-    ):
-        """Instantiate GPT-2 model from preset architecture and weights.
-
-        Args:
-            preset: string. Must be one of {{names}}.
-            load_weights: Whether to load pre-trained weights into model.
-                Defaults to `True`.
-
-        Examples:
-        ```python
-        input_data = {
-            "token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
-            "padding_mask": tf.constant(
-                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-        }
+    def from_preset(cls, preset, load_weights=True, **kwargs):
+        return super().from_preset(preset, load_weights, **kwargs)
 
-        # Load architecture and weights from preset
-        model = GPT2Backbone.from_preset("gpt2_base_en")
-        output = model(input_data)
-
-        # Load randomly initialized model from preset architecture
-        model = GPT2Backbone.from_preset("gpt2_base", load_weights=False)
-        output = model(input_data)
-        ```
-        """
-
-        if preset not in cls.presets:
-            raise ValueError(
-                "`preset` must be one of "
-                f"""{", ".join(cls.presets)}. Received: {preset}."""
-            )
-        metadata = cls.presets[preset]
-        config = metadata["config"]
-        model = cls.from_config({**config, **kwargs})
-
-        if not load_weights:
-            return model
-
-        weights = keras.utils.get_file(
-            "model.h5",
-            metadata["weights_url"],
-            cache_subdir=os.path.join("models", preset),
-            file_hash=metadata["weights_hash"],
-        )
 
-        model.load_weights(weights)
-        return model
+GPT2Backbone.from_preset.__func__.__doc__ = Backbone.from_preset.__doc__
+format_docstring(
+    model_name=GPT2Backbone.__name__,
+    example_preset_name="gpt2_base_en",
+    preset_names='", "'.join(GPT2Backbone.presets),
+)(GPT2Backbone.from_preset.__func__)
diff --git a/keras_nlp/models/roberta/roberta_backbone.py b/keras_nlp/models/roberta/roberta_backbone.py
index 99a322d0..28a823c9 100644
--- a/keras_nlp/models/roberta/roberta_backbone.py
+++ b/keras_nlp/models/roberta/roberta_backbone.py
@@ -12,16 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""RoBERTa backbone models."""
+"""RoBERTa backbone model."""
 
 import copy
-import os
 
 import tensorflow as tf
 from tensorflow import keras
 
 from keras_nlp.layers import TokenAndPositionEmbedding
 from keras_nlp.layers import TransformerEncoder
+from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.roberta.roberta_presets import backbone_presets
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
@@ -32,7 +32,7 @@ def roberta_kernel_initializer(stddev=0.02):
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class RobertaBackbone(keras.Model):
+class RobertaBackbone(Backbone):
     """RoBERTa encoder.
 
     This network implements a bi-directional Transformer-based encoder as
@@ -42,7 +42,7 @@ class RobertaBackbone(keras.Model):
 
     The default constructor gives a fully customizable, randomly initialized
     RoBERTa encoder with any number of layers, heads, and embedding
-    dimensions. To load preset architectures and weights, use the `from_presets`
+    dimensions. To load preset architectures and weights, use the `from_preset`
     constructor.
 
     Disclaimer: Pre-trained models are provided on an "as is" basis, without
@@ -71,7 +71,11 @@ class RobertaBackbone(keras.Model):
             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
     }
 
-    # Randomly initialized RoBERTa model
+    # Pretrained RoBERTa encoder
+    model = keras_nlp.models.RobertaBackbone.from_preset("roberta_base_en")
+    output = model(input_data)
+
+    # Randomly initialized RoBERTa model with custom config
     model = keras_nlp.models.RobertaBackbone(
         vocabulary_size=50265,
         num_layers=12,
@@ -80,8 +84,6 @@ class RobertaBackbone(keras.Model):
         intermediate_dim=3072,
         max_sequence_length=12
     )
-
-    # Call the model on the input data.
     output = model(input_data)
     ```
     """
@@ -171,67 +173,18 @@ class RobertaBackbone(keras.Model):
             "trainable": self.trainable,
         }
 
-    @classmethod
-    def from_config(cls, config):
-        return cls(**config)
-
     @classproperty
     def presets(cls):
         return copy.deepcopy(backbone_presets)
 
     @classmethod
-    @format_docstring(names=", ".join(backbone_presets))
-    def from_preset(
-        cls,
-        preset,
-        load_weights=True,
-        **kwargs,
-    ):
-        """Instantiate RoBERTa model from preset architecture and weights.
-
-        Args:
-            preset: string. Must be one of {{names}}.
-            load_weights: Whether to load pre-trained weights into model.
-                Defaults to `True`.
-
-        Examples:
-        ```python
-        input_data = {
-            "token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
-            "padding_mask": tf.constant(
-                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-        }
-
-        # Load architecture and weights from preset
-        model = keras_nlp.models.RobertaBackbone.from_preset("roberta_base_en")
-        output = model(input_data)
+    def from_preset(cls, preset, load_weights=True, **kwargs):
+        return super().from_preset(preset, load_weights, **kwargs)
 
-        # Load randomly initialized model from preset architecture
-        model = keras_nlp.models.RobertaBackbone.from_preset(
-            "roberta_base_en", load_weights=False
-        )
-        output = model(input_data)
-        ```
-        """
-        if preset not in cls.presets:
-            raise ValueError(
-                "`preset` must be one of "
-                f"""{", ".join(cls.presets)}. Received: {preset}."""
-            )
-        metadata = cls.presets[preset]
-        config = metadata["config"]
-        model = cls.from_config({**config, **kwargs})
-
-        if not load_weights:
-            return model
-
-        weights = keras.utils.get_file(
-            "model.h5",
-            metadata["weights_url"],
-            cache_subdir=os.path.join("models", preset),
-            file_hash=metadata["weights_hash"],
-        )
 
-        model.load_weights(weights)
-        return model
+RobertaBackbone.from_preset.__func__.__doc__ = Backbone.from_preset.__doc__
+format_docstring(
+    model_name=RobertaBackbone.__name__,
+    example_preset_name="roberta_base_en",
+    preset_names='", "'.join(RobertaBackbone.presets),
+)(RobertaBackbone.from_preset.__func__)
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py
index 820e5bf1..741ed02e 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone.py
@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""XLM-RoBERTa backbone models."""
+"""XLM-RoBERTa backbone model."""
 
 import copy
-import os
 
 from tensorflow import keras
 
+from keras_nlp.models.backbone import Backbone
 from keras_nlp.models.roberta import roberta_backbone
 from keras_nlp.models.xlm_roberta.xlm_roberta_presets import backbone_presets
 from keras_nlp.utils.python_utils import classproperty
@@ -37,7 +37,7 @@ class XLMRobertaBackbone(roberta_backbone.RobertaBackbone):
 
     The default constructor gives a fully customizable, randomly initialized
     RoBERTa encoder with any number of layers, heads, and embedding
-    dimensions. To load preset architectures and weights, use the `from_presets`
+    dimensions. To load preset architectures and weights, use the `from_preset`
     constructor.
 
     Disclaimer: Pre-trained models are provided on an "as is" basis, without
@@ -66,7 +66,13 @@ class XLMRobertaBackbone(roberta_backbone.RobertaBackbone):
             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
     }
 
-    # Randomly initialized XLM-R model
+    # Pretrained XLM-R encoder
+    model = keras_nlp.models.XLMRobertaBackbone.from_preset(
+        "xlm_roberta_base_multi",
+    )
+    output = model(input_data)
+
+    # Randomly initialized XLM-R model with custom config
     model = keras_nlp.models.XLMRobertaBackbone(
         vocabulary_size=250002,
         num_layers=12,
@@ -86,60 +92,13 @@ class XLMRobertaBackbone(roberta_backbone.RobertaBackbone):
         return copy.deepcopy(backbone_presets)
 
     @classmethod
-    @format_docstring(names=", ".join(backbone_presets))
-    def from_preset(
-        cls,
-        preset,
-        load_weights=True,
-        **kwargs,
-    ):
-        """Instantiate XLM-RoBERTa model from preset architecture and weights.
-
-        Args:
-            preset: string. Must be one of {{names}}.
-            load_weights: Whether to load pre-trained weights into model.
-                Defaults to `True`.
-
-        Examples:
-        ```python
-        input_data = {
-            "token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
-            "padding_mask": tf.constant(
-                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-        }
-
-        # Load architecture and weights from preset
-        model = keras_nlp.models.XLMRobertaBackbone.from_preset(
-            "xlm_roberta_base_multi",
-        )
-        output = model(input_data)
-
-        # Load randomly initialized model from preset architecture
-        model = keras_nlp.models.XLMRobertaBackbone.from_preset(
-            "xlm_roberta_base_multi", load_weights=False
-        )
-        output = model(input_data)
-        ```
-        """
-        if preset not in cls.presets:
-            raise ValueError(
-                "`preset` must be one of "
-                f"""{", ".join(cls.presets)}. Received: {preset}."""
-            )
-        metadata = cls.presets[preset]
-        config = metadata["config"]
-        model = cls.from_config({**config, **kwargs})
-
-        if not load_weights:
-            return model
-
-        weights = keras.utils.get_file(
-            "model.h5",
-            metadata["weights_url"],
-            cache_subdir=os.path.join("models", preset),
-            file_hash=metadata["weights_hash"],
-        )
-
-        model.load_weights(weights)
-        return model
+    def from_preset(cls, preset, load_weights=True, **kwargs):
+        return super().from_preset(preset, load_weights, **kwargs)
+
+
+XLMRobertaBackbone.from_preset.__func__.__doc__ = Backbone.from_preset.__doc__
+format_docstring(
+    model_name=XLMRobertaBackbone.__name__,
+    example_preset_name="xlm_roberta_base_multi",
+    preset_names=", ".join(XLMRobertaBackbone.presets),
+)(XLMRobertaBackbone.from_preset.__func__)
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index d78a0d34..97225c45 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -93,7 +93,7 @@ WHITESPACE_PUNCTUATION_AND_CJK_REGEX = r"|".join(
 
 def pretokenize(
     text,
-    lowercase=True,
+    lowercase=False,
     strip_accents=True,
     split=True,
     split_on_cjk=True,
diff --git a/tools/checkpoint_conversion/convert_deberta_v3_checkpoints.py b/tools/checkpoint_conversion/convert_deberta_v3_checkpoints.py
new file mode 100644
index 00000000..f818aa53
--- /dev/null
+++ b/tools/checkpoint_conversion/convert_deberta_v3_checkpoints.py
@@ -0,0 +1,326 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import os
+
+import numpy as np
+import requests
+import tensorflow as tf
+import transformers
+from absl import app
+from absl import flags
+
+from keras_nlp.models.deberta_v3.deberta_v3_backbone import DebertaV3Backbone
+from keras_nlp.models.deberta_v3.deberta_v3_preprocessor import (
+    DebertaV3Preprocessor,
+)
+from keras_nlp.models.deberta_v3.deberta_v3_tokenizer import DebertaV3Tokenizer
+from tools.checkpoint_conversion.checkpoint_conversion_utils import (
+    get_md5_checksum,
+)
+
+PRESET_MAP = {
+    "deberta_v3_extra_small_en": "microsoft/deberta-v3-xsmall",
+    "deberta_v3_small_en": "microsoft/deberta-v3-small",
+    "deberta_v3_base_en": "microsoft/deberta-v3-base",
+    "deberta_v3_large_en": "microsoft/deberta-v3-large",
+    "deberta_v3_base_multi": "microsoft/mdeberta-v3-base",
+}
+
+EXTRACT_DIR = "./{}"
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string(
+    "preset", None, f'Must be one of {",".join(PRESET_MAP.keys())}'
+)
+
+
+def download_files(preset, hf_model_name):
+    print("-> Download original vocabulary and config.")
+
+    extract_dir = EXTRACT_DIR.format(preset)
+    if not os.path.exists(extract_dir):
+        os.makedirs(extract_dir)
+
+    # Config.
+    config_path = os.path.join(extract_dir, "config.json")
+    response = requests.get(
+        f"https://huggingface.co/{hf_model_name}/raw/main/config.json"
+    )
+    open(config_path, "wb").write(response.content)
+    print(f"`{config_path}`")
+
+    # Vocab.
+    spm_path = os.path.join(extract_dir, "spm.model")
+    response = requests.get(
+        f"https://huggingface.co/{hf_model_name}/resolve/main/spm.model"
+    )
+    open(spm_path, "wb").write(response.content)
+    print(f"`{spm_path}`")
+
+
+def define_preprocessor(preset, hf_model_name):
+    print("\n-> Define the tokenizers.")
+    extract_dir = EXTRACT_DIR.format(preset)
+    spm_path = os.path.join(extract_dir, "spm.model")
+
+    keras_nlp_tokenizer = DebertaV3Tokenizer(proto=spm_path)
+
+    # Avoid having padding tokens. This is because the representations of the
+    # padding token may be vastly different from the representations computed in
+    # the original model. See https://github.com/keras-team/keras/pull/16619#issuecomment-1156338394.
+    sequence_length = 14
+    if preset == "deberta_v3_base_multi":
+        sequence_length = 17
+    keras_nlp_preprocessor = DebertaV3Preprocessor(
+        keras_nlp_tokenizer, sequence_length=sequence_length
+    )
+
+    hf_tokenizer = transformers.AutoTokenizer.from_pretrained(hf_model_name)
+
+    print("\n-> Print MD5 checksum of the vocab files.")
+    print(f"`{spm_path}` md5sum: ", get_md5_checksum(spm_path))
+
+    return keras_nlp_preprocessor, hf_tokenizer
+
+
+def convert_checkpoints(preset, keras_nlp_model, hf_model):
+    print("\n-> Convert original weights to KerasNLP format.")
+
+    extract_dir = EXTRACT_DIR.format(preset)
+    config_path = os.path.join(extract_dir, "config.json")
+
+    # Build config.
+    cfg = {}
+    with open(config_path, "r") as pt_cfg_handler:
+        pt_cfg = json.load(pt_cfg_handler)
+    cfg["vocabulary_size"] = pt_cfg["vocab_size"]
+    cfg["num_layers"] = pt_cfg["num_hidden_layers"]
+    cfg["num_heads"] = pt_cfg["num_attention_heads"]
+    cfg["hidden_dim"] = pt_cfg["hidden_size"]
+    cfg["intermediate_dim"] = pt_cfg["intermediate_size"]
+    cfg["dropout"] = pt_cfg["hidden_dropout_prob"]
+    cfg["max_sequence_length"] = pt_cfg["max_position_embeddings"]
+    cfg["bucket_size"] = pt_cfg["position_buckets"]
+    print("Config:", cfg)
+
+    hf_wts = hf_model.state_dict()
+    print("Original weights:")
+    print(
+        str(hf_wts.keys())
+        .replace(", ", "\n")
+        .replace("odict_keys([", "")
+        .replace("]", "")
+        .replace(")", "")
+    )
+
+    keras_nlp_model.get_layer("token_embedding").embeddings.assign(
+        hf_wts["embeddings.word_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("rel_embedding").rel_embeddings.assign(
+        hf_wts["encoder.rel_embeddings.weight"]
+    )
+    keras_nlp_model.get_layer("rel_embedding").layer_norm.gamma.assign(
+        hf_wts["encoder.LayerNorm.weight"]
+    )
+    keras_nlp_model.get_layer("rel_embedding").layer_norm.beta.assign(
+        hf_wts["encoder.LayerNorm.bias"]
+    )
+
+    for i in range(keras_nlp_model.num_layers):
+        # Q,K,V
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._self_attention_layer._query_dense.kernel.assign(
+            hf_wts[f"encoder.layer.{i}.attention.self.query_proj.weight"]
+            .numpy()
+            .T.reshape((cfg["hidden_dim"], cfg["num_heads"], -1))
+        )
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._self_attention_layer._query_dense.bias.assign(
+            hf_wts[f"encoder.layer.{i}.attention.self.query_proj.bias"]
+            .reshape((cfg["num_heads"], -1))
+            .numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._self_attention_layer._key_dense.kernel.assign(
+            hf_wts[f"encoder.layer.{i}.attention.self.key_proj.weight"]
+            .numpy()
+            .T.reshape((cfg["hidden_dim"], cfg["num_heads"], -1))
+        )
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._self_attention_layer._key_dense.bias.assign(
+            hf_wts[f"encoder.layer.{i}.attention.self.key_proj.bias"]
+            .reshape((cfg["num_heads"], -1))
+            .numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._self_attention_layer._value_dense.kernel.assign(
+            hf_wts[f"encoder.layer.{i}.attention.self.value_proj.weight"]
+            .numpy()
+            .T.reshape((cfg["hidden_dim"], cfg["num_heads"], -1))
+        )
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._self_attention_layer._value_dense.bias.assign(
+            hf_wts[f"encoder.layer.{i}.attention.self.value_proj.bias"]
+            .reshape((cfg["num_heads"], -1))
+            .numpy()
+        )
+
+        # Attn output.
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._self_attention_layer._output_dense.kernel.assign(
+            hf_wts[f"encoder.layer.{i}.attention.output.dense.weight"]
+            .transpose(1, 0)
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._self_attention_layer._output_dense.bias.assign(
+            hf_wts[f"encoder.layer.{i}.attention.output.dense.bias"].numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._self_attention_layernorm.gamma.assign(
+            hf_wts[
+                f"encoder.layer.{i}.attention.output.LayerNorm.weight"
+            ].numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._self_attention_layernorm.beta.assign(
+            hf_wts[f"encoder.layer.{i}.attention.output.LayerNorm.bias"].numpy()
+        )
+
+        # Intermediate FF layer.
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._feedforward_intermediate_dense.kernel.assign(
+            hf_wts[f"encoder.layer.{i}.intermediate.dense.weight"]
+            .transpose(1, 0)
+            .numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._feedforward_intermediate_dense.bias.assign(
+            hf_wts[f"encoder.layer.{i}.intermediate.dense.bias"].numpy()
+        )
+
+        # Output FF layer.
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._feedforward_output_dense.kernel.assign(
+            hf_wts[f"encoder.layer.{i}.output.dense.weight"].numpy().T
+        )
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._feedforward_output_dense.bias.assign(
+            hf_wts[f"encoder.layer.{i}.output.dense.bias"].numpy()
+        )
+
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._feedforward_layernorm.gamma.assign(
+            hf_wts[f"encoder.layer.{i}.output.LayerNorm.weight"].numpy()
+        )
+        keras_nlp_model.get_layer(
+            f"disentangled_attention_encoder_layer_{i}"
+        )._feedforward_layernorm.beta.assign(
+            hf_wts[f"encoder.layer.{i}.output.LayerNorm.bias"].numpy()
+        )
+
+    # Save the model.
+    print(f"\n-> Save KerasNLP model weights to `{preset}.h5`.")
+    keras_nlp_model.save_weights(f"{preset}.h5")
+
+    return keras_nlp_model
+
+
+def check_output(
+    preset,
+    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)
+
+    # HF
+    hf_inputs = hf_tokenizer(
+        sample_text, padding="longest", 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(f"./{preset}.h5"))
+
+
+def main(_):
+    hf_model_name = PRESET_MAP[FLAGS.preset]
+
+    download_files(FLAGS.preset, hf_model_name)
+
+    keras_nlp_preprocessor, hf_tokenizer = define_preprocessor(
+        FLAGS.preset, hf_model_name
+    )
+
+    print("\n-> Load KerasNLP model.")
+    keras_nlp_model = DebertaV3Backbone.from_preset(
+        FLAGS.preset, load_weights=False
+    )
+
+    print("\n-> Load HF model.")
+    hf_model = transformers.AutoModel.from_pretrained(hf_model_name)
+    hf_model.eval()
+
+    keras_nlp_model = convert_checkpoints(
+        FLAGS.preset, keras_nlp_model, hf_model
+    )
+
+    check_output(
+        FLAGS.preset,
+        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/transformer_decoder_test.py b/keras_nlp/layers/transformer_decoder_test.py
index aae0015c..52f1c1f1 100644
--- a/keras_nlp/layers/transformer_decoder_test.py
+++ b/keras_nlp/layers/transformer_decoder_test.py
@@ -132,7 +132,7 @@ class TransformerDecoderTest(tf.test.TestCase, parameterized.TestCase):
     def test_one_training_step_of_transformer_with_cross_attention(self):
         class MyModel(keras.Model):
             def __init__(self):
-                super(MyModel, self).__init__()
+                super().__init__()
                 self._decoder = transformer_decoder.TransformerDecoder(
                     intermediate_dim=4, num_heads=2
                 )
@@ -160,7 +160,7 @@ class TransformerDecoderTest(tf.test.TestCase, parameterized.TestCase):
     def test_one_training_step_of_transformer_without_cross_attention(self):
         class MyModel(keras.Model):
             def __init__(self):
-                super(MyModel, self).__init__()
+                super().__init__()
                 self._decoder = transformer_decoder.TransformerDecoder(
                     intermediate_dim=4,
                     num_heads=2,
diff --git a/keras_nlp/models/albert/albert_preprocessor_test.py b/keras_nlp/models/albert/albert_preprocessor_test.py
new file mode 100644
index 00000000..01267a6e
--- /dev/null
+++ b/keras_nlp/models/albert/albert_preprocessor_test.py
@@ -0,0 +1,174 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for ALBERT 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_preprocessor import AlbertPreprocessor
+from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
+
+
+class AlbertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        bytes_io = io.BytesIO()
+        vocab_data = tf.data.Dataset.from_tensor_slices(
+            ["the quick brown fox", "the earth is round"]
+        )
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=vocab_data.as_numpy_iterator(),
+            model_writer=bytes_io,
+            vocab_size=10,
+            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]",
+        )
+        self.proto = bytes_io.getvalue()
+
+        self.preprocessor = AlbertPreprocessor(
+            tokenizer=AlbertTokenizer(proto=self.proto),
+            sequence_length=12,
+        )
+
+    def test_tokenize_strings(self):
+        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]
+        )
+        self.assertAllEqual(
+            output["segment_ids"], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
+        )
+        self.assertAllEqual(
+            output["padding_mask"], [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]
+        )
+
+    def test_tokenize_list_of_strings(self):
+        # We should handle a list of strings as as batch.
+        input_data = ["the quick brown fox"] * 4
+        output = self.preprocessor(input_data)
+        self.assertAllEqual(
+            output["token_ids"],
+            [[2, 4, 9, 5, 7, 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
+        )
+        self.assertAllEqual(
+            output["padding_mask"], [[1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]] * 4
+        )
+
+    def test_tokenize_labeled_batch(self):
+        x = tf.constant(["the quick brown fox"] * 4)
+        y = tf.constant([1] * 4)
+        sw = tf.constant([1.0] * 4)
+        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,
+        )
+        self.assertAllEqual(
+            x_out["segment_ids"], [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(
+            x_out["padding_mask"], [[1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(y_out, y)
+        self.assertAllEqual(sw_out, sw)
+
+    def test_tokenize_labeled_dataset(self):
+        x = tf.constant(["the quick brown fox"] * 4)
+        y = tf.constant([1] * 4)
+        sw = tf.constant([1.0] * 4)
+        ds = tf.data.Dataset.from_tensor_slices((x, y, sw))
+        ds = ds.map(self.preprocessor)
+        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,
+        )
+        self.assertAllEqual(
+            x_out["segment_ids"], [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(
+            x_out["padding_mask"], [[1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(y_out, y)
+        self.assertAllEqual(sw_out, sw)
+
+    def test_tokenize_multiple_sentences(self):
+        sentence_one = tf.constant("the quick brown fox")
+        sentence_two = tf.constant("the earth")
+        output = self.preprocessor((sentence_one, sentence_two))
+        self.assertAllEqual(
+            output["token_ids"],
+            [2, 4, 9, 5, 7, 3, 4, 6, 3, 0, 0, 0],
+        )
+        self.assertAllEqual(
+            output["segment_ids"], [0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0]
+        )
+        self.assertAllEqual(
+            output["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0]
+        )
+
+    def test_tokenize_multiple_batched_sentences(self):
+        sentence_one = tf.constant(["the quick brown fox"] * 4)
+        sentence_two = tf.constant(["the earth"] * 4)
+        # The first tuple or list is always interpreted as an enumeration of
+        # separate sequences to concatenate.
+        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,
+        )
+        self.assertAllEqual(
+            output["segment_ids"], [[0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(
+            output["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0]] * 4
+        )
+
+    def test_errors_for_2d_list_input(self):
+        ambiguous_input = [["one", "two"], ["three", "four"]]
+        with self.assertRaises(ValueError):
+            self.preprocessor(ambiguous_input)
+
+    @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)
+        self.assertAllEqual(
+            model(input_data)["token_ids"],
+            restored_model(input_data)["token_ids"],
+        )
diff --git a/keras_nlp/models/albert/albert_tokenizer_test.py b/keras_nlp/models/albert/albert_tokenizer_test.py
new file mode 100644
index 00000000..0116f11e
--- /dev/null
+++ b/keras_nlp/models/albert/albert_tokenizer_test.py
@@ -0,0 +1,89 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for ALBERT tokenizer."""
+
+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_tokenizer import AlbertTokenizer
+
+
+class AlbertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        bytes_io = io.BytesIO()
+        vocab_data = tf.data.Dataset.from_tensor_slices(
+            ["the quick brown fox", "the earth is round"]
+        )
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=vocab_data.as_numpy_iterator(),
+            model_writer=bytes_io,
+            vocab_size=10,
+            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]",
+        )
+        self.proto = bytes_io.getvalue()
+
+        self.tokenizer = AlbertTokenizer(proto=self.proto)
+
+    def test_tokenize(self):
+        input_data = "the quick brown fox"
+        output = self.tokenizer(input_data)
+        self.assertAllEqual(output, [4, 9, 5, 7])
+
+    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]])
+
+    def test_detokenize(self):
+        input_data = tf.constant([[4, 9, 5, 7]])
+        output = self.tokenizer.detokenize(input_data)
+        self.assertEqual(output, tf.constant(["the quick brown fox"]))
+
+    def test_vocabulary_size(self):
+        tokenizer = AlbertTokenizer(proto=self.proto)
+        self.assertEqual(tokenizer.vocabulary_size(), 10)
+
+    @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.tokenizer(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)
+        self.assertAllEqual(
+            model(input_data),
+            restored_model(input_data),
+        )