Repository: keras-team/keras-nlp

Issue:
Port DistilBert to follow Bert model changes
We had a few larger changes to the bert API:
https://github.com/keras-team/keras-nlp/pull/387
https://github.com/keras-team/keras-nlp/pull/390

As well as some minor cleanups:
https://github.com/keras-team/keras-nlp/pull/397
https://github.com/keras-team/keras-nlp/pull/398
https://github.com/keras-team/keras-nlp/pull/399

We should mirror these where applicable to the DistilBert model code. Since we do not yet host pretrained weight for this model, we can just raise a `NotImplementedError` in `from_preset` for now.

Discussion:

Patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index cb77fead..da6035ec 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -15,22 +15,14 @@
 from keras_nlp.models.bert.bert_models import Bert
 from keras_nlp.models.bert.bert_preprocessing import BertPreprocessor
 from keras_nlp.models.bert.bert_tasks import BertClassifier
-from keras_nlp.models.distilbert.distilbert_models import DistilBertBase
-from keras_nlp.models.distilbert.distilbert_models import DistilBertCustom
+from keras_nlp.models.distilbert.distilbert_models import DistilBert
 from keras_nlp.models.distilbert.distilbert_preprocessing import (
     DistilBertPreprocessor,
 )
-from keras_nlp.models.gpt2.gpt2_models import Gpt2Base
-from keras_nlp.models.gpt2.gpt2_models import Gpt2Custom
-from keras_nlp.models.gpt2.gpt2_models import Gpt2ExtraLarge
-from keras_nlp.models.gpt2.gpt2_models import Gpt2Large
-from keras_nlp.models.gpt2.gpt2_models import Gpt2Medium
-from keras_nlp.models.roberta.roberta_models import RobertaBase
-from keras_nlp.models.roberta.roberta_models import RobertaCustom
+from keras_nlp.models.gpt2.gpt2_models import Gpt2
+from keras_nlp.models.roberta.roberta_models import Roberta
 from keras_nlp.models.roberta.roberta_tasks import RobertaClassifier
-from keras_nlp.models.xlm_roberta.xlm_roberta_models import XLMRobertaBase
-from keras_nlp.models.xlm_roberta.xlm_roberta_models import XLMRobertaCustom
-from keras_nlp.models.xlm_roberta.xlm_roberta_models import XLMRobertaLarge
+from keras_nlp.models.xlm_roberta.xlm_roberta_models import XLMRoberta
 from keras_nlp.models.xlm_roberta.xlm_roberta_preprocessing import (
     XLMRobertaPreprocessor,
 )
diff --git a/keras_nlp/models/bert/bert_preprocessing.py b/keras_nlp/models/bert/bert_preprocessing.py
index 5c06977f..c5abdf6f 100644
--- a/keras_nlp/models/bert/bert_preprocessing.py
+++ b/keras_nlp/models/bert/bert_preprocessing.py
@@ -24,6 +24,7 @@ from keras_nlp.models.utils import classproperty
 from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
 
 
+@keras.utils.register_keras_serializable(package="keras_nlp")
 class BertPreprocessor(keras.layers.Layer):
     """BERT preprocessing layer.
 
@@ -62,7 +63,8 @@ class BertPreprocessor(keras.layers.Layer):
 
     Examples:
     ```python
-    vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "tripped"]
+    vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]"]
+    vocab += ["the", "qu", "##ick", "br", "##own", "fox", "tripped"]
     vocab += ["call", "me", "ish", "##mael", "."]
     vocab += ["oh", "look", "a", "whale"]
     vocab += ["i", "forgot", "my", "home", "##work"]
diff --git a/keras_nlp/models/distilbert/distilbert_models.py b/keras_nlp/models/distilbert/distilbert_models.py
index 72d54b14..7e5cdec2 100644
--- a/keras_nlp/models/distilbert/distilbert_models.py
+++ b/keras_nlp/models/distilbert/distilbert_models.py
@@ -29,8 +29,8 @@ def distilbert_kernel_initializer(stddev=0.02):
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class DistilBertCustom(keras.Model):
-    """DistilBERT encoder network with custom hyperparmeters.
+class DistilBert(keras.Model):
+    """DistilBERT encoder network.
 
     This network implements a bi-directional Transformer-based encoder as
     described in ["DistilBERT, a distilled version of BERT: smaller, faster,
@@ -38,9 +38,10 @@ class DistilBertCustom(keras.Model):
     embedding lookups and transformer layers, but not the masked language model
     or classification task networks.
 
-    This class gives a fully customizable DistilBERT model with any number of layers,
-    heads, and embedding dimensions. For specific DistilBERT architectures defined in
-    the paper, see, for example, `keras_nlp.models.DistilBertBase`.
+    The default constructor gives a fully customizable, randomly initalized
+    DistilBERT encoder with any number of layers, heads, and embedding
+    dimensions. To load preset architectures and weights, use the `from_presets`
+    constructor.
 
     Args:
         vocabulary_size: int. The size of the token vocabulary.
@@ -55,32 +56,27 @@ class DistilBertCustom(keras.Model):
             can consume. If None, `max_sequence_length` uses the value from
             sequence length. This determines the variable shape for positional
             embeddings.
-        name: string, optional. Name of the model.
-        trainable: boolean, optional. If the model's variables should be
-            trainable.
 
     Examples:
     ```python
-    # Randomly initialized DistilBERT encoder
-    model = keras_nlp.models.DistilBertCustom(
-        vocabulary_size=30522,
-        num_layers=12,
-        num_heads=12,
-        hidden_dim=768,
-        intermediate_dim=3072,
-        max_sequence_length=12,
-        name="encoder",
-    )
-
-    # Call encoder on the inputs
     input_data = {
         "token_ids": tf.random.uniform(
-            shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
+            shape=(1, 12), dtype=tf.int64, maxval=30522
         ),
         "padding_mask": tf.constant(
             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
         ),
     }
+
+    # Randomly initialized DistilBERT encoder
+    model = keras_nlp.models.DistilBert(
+        vocabulary_size=30522,
+        num_layers=6,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=12,
+    )
     output = model(input_data)
     ```
     """
@@ -94,8 +90,7 @@ class DistilBertCustom(keras.Model):
         intermediate_dim,
         dropout=0.1,
         max_sequence_length=512,
-        name=None,
-        trainable=True,
+        **kwargs,
     ):
 
         # Inputs
@@ -138,6 +133,10 @@ class DistilBertCustom(keras.Model):
                 name=f"transformer_layer_{i}",
             )(x, padding_mask=padding_mask)
 
+        # Set default for `name` if none given
+        if "name" not in kwargs:
+            kwargs["name"] = "backbone"
+
         # Instantiate using Functional API Model constructor
         super().__init__(
             inputs={
@@ -145,8 +144,7 @@ class DistilBertCustom(keras.Model):
                 "padding_mask": padding_mask,
             },
             outputs=x,
-            name=name,
-            trainable=trainable,
+            **kwargs,
         )
         # All references to `self` below this line
         self.vocabulary_size = vocabulary_size
@@ -174,64 +172,11 @@ class DistilBertCustom(keras.Model):
     def from_config(cls, config):
         return cls(**config)
 
-
-MODEL_DOCSTRING = """DistilBert "{type}" architecture.
-
-This network implements a bi-directional Transformer-based encoder as
-described in ["DistilBERT, a distilled version of BERT: smaller, faster,
-cheaper and lighter"](https://arxiv.org/abs/1910.01108). It includes the
-embedding lookups and transformer layers, but not the masked language model
-or classification task networks.
-
-Args:
-    vocabulary_size: int, optional. The size of the token vocabulary.
-    name: string, optional. Name of the model.
-    trainable: boolean, optional. If the model's variables should be
-        trainable.
-
-Examples:
-```python
-# Randomly initialized DistilBert{type} encoder
-model = keras_nlp.models.DistilBert{type}(vocabulary_size=10000)
-
-# Call encoder on the inputs.
-input_data = {{
-    "token_ids": tf.random.uniform(
-        shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
-    ),
-    "padding_mask": tf.constant([1] * 512, shape=(1, 512)),
-}}
-output = model(input_data)
-```
-"""
-
-
-def DistilBertBase(vocabulary_size, name=None, trainable=True):
-
-    model = DistilBertCustom(
-        vocabulary_size=vocabulary_size,
-        num_layers=6,
-        num_heads=12,
-        hidden_dim=768,
-        intermediate_dim=3072,
-        dropout=0.1,
-        max_sequence_length=512,
-        name=name,
-        trainable=trainable,
-    )
-
-    return model
-
-
-def model_class_by_name(classname):
-    """Return model class given the class name."""
-    return {
-        "DistilBertBase": DistilBertBase,
-    }[classname]
-
-
-setattr(
-    DistilBertBase,
-    "__doc__",
-    MODEL_DOCSTRING.format(type="Base"),
-)
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        load_weights=True,
+        **kwargs,
+    ):
+        raise NotImplementedError
diff --git a/keras_nlp/models/distilbert/distilbert_preprocessing.py b/keras_nlp/models/distilbert/distilbert_preprocessing.py
index aa8ed961..39045f39 100644
--- a/keras_nlp/models/distilbert/distilbert_preprocessing.py
+++ b/keras_nlp/models/distilbert/distilbert_preprocessing.py
@@ -18,78 +18,80 @@ from tensorflow import keras
 from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
 from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
 
-PREPROCESSOR_DOCSTRING = """DistilBERT preprocessor with pretrained vocabularies.
-
-This preprocessing layer will do three things:
-
- - Tokenize any number of inputs using a
-   `keras_nlp.tokenizers.WordPieceTokenizer`.
- - Pack the inputs together using a `keras_nlp.layers.MultiSegmentPacker`.
-   with the appropriate `"[CLS]"`, `"[SEP]"` and `"[PAD]"` tokens.
- - Construct a dictionary of with keys `"token_ids"`, `"padding_mask"`, that can
-   be passed directly to a DistilBERT model.
-
-This layer will accept either a tuple of (possibly batched) inputs, or a single
-input tensor. If a single tensor is passed, it will be packed equivalently to
-a tuple with a single element.
-
-The WordPiece tokenizer can be accessed via the `tokenizer` property on this
-layer, and can be used directly for custom packing on inputs.
-
-Args:
-    vocabulary: One of a list of vocabulary terms, a vocabulary filename, or
-        the name of the pretrained vocabulary.
-    lowercase: If `True`, input will be lowercase before tokenization. If
-        `vocabulary` is set to a pretrained vocabulary, this parameter will
-        be inferred.
-    sequence_length: The length of the packed inputs. Only used if
-        `pack_inputs` is True.
-    truncate: The algorithm to truncate a list of batched segments to fit
-        within `sequence_length`. Only used if
-        `pack_inputs` is True. 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
-preprocessor = keras_nlp.models.DistilBertPreprocessor(vocabulary="uncased_en")
-
-# Tokenize and pack a single sentence directly.
-preprocessor("The quick brown fox jumped.")
-
-# Tokenize and pack a multiple sentence directly.
-preprocessor(("The quick brown fox jumped.", "Call me Ishmael."))
-
-# Map a dataset to preprocess a single sentence.
-features = ["The quick brown fox jumped.", "I forgot my homework."]
-labels = [0, 1]
-ds = tf.data.Dataset.from_tensor_slices((features, labels))
-ds = ds.map(
-    lambda x, y: (preprocessor(x), y),
-    num_parallel_calls=tf.data.AUTOTUNE,
-)
-
-# Map a dataset to preprocess a multiple sentences.
-first_sentences = ["The quick brown fox jumped.", "Call me Ishmael."]
-second_sentences = ["The fox tripped.", "Oh look, a whale."]
-labels = [1, 1]
-ds = tf.data.Dataset.from_tensor_slices(((first_sentences, second_sentences), labels))
-ds = ds.map(
-    lambda x, y: (preprocessor(x), y),
-    num_parallel_calls=tf.data.AUTOTUNE,
-)
-```
-"""
-
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
 class DistilBertPreprocessor(keras.layers.Layer):
+    """DistilBERT preprocessing layer.
+
+    This preprocessing layer will do three things:
+
+     - Tokenize any number of inputs using a
+       `keras_nlp.tokenizers.WordPieceTokenizer`.
+     - Pack the inputs together using a `keras_nlp.layers.MultiSegmentPacker`.
+       with the appropriate `"[CLS]"`, `"[SEP]"` and `"[PAD]"` tokens.
+     - Construct a dictionary of with keys `"token_ids"`, `"padding_mask"`,
+       that can be passed directly to a DistilBERT model.
+
+    This layer will accept either a tuple of (possibly batched) inputs, or a
+    single input tensor. If a single tensor is passed, it will be packed
+    equivalently to a tuple with a single element.
+
+    The WordPiece tokenizer can be accessed via the `tokenizer` property on this
+    layer, and can be used directly for custom packing on inputs.
+
+    Args:
+        vocabulary: A list of vocabulary terms or a vocabulary filename.
+        lowercase: If `True`, input will be lowercase before tokenization. If
+            `vocabulary` is set to a pretrained vocabulary, this parameter will
+            be inferred.
+        sequence_length: The length of the packed inputs. Only used if
+            `pack_inputs` is True.
+        truncate: The algorithm to truncate a list of batched segments to fit
+            within `sequence_length`. Only used if
+            `pack_inputs` is True. 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
+    preprocessor = keras_nlp.models.DistilBertPreprocessor(
+        vocabulary="./vocab.txt",
+    )
+
+    # Tokenize and pack a single sentence directly.
+    preprocessor("The quick brown fox jumped.")
+
+    # Tokenize and pack a multiple sentence directly.
+    preprocessor(("The quick brown fox jumped.", "Call me Ishmael."))
+
+    # Map a dataset to preprocess a single sentence.
+    features = ["The quick brown fox jumped.", "I forgot my homework."]
+    labels = [0, 1]
+    ds = tf.data.Dataset.from_tensor_slices((features, labels))
+    ds = ds.map(
+        lambda x, y: (preprocessor(x), y),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+
+    # Map a dataset to preprocess a multiple sentences.
+    first_sentences = ["The quick brown fox jumped.", "Call me Ishmael."]
+    second_sentences = ["The fox tripped.", "Oh look, a whale."]
+    labels = [1, 1]
+    ds = tf.data.Dataset.from_tensor_slices(
+        ((first_sentences, second_sentences), labels))
+    ds = ds.map(
+        lambda x, y: (preprocessor(x), y),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+    ```
+    """
+
     def __init__(
         self,
         vocabulary="uncased_en",
@@ -153,9 +155,12 @@ class DistilBertPreprocessor(keras.layers.Layer):
             "padding_mask": token_ids != self.pad_token_id,
         }
 
-
-setattr(
-    DistilBertPreprocessor,
-    "__doc__",
-    PREPROCESSOR_DOCSTRING,
-)
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        sequence_length=None,
+        truncate="round_robin",
+        **kwargs,
+    ):
+        raise NotImplementedError
diff --git a/keras_nlp/models/gpt2/gpt2_models.py b/keras_nlp/models/gpt2/gpt2_models.py
index 04e4a07b..14f0937f 100644
--- a/keras_nlp/models/gpt2/gpt2_models.py
+++ b/keras_nlp/models/gpt2/gpt2_models.py
@@ -26,17 +26,18 @@ def _gpt_2_kernel_initializer(stddev=0.02):
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class Gpt2Custom(keras.Model):
-    """GPT-2 core network with customizable hyperparameters.
+class Gpt2(keras.Model):
+    """GPT-2 core network with hyperparameters.
 
     This network implements a Transformer-based decoder network,
     Generative Pretrained Transformer-2 (GPT-2), as described in
     ["Language Models are Unsupervised Multitask Learners"](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf).
     It includes the embedding lookups and transformer layers.
 
-    This class gives a fully customizable GPT-2 model with any number of layers,
-    heads, and embedding dimensions. For specific GPT-2 architectures
-    defined in the paper, see, for example, `keras_nlp.models.Gpt2Base`.
+    The default constructor gives a fully customizable, randomly initalized
+    GPT-2 model with any number of layers, heads, and embedding
+    dimensions. To load preset architectures and weights, use the `from_presets`
+    constructor.
 
     Args:
         vocabulary_size: int. The size of the token vocabulary.
@@ -51,32 +52,29 @@ class Gpt2Custom(keras.Model):
             can consume. If None, `max_sequence_length` uses the value from
             sequence length. This determines the variable shape for positional
             embeddings.
-        name: string, optional. Name of the model.
-        trainable: boolean, optional. If the model's variables should be
-            trainable.
 
     Example usage:
     ```python
+    input_data = {
+        "token_ids": tf.random.uniform(
+            shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
+        ),
+        "padding_mask": tf.constant(
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
+    }
+
     # Randomly initialized GPT-2 decoder
-    model = keras_nlp.models.Gpt2Custom(
+    model = keras_nlp.models.Gpt2(
         vocabulary_size=50257,
         num_layers=12,
         num_heads=12,
         hidden_dim=768,
         intermediate_dim=3072,
         max_sequence_length=1024,
-        name="encoder",
     )
 
-    # Call encoder on the inputs
-    input_data = {
-        "token_ids": tf.random.uniform(
-            shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
-        ),
-        "padding_mask": tf.constant(
-            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-        ),
-    }
+    # Call the model on the input data.
     output = model(input_data)
     ```
     """
@@ -90,8 +88,7 @@ class Gpt2Custom(keras.Model):
         intermediate_dim,
         dropout=0.1,
         max_sequence_length=1024,
-        name=None,
-        trainable=True,
+        **kwargs,
     ):
 
         # Inputs
@@ -145,6 +142,10 @@ class Gpt2Custom(keras.Model):
             dtype=tf.float32,
         )(x)
 
+        # Set default for `name` if none given
+        if "name" not in kwargs:
+            kwargs["name"] = "backbone"
+
         # Instantiate using Functional API Model constructor
         super().__init__(
             inputs={
@@ -152,8 +153,7 @@ class Gpt2Custom(keras.Model):
                 "padding_mask": padding_mask,
             },
             outputs=sequence_output,
-            name=name,
-            trainable=trainable,
+            **kwargs,
         )
         # All references to `self` below this line
         self.vocabulary_size = vocabulary_size
@@ -181,109 +181,11 @@ class Gpt2Custom(keras.Model):
     def from_config(cls, config):
         return cls(**config)
 
-
-MODEL_DOCSTRING = """GPT-2 "{type}" architecture.
-
-    This network implements a Transformer-based decoder as
-    described in
-    ["Language Models are Unsupervised Multitask Learners"](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf).
-    It includes the embedding lookups and transformer layers.
-
-    Args:
-        vocabulary_size: int, optional. The size of the token vocabulary.
-        name: String, optional. Name of the model.
-        trainable: boolean, optional. If the model's variables should be
-            trainable.
-
-    Example usage:
-    ```python
-    # Randomly initialized Gpt2{type} encoder
-    model = keras_nlp.models.Gpt2{type}(vocabulary_size=10000)
-
-    # Call encoder on the inputs.
-    input_data = {{
-        "token_ids": tf.random.uniform(
-            shape=(1, 1024), dtype=tf.int64, maxval=model.vocabulary_size
-        ),
-        "padding_mask": tf.constant([1] * 1024, shape=(1, 1024)),
-    }}
-    output = model(input_data)
-"""
-
-
-def Gpt2Base(vocabulary_size, name=None, trainable=True):
-    return Gpt2Custom(
-        vocabulary_size=vocabulary_size,
-        num_layers=12,
-        num_heads=12,
-        hidden_dim=768,
-        intermediate_dim=3072,
-        dropout=0.1,
-        max_sequence_length=1024,
-        name=name,
-        trainable=trainable,
-    )
-
-
-def Gpt2Medium(vocabulary_size, name=None, trainable=True):
-    return Gpt2Custom(
-        vocabulary_size=vocabulary_size,
-        num_layers=24,
-        num_heads=16,
-        hidden_dim=1024,
-        intermediate_dim=4096,
-        dropout=0.1,
-        max_sequence_length=1024,
-        name=name,
-        trainable=trainable,
-    )
-
-
-def Gpt2Large(vocabulary_size, name=None, trainable=True):
-    return Gpt2Custom(
-        vocabulary_size=vocabulary_size,
-        num_layers=36,
-        num_heads=20,
-        hidden_dim=1280,
-        intermediate_dim=5120,
-        dropout=0.1,
-        max_sequence_length=1024,
-        name=name,
-        trainable=trainable,
-    )
-
-
-def Gpt2ExtraLarge(vocabulary_size, name=None, trainable=True):
-    return Gpt2Custom(
-        vocabulary_size=vocabulary_size,
-        num_layers=48,
-        num_heads=25,
-        hidden_dim=1600,
-        intermediate_dim=6400,
-        dropout=0.1,
-        max_sequence_length=1024,
-        name=name,
-        trainable=trainable,
-    )
-
-
-setattr(
-    Gpt2Base,
-    "__doc__",
-    MODEL_DOCSTRING.format(type="Base"),
-)
-setattr(
-    Gpt2Medium,
-    "__doc__",
-    MODEL_DOCSTRING.format(type="Medium"),
-)
-setattr(
-    Gpt2Large,
-    "__doc__",
-    MODEL_DOCSTRING.format(type="Large"),
-)
-setattr(
-    Gpt2ExtraLarge,
-    "__doc__",
-    MODEL_DOCSTRING.format(type="ExtraLarge"),
-)
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        load_weights=True,
+        **kwargs,
+    ):
+        raise NotImplementedError
diff --git a/keras_nlp/models/roberta/roberta_models.py b/keras_nlp/models/roberta/roberta_models.py
index 6520c516..bebb425a 100644
--- a/keras_nlp/models/roberta/roberta_models.py
+++ b/keras_nlp/models/roberta/roberta_models.py
@@ -26,17 +26,18 @@ def roberta_kernel_initializer(stddev=0.02):
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class RobertaCustom(keras.Model):
-    """RoBERTa encoder with a customizable set of hyperparameters.
+class Roberta(keras.Model):
+    """RoBERTa encoder.
 
     This network implements a bi-directional Transformer-based encoder as
     described in ["RoBERTa: A Robustly Optimized BERT Pretraining Approach"](https://arxiv.org/abs/1907.11692).
     It includes the embedding lookups and transformer layers, but does not
     include the masked language modeling head used during pretraining.
 
-    This class gives a fully configurable RoBERTa model with any number of
-    layers, heads, and embedding dimensions. For specific RoBERTa architectures
-    defined in the paper, see, for example, `keras_nlp.models.RobertaBase`.
+    The default constructor gives a fully customizable, randomly initalized
+    RoBERTa encoder with any number of layers, heads, and embedding
+    dimensions. To load preset architectures and weights, use the `from_presets`
+    constructor.
 
     Args:
         vocabulary_size: int. The size of the token vocabulary.
@@ -50,14 +51,18 @@ class RobertaCustom(keras.Model):
         max_sequence_length: int. The maximum sequence length this encoder can
             consume. The sequence length of the input must be less than
             `max_sequence_length`.
-        name: string, optional. Name of the model.
-        trainable: boolean, optional. If the model's variables should be
-            trainable.
 
     Example usage:
     ```python
+    input_data = {
+        "token_ids": tf.random.uniform(
+            shape=(1, 12), dtype=tf.int64, maxval=50265),
+        "padding_mask": tf.constant(
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+    }
+
     # Randomly initialized RoBERTa model
-    model = keras_nlp.models.RobertaCustom(
+    model = keras_nlp.models.Roberta(
         vocabulary_size=50265,
         num_layers=12,
         num_heads=12,
@@ -66,13 +71,7 @@ class RobertaCustom(keras.Model):
         max_sequence_length=12
     )
 
-    # Call encoder on the inputs.
-    input_data = {
-        "token_ids": tf.random.uniform(
-            shape=(1, 12), dtype=tf.int64, maxval=50265),
-        "padding_mask": tf.constant(
-            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
-    }
+    # Call the model on the input data.
     output = model(input_data)
     ```
     """
@@ -86,12 +85,9 @@ class RobertaCustom(keras.Model):
         intermediate_dim,
         dropout=0.1,
         max_sequence_length=512,
-        name=None,
-        trainable=True,
+        **kwargs,
     ):
 
-        # Index of classification token in the vocabulary
-        cls_token_index = 0
         # Inputs
         token_id_input = keras.Input(
             shape=(None,), dtype=tf.int32, name="token_ids"
@@ -133,6 +129,10 @@ class RobertaCustom(keras.Model):
                 name=f"transformer_layer_{i}",
             )(x, padding_mask=padding_mask)
 
+        # Set default for `name` if none given
+        if "name" not in kwargs:
+            kwargs["name"] = "backbone"
+
         # Instantiate using Functional API Model constructor
         super().__init__(
             inputs={
@@ -140,29 +140,27 @@ class RobertaCustom(keras.Model):
                 "padding_mask": padding_mask,
             },
             outputs=x,
-            name=name,
-            trainable=trainable,
+            **kwargs,
         )
         # All references to `self` below this line
         self.vocabulary_size = vocabulary_size
-        self.hidden_dim = hidden_dim
-        self.intermediate_dim = intermediate_dim
         self.num_layers = num_layers
         self.num_heads = num_heads
-        self.max_sequence_length = max_sequence_length
+        self.hidden_dim = hidden_dim
         self.intermediate_dim = intermediate_dim
         self.dropout = dropout
-        self.cls_token_index = cls_token_index
+        self.max_sequence_length = max_sequence_length
+        self.start_token_index = 0
 
     def get_config(self):
         return {
             "vocabulary_size": self.vocabulary_size,
-            "hidden_dim": self.hidden_dim,
-            "intermediate_dim": self.intermediate_dim,
             "num_layers": self.num_layers,
             "num_heads": self.num_heads,
-            "max_sequence_length": self.max_sequence_length,
+            "hidden_dim": self.hidden_dim,
+            "intermediate_dim": self.intermediate_dim,
             "dropout": self.dropout,
+            "max_sequence_length": self.max_sequence_length,
             "name": self.name,
             "trainable": self.trainable,
         }
@@ -171,45 +169,11 @@ class RobertaCustom(keras.Model):
     def from_config(cls, config):
         return cls(**config)
 
-
-def RobertaBase(vocabulary_size, name=None, trainable=True):
-    """RoBERTa implementation using "Base" architecture.
-
-    This network implements a bi-directional Transformer-based encoder as
-    described in ["RoBERTa: A Robustly Optimized BERT Pretraining
-    Approach"](https://arxiv.org/abs/1907.11692). It includes the
-    embedding lookups and transformer layers, but does not include the masked
-    language modeling head used during pretraining.
-
-    Args:
-        vocabulary_size: int, optional. The size of the token vocabulary.
-        name: string, optional. Name of the model.
-        trainable: boolean, optional. If the model's variables should be
-            trainable.
-
-    Example usage:
-    ```python
-    # Randomly initialized RobertaBase encoder
-    model = keras_nlp.models.RobertaBase(vocabulary_size=10000)
-
-    # Call encoder on the inputs.
-    input_data = {
-        "token_ids": tf.random.uniform(
-            shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size),
-        "padding_mask": tf.ones((1, 512)),
-    }
-    output = model(input_data)
-    ```
-    """
-
-    return RobertaCustom(
-        vocabulary_size=vocabulary_size,
-        num_layers=12,
-        num_heads=12,
-        hidden_dim=768,
-        intermediate_dim=3072,
-        dropout=0.1,
-        max_sequence_length=512,
-        name=name,
-        trainable=trainable,
-    )
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        load_weights=True,
+        **kwargs,
+    ):
+        raise NotImplementedError
diff --git a/keras_nlp/models/roberta/roberta_tasks.py b/keras_nlp/models/roberta/roberta_tasks.py
index 88cd0422..977ad621 100644
--- a/keras_nlp/models/roberta/roberta_tasks.py
+++ b/keras_nlp/models/roberta/roberta_tasks.py
@@ -23,17 +23,14 @@ class RobertaClassifier(keras.Model):
     """RoBERTa encoder model with a classification head.
 
     Args:
-        backbone: A `keras_nlp.models.Roberta` to encode inputs.
+        backbone: A `keras_nlp.models.Roberta` instance.
         num_classes: int. Number of classes to predict.
         hidden_dim: int. The size of the pooler layer.
-        name: string, optional. Name of the model.
-        trainable: boolean, optional. If the model's variables should be
-            trainable.
 
     Example usage:
     ```python
     # Randomly initialized RoBERTa encoder
-    model = keras_nlp.models.RobertaCustom(
+    model = keras_nlp.models.Roberta(
         vocabulary_size=50265,
         num_layers=12,
         num_heads=12,
@@ -60,14 +57,13 @@ class RobertaClassifier(keras.Model):
         num_classes,
         hidden_dim=None,
         dropout=0.0,
-        name=None,
-        trainable=True,
+        **kwargs,
     ):
         inputs = backbone.input
         if hidden_dim is None:
             hidden_dim = backbone.hidden_dim
 
-        x = backbone(inputs)[:, backbone.cls_token_index, :]
+        x = backbone(inputs)[:, backbone.start_token_index, :]
         x = keras.layers.Dropout(dropout, name="pooled_dropout")(x)
         x = keras.layers.Dense(
             hidden_dim, activation="tanh", name="pooled_dense"
@@ -81,14 +77,21 @@ class RobertaClassifier(keras.Model):
 
         # Instantiate using Functional API Model constructor
         super().__init__(
-            inputs=inputs, outputs=outputs, name=name, trainable=trainable
+            inputs=inputs,
+            outputs=outputs,
+            **kwargs,
         )
         # All references to `self` below this line
-        self.backbone = backbone
+        self._backbone = backbone
         self.num_classes = num_classes
         self.hidden_dim = hidden_dim
         self.dropout = dropout
 
+    @property
+    def backbone(self):
+        """A `keras_nlp.models.Roberta` instance providing the encoder submodel."""
+        return self._backbone
+
     def get_config(self):
         return {
             "backbone": keras.layers.serialize(self.backbone),
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_models.py b/keras_nlp/models/xlm_roberta/xlm_roberta_models.py
index 67fb0e00..02f293f3 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_models.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_models.py
@@ -21,8 +21,8 @@ from keras_nlp.models.roberta import roberta_models
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class XLMRobertaCustom(roberta_models.RobertaCustom):
-    """XLM-RoBERTa encoder with a customizable set of hyperparameters.
+class XLMRoberta(roberta_models.Roberta):
+    """XLM-RoBERTa encoder.
 
     This network implements a bi-directional Transformer-based encoder as
     described in
@@ -30,10 +30,10 @@ class XLMRobertaCustom(roberta_models.RobertaCustom):
     It includes the embedding lookups and transformer layers, but does not
     include the masked language modeling head used during pretraining.
 
-    This class gives a fully configurable XLM-R model with any number of
-    layers, heads, and embedding dimensions. The graph of XLM-R is
-    exactly the same as RoBERTa's. For specific XLM-R architectures
-    defined in the paper, see, for example, `keras_nlp.models.XLMRobertaBase`.
+    The default constructor gives a fully customizable, randomly initalized
+    RoBERTa encoder with any number of layers, heads, and embedding
+    dimensions. To load preset architectures and weights, use the `from_presets`
+    constructor.
 
     Args:
         vocabulary_size: int. The size of the token vocabulary.
@@ -47,120 +47,36 @@ class XLMRobertaCustom(roberta_models.RobertaCustom):
         max_sequence_length: int. The maximum sequence length this encoder can
             consume. The sequence length of the input must be less than
             `max_sequence_length`.
-        name: string, optional. Name of the model.
-        trainable: boolean, optional. If the model's variables should be
-            trainable.
 
     Example usage:
     ```python
-    # Randomly initialized XLM-R model
-    model = keras_nlp.models.XLMRobertaCustom(
-        vocabulary_size=50265,
-        num_layers=12,
-        num_heads=12,
-        hidden_dim=768,
-        intermediate_dim=3072,
-        max_sequence_length=12
-    )
-
-    # Call encoder on the inputs.
     input_data = {
         "token_ids": tf.random.uniform(
             shape=(1, 12), dtype=tf.int64, maxval=50265),
         "padding_mask": tf.constant(
             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
     }
-    output = model(input_data)
-    ```
-    """
-
-    pass
-
-
-MODEL_DOCSTRING = """XLM-RoBERTa "{type}" architecture.
-
-This network implements a bi-directional Transformer-based encoder as
-described in
-["Unsupervised Cross-lingual Representation Learning at Scale"](https://arxiv.org/abs/1911.02116).
-It includes the embedding lookups and transformer layers, but does not
-include the masked language modeling head used during pretraining.
-
-Args:
-    vocabulary_size: int, optional. The size of the token vocabulary.
-    name: string, optional. Name of the model.
-    trainable: boolean, optional. If the model's variables should be
-        trainable.
-
-Example usage:
-```python
-# Randomly initialized XLMRoberta{type} encoder
-model = keras_nlp.models.XLMRoberta{type}(weights=None, vocabulary_size=10000)
 
-# Call encoder on the inputs.
-input_data = {{
-    "token_ids": tf.random.uniform(
-        shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size),
-    "padding_mask": tf.ones((1, 512)),
-}}
-output = model(input_data)
-```
-"""
-
-
-def XLMRobertaBase(vocabulary_size, name=None, trainable=True):
-
-    model = XLMRobertaCustom(
-        vocabulary_size=vocabulary_size,
+    # Randomly initialized XLM-R model
+    model = keras_nlp.models.XLMRoberta(
+        vocabulary_size=50265,
         num_layers=12,
         num_heads=12,
         hidden_dim=768,
         intermediate_dim=3072,
-        dropout=0.1,
-        max_sequence_length=512,
-        name=name,
-        trainable=trainable,
-    )
-
-    return model
-
-
-def XLMRobertaLarge(
-    vocabulary_size,
-    name=None,
-    trainable=True,
-):
-
-    model = XLMRobertaCustom(
-        vocabulary_size=vocabulary_size,
-        num_layers=24,
-        num_heads=16,
-        hidden_dim=1024,
-        intermediate_dim=4096,
-        dropout=0.1,
-        max_sequence_length=512,
-        name=name,
-        trainable=trainable,
+        max_sequence_length=12
     )
 
-    return model
-
-
-def model_class_by_name(classname):
-    """Return model class given the class name."""
-    return {
-        "XLMRobertaBase": XLMRobertaBase,
-        "XLMRobertaLarge": XLMRobertaLarge,
-    }[classname]
-
-
-setattr(
-    XLMRobertaBase,
-    "__doc__",
-    MODEL_DOCSTRING.format(type="Base"),
-)
+    # Call the model on the input data.
+    output = model(input_data)
+    ```
+    """
 
-setattr(
-    XLMRobertaLarge,
-    "__doc__",
-    MODEL_DOCSTRING.format(type="Large"),
-)
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        load_weights=True,
+        **kwargs,
+    ):
+        raise NotImplementedError
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessing.py b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessing.py
index f8683461..ddf28252 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessing.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessing.py
@@ -22,83 +22,81 @@ from keras_nlp.models.roberta.roberta_preprocessing import (
 )
 from keras_nlp.tokenizers.sentence_piece_tokenizer import SentencePieceTokenizer
 
-PREPROCESSOR_DOCSTRING = """XLM-RoBERTa preprocessor with pretrained vocabularies.
-
-This preprocessing layer will do three things:
-
- - Tokenize any number of inputs using a
-   `keras_nlp.tokenizers.SentencePieceTokenizer`.
- - Pack the inputs together with the appropriate `"
"`, `"
"` and `"<pad>"`
-   tokens, i.e., adding a single `"
"` at the start of the entire sequence,
-   `"

"` at the end of each segment, save the last and a `"
"` at the
-   end of the entire sequence.
- - Construct a dictionary of with keys `"token_ids"`, `"padding_mask"`, that can
-   be passed directly to a XLM-RoBERTa model.
-
-Note that the original fairseq implementation modifies the indices of the
-SentencePiece tokenizer output. To preserve compatibility, we make the same
-changes, i.e., `"
"`, `"<pad>"`, `"
"` and `"<unk>"` are mapped to
-0, 1, 2, 3, respectively, and non-special tokens' indices are shifted right by
-one. Keep this in mind if generating your own vocabulary for tokenization.
-
-This layer will accept either a tuple of (possibly batched) inputs, or a single
-input tensor. If a single tensor is passed, it will be packed equivalently to
-a tuple with a single element.
-
-The SentencePiece tokenizer can be accessed via the `tokenizer` property on this
-layer, and can be used directly for custom packing on inputs.
-
-Args:
-    proto: Either a `string` path to a SentencePiece proto file, a
-        `bytes` object with a serialized SentencePiece proto.
-    sequence_length: The length of the packed inputs. Only used if
-        `pack_inputs` is True.
-    truncate: The algorithm to truncate a list of batched segments to fit
-        within `sequence_length`. Only used if
-        `pack_inputs` is True. 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
-preprocessor = keras_nlp.models.XLMRobertaPreprocessor(proto="model.spm")
-
-# Tokenize and pack a single sentence directly.
-preprocessor("The quick brown fox jumped.")
-
-# Tokenize and pack a multiple sentence directly.
-preprocessor(("The quick brown fox jumped.", "Call me Ishmael."))
-
-# Map a dataset to preprocess a single sentence.
-features = ["The quick brown fox jumped.", "I forgot my homework."]
-labels = [0, 1]
-ds = tf.data.Dataset.from_tensor_slices((features, labels))
-ds = ds.map(
-    lambda x, y: (preprocessor(x), y),
-    num_parallel_calls=tf.data.AUTOTUNE,
-)
-
-# Map a dataset to preprocess a multiple sentences.
-first_sentences = ["The quick brown fox jumped.", "Call me Ishmael."]
-second_sentences = ["The fox tripped.", "Oh look, a whale."]
-labels = [1, 1]
-ds = tf.data.Dataset.from_tensor_slices(((first_sentences, second_sentences), labels))
-ds = ds.map(
-    lambda x, y: (preprocessor(x), y),
-    num_parallel_calls=tf.data.AUTOTUNE,
-)
-```
-"""
-
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
 class XLMRobertaPreprocessor(keras.layers.Layer):
+    """XLM-RoBERTa preprocessing layer.
+
+    This preprocessing layer will do three things:
+
+    - Tokenize any number of inputs using a
+      `keras_nlp.tokenizers.SentencePieceTokenizer`.
+    - Pack the inputs together with the appropriate `"
"`, `"
"` and
+      `"<pad>"` tokens, i.e., adding a single `"
"` at the start of the
+      entire sequence, `"

"` at the end of each segment, save the last
+      and a `"
"` at the end of the entire sequence.
+    - Construct a dictionary with keys `"token_ids"` and `"padding_mask"`
+      that can be passed directly to a XLM-RoBERTa model.
+
+    Note that the original fairseq implementation modifies the indices of the
+    SentencePiece tokenizer output. To preserve compatibility, we make the same
+    changes, i.e., `"
"`, `"<pad>"`, `"
"` and `"<unk>"` are mapped to
+    0, 1, 2, 3, respectively, and non-special tokens' indices are shifted right
+    by one. Keep this in mind if generating your own vocabulary for tokenization.
+
+    This layer will accept either a tuple of (possibly batched) inputs, or a
+    single input tensor. If a single tensor is passed, it will be packed
+    equivalently to a tuple with a single element.
+
+    The SentencePiece tokenizer can be accessed via the `tokenizer` property on
+    this layer, and can be used directly for custom packing on inputs.
+
+    Args:
+        proto: Either a `string` path to a SentencePiece proto file, a `bytes`
+            object with a serialized SentencePiece proto.
+        sequence_length: The length of the packed inputs.
+        truncate: 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
+    preprocessor = keras_nlp.models.XLMRobertaPreprocessor(proto="model.spm")
+
+    # Tokenize and pack a single sentence directly.
+    preprocessor("The quick brown fox jumped.")
+
+    # Tokenize and pack a multiple sentence directly.
+    preprocessor(("The quick brown fox jumped.", "Call me Ishmael."))
+
+    # Map a dataset to preprocess a single sentence.
+    features = ["The quick brown fox jumped.", "I forgot my homework."]
+    labels = [0, 1]
+    ds = tf.data.Dataset.from_tensor_slices((features, labels))
+    ds = ds.map(
+        lambda x, y: (preprocessor(x), y),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+
+    # Map a dataset to preprocess a multiple sentences.
+    first_sentences = ["The quick brown fox jumped.", "Call me Ishmael."]
+    second_sentences = ["The fox tripped.", "Oh look, a whale."]
+    labels = [1, 1]
+    ds = tf.data.Dataset.from_tensor_slices(((first_sentences, second_sentences), labels))
+    ds = ds.map(
+        lambda x, y: (preprocessor(x), y),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+    ```
+    """
+
     def __init__(
         self,
         proto,
@@ -162,5 +160,12 @@ class XLMRobertaPreprocessor(keras.layers.Layer):
             "padding_mask": token_ids != self.pad_token_id,
         }
 
-
-setattr(XLMRobertaPreprocessor, "__doc__", PREPROCESSOR_DOCSTRING)
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        sequence_length=None,
+        truncate="round_robin",
+        **kwargs,
+    ):
+        raise NotImplementedError
diff --git a/keras_nlp/tokenizers/__init__.py b/keras_nlp/tokenizers/__init__.py
index 048ea15d..6c0d1972 100644
--- a/keras_nlp/tokenizers/__init__.py
+++ b/keras_nlp/tokenizers/__init__.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_nlp.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 from keras_nlp.tokenizers.byte_tokenizer import ByteTokenizer
 from keras_nlp.tokenizers.sentence_piece_tokenizer import SentencePieceTokenizer
 from keras_nlp.tokenizers.sentence_piece_tokenizer_trainer import (
diff --git a/keras_nlp/tokenizers/byte_pair_tokenizer.py b/keras_nlp/tokenizers/byte_pair_tokenizer.py
new file mode 100644
index 00000000..def254e8
--- /dev/null
+++ b/keras_nlp/tokenizers/byte_pair_tokenizer.py
@@ -0,0 +1,547 @@
+# 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.
+
+""" Byte-pair encoder implementation.
+
+This file implements the same logic as openai BPE:
+https://github.com/openai/gpt-2/blob/master/src/encoder.py,
+but is TF graph compatible.
+"""
+
+import json
+from typing import Iterable
+from typing import List
+
+import tensorflow as tf
+import tensorflow_text as tf_text
+from tensorflow import keras
+
+from keras_nlp.tokenizers import tokenizer
+
+# As python and TF handles special spaces differently, we need to
+# manually handle special spaces during string split.
+SPECIAL_WHITESPACES = r"\x{a0}\x{2009}\x{202f}\x{3000}"
+
+# String splitting regex pattern.
+SPLIT_PATTERN_1 = r"""'s|'t|'re|'ve|'m|'ll|'d
+    |[\s{special_spaces}]+[\n\r\t\f६{special_spaces}]| ?\p{L}+
+    | ?[\p{N}]+| ?[^\s\p{L}\p{N}{special_spaces}]+""".replace(
+    "{special_spaces}", SPECIAL_WHITESPACES
+)
+
+SPLIT_PATTERN_2 = rf"""[\s६{SPECIAL_WHITESPACES}]$"""
+
+
+def bytes_to_unicode():
+    bs = (
+        list(range(ord("!"), ord("~") + 1))
+        + list(range(ord("¡"), ord("¬") + 1))
+        + list(range(ord("®"), ord("ÿ") + 1))
+    )
+    cs = bs[:]
+    n = 0
+    # removes mapping an int to a whitespace character
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    bs = [n.to_bytes(1, "little") for n in bs]
+    return bs, cs  # int to string mapping
+
+
+def remove_strings_from_inputs(tensor, string_to_remove):
+    """Remove certain strings from input tensor."""
+    non_empty_mask = tensor != string_to_remove
+    flatten_indexes = tf.where(non_empty_mask)
+    flatten_result = tf.gather_nd(tensor, flatten_indexes)
+    row_lengths = tf.reduce_sum(tf.cast(non_empty_mask, tf.int64), axis=1)
+    result = tf.RaggedTensor.from_row_lengths(
+        values=flatten_result,
+        row_lengths=row_lengths,
+    )
+    return result
+
+
+def split_strings_for_bpe(inputs):
+    # We need to recreate the exact behavior of token presplitting in the
+    # original gpt2 tokenizer which uses a lookahead. As re2 does not
+    # support lookahead match, we are using an alternative insert a special
+    # token "६" before leading space of non-space characters and after the
+    # trailing space, e.g., " keras" will be "६ keras".
+    inputs = tf.strings.regex_replace(
+        inputs, rf"( )([^\s{SPECIAL_WHITESPACES}])", r"६\1\2"
+    )
+    inputs = tf.strings.regex_replace(
+        inputs, rf"(\s{SPECIAL_WHITESPACES})$", r"\1६"
+    )
+    raw_tokens = tf_text.regex_split(inputs, SPLIT_PATTERN_1, SPLIT_PATTERN_1)
+    # Second pass splits out the last whilespace char or "६".
+    raw_tokens = tf_text.regex_split(
+        raw_tokens, SPLIT_PATTERN_2, SPLIT_PATTERN_2
+    )
+    if raw_tokens.shape.rank > 2:
+        raw_tokens = raw_tokens.merge_dims(1, 2)
+    return remove_strings_from_inputs(raw_tokens, "६")
+
+
+class BytePairTokenizerCache(tf.Module):
+    """Cache that stores the encoded result of seen tokens.
+
+    The cache key is string tensor or python strings, and the value is split
+    tokens joined by whitespace. For example, "dragonfly" => "dragon fly"
+
+    Examples:
+    ```
+    cache = BytePairTokenizerCache()
+    cache.insert(["butterfly", "dragonfly"], ["but ter fly", "dragon fly"])
+    cache.lookup(["butterfly"])
+    ```
+    """
+
+    def __init__(self):
+        # `tf.lookup.experimental.MutableHashTable` does not support string to
+        # string mapping. So we first convert to string to an integer key, and
+        # use the integer key to find the value.
+        self.factors = tf.pow(256, tf.range(0, 8, dtype=tf.int64))
+        self.id2value = tf.lookup.experimental.MutableHashTable(
+            tf.int64, tf.string, ""
+        )
+
+    def _get_key(self, keys):
+        """Get the hash key for given inputs."""
+        # `tf.fingerprint` converts token to a array of uint8 of length 8, we
+        # need to convert it to a uint64.
+        return tf.squeeze(
+            tf.matmul(
+                tf.cast(tf.fingerprint(keys), dtype=tf.int64),
+                self.factors[:, tf.newaxis],
+            ),
+            -1,
+        )
+
+    def lookup(self, keys):
+        """Look up the encoded outputs of given tokens."""
+        ids = self._get_key(keys)
+        result = self.id2value.lookup(ids)
+        # Ensure output shape for graph mode.
+        result.set_shape([None])
+        return result
+
+    def insert(self, keys, values):
+        """Insert token <=> encoded outputs pairs."""
+        self.id2value.insert(self._get_key(keys), values)
+
+
+def create_static_hashtable(keys, values, default):
+    return tf.lookup.StaticHashTable(
+        tf.lookup.KeyValueTensorInitializer(
+            tf.convert_to_tensor(keys),
+            tf.convert_to_tensor(values),
+        ),
+        default_value=default,
+    )
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class BytePairTokenizer(tokenizer.Tokenizer):
+    """Bype-pair encoding tokenizer layer.
+
+    This BPE tokenizer provides the same funtionality as official GPT2
+    tokenizer. Given the same `vocabulary` which maps tokens to ids, and `merges`
+    which describes BPE merge rules, it should provide the same output
+    as openai implementation (https://github.com/openai/gpt-2/blob/master/src/encoder.py).
+    Different from openai, this implementation is graph-compatible, so you can
+    use it within a tf.data pipeline.
+
+    If input is a batch of strings (rank > 0):
+    By default, the layer will output a `tf.RaggedTensor` where the last
+    dimension of the output is ragged. If `sequence_length` is set, the layer
+    will output a dense `tf.Tensor` where all inputs have been padded or
+    truncated to `sequence_length`.
+    If input is a scalar string (rank == 0):
+    By default, the layer will output a dense `tf.Tensor` with static shape
+    `[None]`. If `sequence_length` is set, the output will be
+    a dense `tf.Tensor` of shape `[sequence_length]`.
+
+    Args:
+        vocabulary: string or dict, maps token to integer ids. If it is a
+            string, it should be the file path to a json file.
+        merges: string or list, contains the merge rule. If it is a string,
+            it should be the file path to merge rules. The merge rule file
+            should have one merge rule per line.
+        sequence_length: int, defaults to None. If set, the output will be
+            padded or truncated to the `sequence_length`.
+
+    Examples:
+
+    Use in-momery vocabulary and merge list.
+
+    >>> vocab = {"butter": 1, "fly": 2}
+    >>> merge = ["b u", "t t", "e r", "bu tt", "butt er", "f l", "fl y"]
+    >>> tokenizer = keras_nlp.tokenizers.BytePairTokenizer(vocab, merge)
+    >>> tokenizer("butterfly")
+    <tf.Tensor: shape=(2,), dtype=int32, numpy=array([1, 2], dtype=int32)>
+    >>> tokenizer(["butterfly"])
+    <tf.RaggedTensor [[1, 2]]>
+    >>> tokenizer(["butterfly", "butter"])
+    <tf.RaggedTensor [[1, 2], [1]]>
+    >>> tokenizer = keras_nlp.tokenizers.BytePairTokenizer(
+    ...     vocab, merge, sequence_length=2)
+    >>> tokenizer(["butterfly", "butter"])
+    <tf.Tensor: shape=(2, 2), dtype=int32, numpy=
+    array([[1, 2],
+           [1, 0]], dtype=int32)>
+
+    Use hosted vocabluary and merge list.
+
+    ```python
+    vocab_path = tf.keras.utils.get_file(
+        "vocab.json",
+        "https://storage.googleapis.com/keras-nlp/models/roberta_base/vocab.json",
+    )
+    merge_path = tf.keras.utils.get_file(
+        "merges.txt",
+        "https://storage.googleapis.com/keras-nlp/models/roberta_base/merges.txt",
+    )
+    tokenizer = BytePairTokenizer(
+        vocabulary=vocab_path, merges=merge_path
+    )
+    tokenizer("Butterfly is not flying butter!")
+    ```
+
+    Detokenize
+    >>> vocab = {"butter": 1, "fly": 2}
+    >>> merge = ["b u", "t t", "e r", "bu tt", "butt er", "f l", "fl y"]
+    >>> tokenizer = keras_nlp.tokenizers.BytePairTokenizer(vocab, merge)
+    >>> tokenizer.detokenize([[1, 2]])
+    <tf.Tensor: shape=(1,), dtype=string, numpy=array([b'butterfly'],
+    dtype=object)>
+
+    """
+
+    def __init__(
+        self,
+        vocabulary,
+        merges,
+        sequence_length=None,
+        **kwargs,
+    ) -> None:
+        # Check dtype and provide a default.
+        if "dtype" not in kwargs or kwargs["dtype"] is None:
+            kwargs["dtype"] = tf.int32
+        else:
+            dtype = tf.dtypes.as_dtype(kwargs["dtype"])
+            if not dtype.is_integer:
+                raise ValueError(
+                    "Output dtype must be an integer type or a string. "
+                    f"Received: `dtype={dtype}`"
+                )
+
+        super().__init__(**kwargs)
+
+        if isinstance(vocabulary, str):
+            with open(vocabulary, "r") as f:
+                self.vocabulary = json.load(f)
+        elif isinstance(vocabulary, dict):
+            self.vocabulary = vocabulary.copy()
+        else:
+            raise ValueError(
+                "Vocabulary must be an file path or dictionary mapping string "
+                f"token to int ids. Received: `type(vocabulary)={type(vocabulary)}`."
+            )
+        if isinstance(merges, str):
+            self.merges = [bp.rstrip() for bp in tf.io.gfile.GFile(merges)]
+        elif isinstance(merges, Iterable):
+            self.merges = list(merges)
+        else:
+            raise ValueError(
+                "Merges must be a file path or a list of merge rules. "
+                f"Received: `type(merges)={type(merges)}`"
+            )
+        self.sequence_length = sequence_length
+
+        # Create byte <=> unicode mapping. This is useful for handling
+        # whitespace tokens.
+        byte_list, unicode_list = bytes_to_unicode()
+        self.byte2unicode = create_static_hashtable(
+            byte_list, unicode_list, default=""
+        )
+        self.unicode2byte = create_static_hashtable(
+            unicode_list, byte_list, default=""
+        )
+
+        self.cache = BytePairTokenizerCache()
+
+        # Create mapping between string tokens to int ids, and vice versa.
+        byte_pairs = [x[0] for x in self.vocabulary.items()]
+        byte_pair_encoding_indices = [x[1] for x in self.vocabulary.items()]
+        self.token_to_id_map = create_static_hashtable(
+            byte_pairs,
+            byte_pair_encoding_indices,
+            default=-1,
+        )
+        self.id_to_token_map = create_static_hashtable(
+            byte_pair_encoding_indices,
+            byte_pairs,
+            default="",
+        )
+
+        # Create ranking of merge rules, this is the same as order of merge
+        # pairs in `self.merges`.
+        self.merge_ranks_lookup_default = len(self.merges) + 1
+        self.merge_ranks = create_static_hashtable(
+            self.merges,
+            list(range(len(self.merges))),
+            default=self.merge_ranks_lookup_default,
+        )
+
+    def get_vocabulary(self) -> List[str]:
+        """Get the tokenizer vocabulary as a list of strings tokens."""
+        return self.vocabulary.keys()
+
+    def vocabulary_size(self) -> int:
+        """Get the size of the tokenizer vocabulary."""
+        return len(self.vocabulary)
+
+    def id_to_token(self, id: int) -> str:
+        """Convert an integer id to a string token."""
+        # This will be slow, but keep memory usage down compared to building a
+        # dict. Assuming the main use case is looking up a few special tokens
+        # early in the vocab, this should be fine.
+
+        keys = self.get_vocabulary()
+        for token in keys:
+            if self.vocabulary[token] == id:
+                return token
+        return None
+
+    def token_to_id(self, token: str) -> int:
+        """Convert a string token to an integer id."""
+        return self.vocabulary[token]
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                # Ideally vocabulary and merge list would be saved as plain text
+                # assets in the saved model. We have no good way to support
+                # this currently, so we save the vocabulary in the config.
+                "vocabulary": self.vocabulary,
+                "merges": self.merges,
+                "sequence_length": self.sequence_length,
+            }
+        )
+        return config
+
+    @tf.function
+    def _bpe_merge_one_step(self, words, mask):
+        """Perform one step of byte-pair merge."""
+        # Get all word pairs.
+        first, second = words[:, :-1], words[:, 1:]
+
+        # Mask empty.
+        non_empty_mask = second.nested_row_lengths()[0] != 0
+        mask = mask & non_empty_mask
+        if not tf.reduce_any(mask):
+            return [words, mask]
+        non_empty_indices = tf.boolean_mask(tf.range(tf.shape(mask)[0]), mask)
+        filterd_first = tf.ragged.boolean_mask(first, mask)
+        filtered_second = tf.ragged.boolean_mask(second, mask)
+
+        # Get byte pair ranking in merge rules.
+        pairs = tf.strings.join([filterd_first, filtered_second], separator=" ")
+        pair_rank = self.merge_ranks.lookup(pairs)
+
+        # Get BPE pair ranks.
+        min_pair_rank = tf.reduce_min(pair_rank, axis=1)
+        pair_found_mask = min_pair_rank != self.merge_ranks_lookup_default
+
+        # Tokens that cannot be further merged are marked as finished.
+        mask = tf.tensor_scatter_nd_update(
+            mask, tf.expand_dims(non_empty_indices, axis=1), pair_found_mask
+        )
+        if not tf.math.reduce_any(mask):
+            return [words, mask]
+
+        masked_pair_rank = tf.ragged.boolean_mask(pair_rank, pair_found_mask)
+        min_pair_rank_indices = tf.math.argmin(
+            masked_pair_rank.to_tensor(self.merge_ranks_lookup_default), axis=1
+        )
+
+        # Get words and pairs to process.
+        unfinished_words = tf.ragged.boolean_mask(words, mask)
+
+        pair_left = tf.gather(
+            unfinished_words, min_pair_rank_indices, batch_dims=1
+        )
+        pair_right = tf.gather(
+            unfinished_words, min_pair_rank_indices + 1, batch_dims=1
+        )
+
+        merged_pairs = tf.strings.join([pair_left, pair_right])
+        empty_strs = tf.fill(tf.shape(merged_pairs), "")
+
+        unfinished_word_indices = tf.cast(
+            tf.boolean_mask(tf.range(tf.shape(mask)[0]), mask), dtype=tf.int64
+        )
+        merged_pair_indices = tf.concat(
+            [
+                unfinished_word_indices[:, tf.newaxis],
+                min_pair_rank_indices[:, tf.newaxis],
+            ],
+            axis=1,
+        )
+        empty_string_indices = tf.concat(
+            [
+                unfinished_word_indices[:, tf.newaxis],
+                min_pair_rank_indices[:, tf.newaxis] + 1,
+            ],
+            axis=1,
+        )
+
+        tensor_words = words.to_tensor(default_value="")
+        tensor_words = tf.tensor_scatter_nd_update(
+            tensor_words,
+            merged_pair_indices,
+            merged_pairs,
+        )
+
+        words = tf.tensor_scatter_nd_update(
+            tensor_words,
+            empty_string_indices,
+            empty_strs,
+        )
+        # Remove empty strings.
+        words = remove_strings_from_inputs(words, "")
+        return [words, mask]
+
+    def _bpe_merge(self, inputs):
+        """Perform byte-pair merge for each word in the inputs."""
+        num_words = tf.shape(inputs)[0]
+
+        # Merge bytes.
+        def loop_condition(_, mask):
+            return tf.math.reduce_any(mask)
+
+        initial_mask = tf.fill((num_words,), True)
+        merged_words, _ = tf.while_loop(
+            loop_condition,
+            self._bpe_merge_one_step,
+            loop_vars=[
+                inputs,
+                initial_mask,
+            ],
+            shape_invariants=[
+                tf.TensorShape([None, None]),
+                tf.TensorShape([None]),
+            ],
+        )
+        return merged_words
+
+    def tokenize(self, inputs):
+        if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
+            inputs = tf.convert_to_tensor(inputs)
+
+        scalar_input = inputs.shape.rank == 0
+        if scalar_input:
+            inputs = tf.expand_dims(inputs, 0)
+
+        raw_tokens = split_strings_for_bpe(inputs)
+        token_row_splits = raw_tokens.row_splits
+        flat_tokens = raw_tokens.flat_values
+
+        # Check cache.
+        cache_lookup = self.cache.lookup(flat_tokens)
+        cache_mask = cache_lookup == ""
+
+        has_unseen_words = tf.math.reduce_any(
+            (cache_lookup == "") & (flat_tokens != "")
+        )
+
+        def process_unseen_tokens():
+            unseen_tokens = tf.boolean_mask(flat_tokens, cache_mask)
+            self._bpe_merge_and_update_cache(unseen_tokens)
+            return self.cache.lookup(flat_tokens)
+
+        # If `has_unseen_words == True`, it means not all tokens are in cache,
+        # we will process the unseen tokens. Otherwise return the cache lookup.
+        tokenized_words = tf.cond(
+            has_unseen_words,
+            process_unseen_tokens,
+            lambda: cache_lookup,
+        )
+
+        # Encode merged tokens.
+        tokenized_words = tf.strings.split(tokenized_words, sep=" ")
+        encoding = self.token_to_id_map.lookup(tokenized_words)
+
+        # Unflatten to match input.
+        encoding = tf.RaggedTensor.from_row_splits(
+            encoding.flat_values,
+            tf.gather(encoding.row_splits, token_row_splits),
+        )
+
+        # Convert to a dense output if `sequence_length` is set.
+        if self.sequence_length:
+            output_shape = encoding.shape.as_list()
+            output_shape[-1] = self.sequence_length
+            encoding = encoding.to_tensor(shape=output_shape)
+
+        # Convert to a dense output if input in scalar
+        if scalar_input:
+            encoding = tf.squeeze(encoding, 0)
+            tf.ensure_shape(encoding, shape=[self.sequence_length])
+
+        return encoding
+
+    def detokenize(self, inputs):
+        if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
+            inputs = tf.convert_to_tensor(inputs)
+
+        scalar_input = inputs.shape.rank == 0
+        if scalar_input:
+            inputs = tf.expand_dims(inputs, 0)
+
+        unicode_text = tf.strings.reduce_join(
+            self.id_to_token_map.lookup(inputs), axis=1
+        )
+        split_unicode_text = tf.strings.unicode_split(unicode_text, "UTF-8")
+        byte_text = tf.strings.reduce_join(
+            self.unicode2byte.lookup(split_unicode_text)
+        )
+
+        if not scalar_input:
+            byte_text = tf.expand_dims(byte_text, 0)
+
+        return byte_text
+
+    def _transform_bytes(self, tokens):
+        """Map token bytes to unicode using `byte2unicode`."""
+        split_bytes = tf.strings.bytes_split(tokens)
+        split_unicode = self.byte2unicode.lookup(split_bytes)
+        return split_unicode
+
+    def _bpe_merge_and_update_cache(self, tokens):
+        """Process unseen tokens and add to cache."""
+        words = self._transform_bytes(tokens)
+        tokenized_words = self._bpe_merge(words)
+
+        # For each word, join all its token by a whitespace,
+        # e.g., ["dragon", "fly"] => "dragon fly" for hash purpose.
+        tokenized_words = tf.strings.reduce_join(
+            tokenized_words, axis=1, separator=" "
+        )
+        self.cache.insert(tokens, tokenized_words)


Test patch:
diff --git a/keras_nlp/models/bert/bert_preprocessing_test.py b/keras_nlp/models/bert/bert_preprocessing_test.py
index 7b5f1295..42b60477 100644
--- a/keras_nlp/models/bert/bert_preprocessing_test.py
+++ b/keras_nlp/models/bert/bert_preprocessing_test.py
@@ -13,14 +13,17 @@
 # limitations under the License.
 """Tests for BERT preprocessing layers."""
 
+import os
 import unittest
 
 import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
 
 from keras_nlp.models.bert.bert_preprocessing import BertPreprocessor
 
 
-class BertPreprocessorTest(tf.test.TestCase):
+class BertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         self.vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
         self.vocab += ["THE", "QUICK", "BROWN", "FOX"]
@@ -163,3 +166,23 @@ class BertPreprocessorTest(tf.test.TestCase):
                 sequence_length=1024,
             )
         get_file_mock.assert_called_once()
+
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
+        input_data = tf.constant(["THE QUICK BROWN FOX."])
+        preprocessor = BertPreprocessor(
+            vocabulary=self.vocab,
+            sequence_length=8,
+        )
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = preprocessor(inputs)
+        model = keras.Model(inputs, outputs)
+        path = os.path.join(self.get_temp_dir(), "model")
+        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/distilbert/distilbert_models_test.py b/keras_nlp/models/distilbert/distilbert_models_test.py
index a7114fcd..89e4a131 100644
--- a/keras_nlp/models/distilbert/distilbert_models_test.py
+++ b/keras_nlp/models/distilbert/distilbert_models_test.py
@@ -19,13 +19,12 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.distilbert.distilbert_models import DistilBertBase
-from keras_nlp.models.distilbert.distilbert_models import DistilBertCustom
+from keras_nlp.models.distilbert.distilbert_models import DistilBert
 
 
 class DistilBertTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = DistilBertCustom(
+        self.model = DistilBert(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
@@ -63,33 +62,19 @@ class DistilBertTest(tf.test.TestCase, parameterized.TestCase):
             }
             self.model(input_data)
 
-    def test_valid_call_distilbert_base(self):
-        model = DistilBertBase(vocabulary_size=1000, name="encoder")
-        input_data = {
-            "token_ids": tf.ones(
-                (self.batch_size, self.model.max_sequence_length), dtype="int32"
-            ),
-            "padding_mask": tf.ones(
-                (self.batch_size, self.model.max_sequence_length), dtype="int32"
-            ),
-        }
-        model(input_data)
-
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
     def test_distilbert_base_compile(self, jit_compile):
-        model = DistilBertBase(vocabulary_size=1000, name="encoder")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_batch)
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_batch)
 
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
     def test_distilbert_base_compile_batched_ds(self, jit_compile):
-        model = DistilBertBase(vocabulary_size=1000, name="encoder")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_dataset)
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_dataset)
 
     @parameterized.named_parameters(
         ("save_format_tf", "tf"), ("save_format_h5", "h5")
@@ -101,7 +86,7 @@ class DistilBertTest(tf.test.TestCase, parameterized.TestCase):
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
-        self.assertIsInstance(restored_model, DistilBertCustom)
+        self.assertIsInstance(restored_model, DistilBert)
 
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
diff --git a/keras_nlp/models/gpt2/gpt2_models_test.py b/keras_nlp/models/gpt2/gpt2_models_test.py
index e361b72c..3b0b3a34 100644
--- a/keras_nlp/models/gpt2/gpt2_models_test.py
+++ b/keras_nlp/models/gpt2/gpt2_models_test.py
@@ -19,20 +19,18 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.gpt2.gpt2_models import Gpt2Base
-from keras_nlp.models.gpt2.gpt2_models import Gpt2Custom
+from keras_nlp.models.gpt2.gpt2_models import Gpt2
 
 
 class Gpt2Test(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = Gpt2Custom(
+        self.model = Gpt2(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
             hidden_dim=64,
             intermediate_dim=128,
             max_sequence_length=128,
-            name="gpt2_test",
         )
         self.batch_size = 8
         self.input_batch = {
@@ -51,6 +49,9 @@ class Gpt2Test(tf.test.TestCase, parameterized.TestCase):
     def test_valid_call_gpt2(self):
         self.model(self.input_batch)
 
+        # Check default name passed through
+        self.assertEqual(self.model.name, "backbone")
+
     def test_variable_sequence_length_call_gpt2(self):
         for seq_length in (25, 50, 75):
             input_data = {
@@ -63,25 +64,19 @@ class Gpt2Test(tf.test.TestCase, parameterized.TestCase):
             }
             self.model(input_data)
 
-    def test_valid_call_gpt2_base(self):
-        model = Gpt2Base(vocabulary_size=1000, name="gpt2_base_test")
-        model(self.input_batch)
-
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
-    def test_gpt2_base_compile(self, jit_compile):
-        model = Gpt2Base(vocabulary_size=1000, name="gpt2_base_test")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_batch)
+    def test_gpt2_compile(self, jit_compile):
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_batch)
 
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
-    def test_gpt2_base_compile_batched_ds(self, jit_compile):
-        model = Gpt2Base(vocabulary_size=1000, name="gpt2_base_test")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_dataset)
+    def test_gpt2_compile_batched_ds(self, jit_compile):
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_dataset)
 
     @parameterized.named_parameters(
         ("save_format_tf", "tf"), ("save_format_h5", "h5")
@@ -93,7 +88,7 @@ class Gpt2Test(tf.test.TestCase, parameterized.TestCase):
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
-        self.assertIsInstance(restored_model, Gpt2Custom)
+        self.assertIsInstance(restored_model, Gpt2)
 
         # Check that output matches.
         restored_output = restored_model(self.input_batch)
diff --git a/keras_nlp/models/roberta/roberta_models_test.py b/keras_nlp/models/roberta/roberta_models_test.py
index d437b22f..2451443a 100644
--- a/keras_nlp/models/roberta/roberta_models_test.py
+++ b/keras_nlp/models/roberta/roberta_models_test.py
@@ -20,21 +20,18 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.roberta.roberta_models import RobertaBase
-from keras_nlp.models.roberta.roberta_models import RobertaCustom
-from keras_nlp.models.roberta.roberta_tasks import RobertaClassifier
+from keras_nlp.models.roberta.roberta_models import Roberta
 
 
 class RobertaTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = RobertaCustom(
+        self.model = Roberta(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
             hidden_dim=64,
             intermediate_dim=128,
             max_sequence_length=128,
-            name="encoder",
         )
         self.batch_size = 8
         self.input_batch = {
@@ -53,53 +50,22 @@ class RobertaTest(tf.test.TestCase, parameterized.TestCase):
     def test_valid_call_roberta(self):
         self.model(self.input_batch)
 
-    def test_valid_call_classifier(self):
-        classifier = RobertaClassifier(self.model, 4, 128, name="classifier")
-        classifier(self.input_batch)
-
-    def test_valid_call_roberta_base(self):
-        model = RobertaBase(vocabulary_size=1000, name="encoder")
-        input_data = {
-            "token_ids": tf.ones(
-                (self.batch_size, model.max_sequence_length), dtype="int32"
-            ),
-            "padding_mask": tf.ones(
-                (self.batch_size, model.max_sequence_length), dtype="int32"
-            ),
-        }
-        model(input_data)
-
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_roberta_base_compile(self, jit_compile):
-        model = RobertaBase(vocabulary_size=1000, name="encoder")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_batch)
-
-    @parameterized.named_parameters(
-        ("jit_compile_false", False), ("jit_compile_true", True)
-    )
-    def test_roberta_base_compile_batched_ds(self, jit_compile):
-        model = RobertaBase(vocabulary_size=1000, name="encoder")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_dataset)
+        # Check default name passed through
+        self.assertEqual(self.model.name, "backbone")
 
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
-    def test_roberta_classifier_compile(self, jit_compile):
-        model = RobertaClassifier(self.model, 4, 128, name="classifier")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_batch)
+    def test_roberta_compile(self, jit_compile):
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_batch)
 
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
-    def test_roberta_classifier_compile_batched_ds(self, jit_compile):
-        model = RobertaClassifier(self.model, 4, 128, name="classifier")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_dataset)
+    def test_roberta_compile_batched_ds(self, jit_compile):
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_dataset)
 
     def test_variable_sequence_length_call_roberta(self):
         for seq_length in (25, 50, 75):
@@ -121,18 +87,14 @@ class RobertaTest(tf.test.TestCase, parameterized.TestCase):
         ("save_format_tf", "tf"), ("save_format_h5", "h5")
     )
     def test_saving_model(self, save_format):
-        model_output = self.model.predict(self.input_batch)
-
+        model_output = self.model(self.input_batch)
         save_path = os.path.join(self.get_temp_dir(), "model")
         self.model.save(save_path, save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
-        self.assertIsInstance(restored_model, RobertaCustom)
+        self.assertIsInstance(restored_model, Roberta)
 
         # Check that output matches.
-        restored_output = restored_model.predict(self.input_batch)
-        self.assertAllClose(
-            model_output,
-            restored_output,
-        )
+        restored_output = restored_model(self.input_batch)
+        self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/models/roberta/roberta_tasks_test.py b/keras_nlp/models/roberta/roberta_tasks_test.py
index 27557f1e..f50b1dbd 100644
--- a/keras_nlp/models/roberta/roberta_tasks_test.py
+++ b/keras_nlp/models/roberta/roberta_tasks_test.py
@@ -19,20 +19,19 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.roberta.roberta_models import RobertaCustom
+from keras_nlp.models.roberta.roberta_models import Roberta
 from keras_nlp.models.roberta.roberta_tasks import RobertaClassifier
 
 
 class RobertaClassifierTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = RobertaCustom(
+        self.model = Roberta(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
             hidden_dim=64,
             intermediate_dim=128,
             max_sequence_length=128,
-            name="encoder",
         )
         self.batch_size = 8
         self.input_batch = {
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_models_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_models_test.py
index 93bba6d4..e6c09dc0 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_models_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_models_test.py
@@ -20,20 +20,18 @@ import tensorflow as tf
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.xlm_roberta.xlm_roberta_models import XLMRobertaBase
-from keras_nlp.models.xlm_roberta.xlm_roberta_models import XLMRobertaCustom
+from keras_nlp.models.xlm_roberta.xlm_roberta_models import XLMRoberta
 
 
 class XLMRobertaTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = XLMRobertaCustom(
+        self.model = XLMRoberta(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
             hidden_dim=64,
             intermediate_dim=128,
             max_sequence_length=128,
-            name="encoder",
         )
         self.batch_size = 8
         self.input_batch = {
@@ -52,33 +50,22 @@ class XLMRobertaTest(tf.test.TestCase, parameterized.TestCase):
     def test_valid_call_xlm_roberta(self):
         self.model(self.input_batch)
 
-    def test_valid_call_xlm_roberta_base(self):
-        model = XLMRobertaBase(vocabulary_size=1000, name="encoder")
-        input_data = {
-            "token_ids": tf.ones(
-                (self.batch_size, model.max_sequence_length), dtype="int32"
-            ),
-            "padding_mask": tf.ones(
-                (self.batch_size, model.max_sequence_length), dtype="int32"
-            ),
-        }
-        model(input_data)
+        # Check default name passed through
+        self.assertEqual(self.model.name, "backbone")
 
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
-    def test_xlm_roberta_base_compile(self, jit_compile):
-        model = XLMRobertaBase(vocabulary_size=1000, name="encoder")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_batch)
+    def test_xlm_roberta_compile(self, jit_compile):
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_batch)
 
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
-    def test_xlm_roberta_base_compile_batched_ds(self, jit_compile):
-        model = XLMRobertaBase(vocabulary_size=1000, name="encoder")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_dataset)
+    def test_xlm_roberta_compile_batched_ds(self, jit_compile):
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_dataset)
 
     def test_variable_sequence_length_call_xlm_roberta(self):
         for seq_length in (25, 50, 75):
@@ -100,18 +87,14 @@ class XLMRobertaTest(tf.test.TestCase, parameterized.TestCase):
         ("save_format_tf", "tf"), ("save_format_h5", "h5")
     )
     def test_saving_model(self, save_format):
-        model_output = self.model.predict(self.input_batch)
-
+        model_output = self.model(self.input_batch)
         save_path = os.path.join(self.get_temp_dir(), "model")
         self.model.save(save_path, save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
-        self.assertIsInstance(restored_model, XLMRobertaCustom)
+        self.assertIsInstance(restored_model, XLMRoberta)
 
         # Check that output matches.
-        restored_output = restored_model.predict(self.input_batch)
-        self.assertAllClose(
-            model_output,
-            restored_output,
-        )
+        restored_output = restored_model(self.input_batch)
+        self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/tokenizers/byte_pair_tokenizer_test.py b/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
new file mode 100644
index 00000000..243891e5
--- /dev/null
+++ b/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
@@ -0,0 +1,131 @@
+# 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 os
+
+import pytest
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.tokenizers.byte_pair_tokenizer import BytePairTokenizer
+
+VOCAB_PATH = keras.utils.get_file(
+    None,
+    "https://storage.googleapis.com/keras-nlp/models/roberta_base/vocab.json",
+)
+MERGE_PATH = keras.utils.get_file(
+    None,
+    "https://storage.googleapis.com/keras-nlp/models/roberta_base/merges.txt",
+)
+
+
+@pytest.mark.slow
+class BytePairTokenizerTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        super().setUp()
+
+        self.tokenizer = BytePairTokenizer(
+            vocabulary=VOCAB_PATH, merges=MERGE_PATH
+        )
+
+    def test_tokenize_list_input(self):
+        input_data = ["brown.", "black."]
+        call_output = self.tokenizer(input_data)
+        tokenize_output = self.tokenizer.tokenize(input_data)
+        expected = tf.ragged.constant([[31876, 4], [14178, 4]])
+        self.assertAllEqual(call_output, expected)
+        self.assertAllEqual(tokenize_output, expected)
+
+        input_data = tf.convert_to_tensor(["brown.", "black."])
+        encoded = self.tokenizer(input_data)
+        self.assertAllEqual(encoded, expected)
+
+    def test_tokenize_scalar_input(self):
+        input_data = "brown."
+        encoded = self.tokenizer.tokenize(input_data)
+        self.assertAllEqual(encoded, [31876, 4])
+
+    def test_detokenize(self):
+        input_data = ["brown."]
+        encoded = self.tokenizer.tokenize(input_data)
+        decoded = self.tokenizer.detokenize(encoded)
+        self.assertAllEqual(input_data, decoded)
+
+    def test_whitespace_split(self):
+        input_data = "\n\n\n  s"
+        encoded = self.tokenizer(input_data)
+        self.assertAllEqual(encoded, [50140, 50118, 1437, 579])
+
+        input_data = "  \n\n\ns"
+        encoded = self.tokenizer(input_data)
+        self.assertAllEqual(encoded, [1437, 1437, 50140, 50118, 29])
+
+    def test_special_whitespace(self):
+        input_data = "\xa0 \xa0 \x3000 s"
+        encoded = self.tokenizer(input_data)
+        self.assertAllEqual(encoded, [50141, 50143, 12096, 579])
+
+    def test_cjk_input(self):
+        input_data = "素晴らしい！芭比Q啦～"
+        # Black formats long list by one element per line, which is bad to read.
+        expected = [36714, 20024, 21402, 37127, 27, 20024, 48945, 47918]
+        expected += [47780, 43251, 4394, 10172, 36484, 27969, 12410, 37127]
+        expected += [10965, 10674, 1864, 42393, 15722, 18164, 43251, 10809]
+        expected += [17772]
+        encoded = self.tokenizer(input_data)
+        self.assertAllEqual(encoded, expected)
+
+    def test_tokenize_with_tf_data(self):
+        data = [
+            "I am just a test string",
+            "I am also a test string",
+            "I am still a test string",
+            "me too",
+            "I am not a test string (joking)",
+            "You guys should add punctuation!",
+            "Period matters!",
+        ]
+        ds = tf.data.Dataset.from_tensor_slices(data)
+        ds = ds.batch(2).map(self.tokenizer)
+        encoded = next(iter(ds))
+        expected = tf.ragged.constant(
+            [[100, 524, 95, 10, 1296, 6755], [100, 524, 67, 10, 1296, 6755]]
+        )
+        self.assertAllEqual(encoded, expected)
+
+    def test_config(self):
+        input_data = ["the quick brown whale."]
+        cloned_tokenizer = BytePairTokenizer.from_config(
+            self.tokenizer.get_config()
+        )
+        self.assertAllEqual(
+            self.tokenizer(input_data),
+            cloned_tokenizer(input_data),
+        )
+
+    @parameterized.named_parameters(("tf_format", "tf"), ("h5_format", "h5"))
+    def test_saving(self, format):
+        input_data = tf.constant(["the quick brown whale."])
+        tokenizer = self.tokenizer
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = tokenizer(inputs)
+        model = keras.Model(inputs, outputs)
+        path = os.path.join(self.get_temp_dir(), "model")
+        model.save(path, save_format=format)
+        restored_model = keras.models.load_model(path)
+        self.assertAllEqual(
+            model(input_data),
+            restored_model(input_data),
+        )