Repository: keras-team/keras-nlp
Pull request #343 (2022-09-06 17:24:13)

Add a BertPreprocessor class
This exposes a tokenizer for Bert, including a way to access pretrained vocabularies, as a standalone class. I am fairly confident we will want access to this standalone (though it may be a power user tool). You might want to prepreprocess your inputs outside of the training script, and, in that case, we should not require users to load a massive tf graph.

I think one of the major design considerations here is wether we want add a parameters for the size of the model we want the vocab for (e.g. `"base"` or `"large"`). For bert there is no difference for pretrained vocabularies across sizes. This is also true for gpt-2, roberta, and xlm-roberta. I think for most model tokenizers we can avoid a "model size" parameter, which possibly some exceptions (though I have no concrete examples).

mattdangerw: > What if we scoped this method down to just loading pretrained tokenizers and left advanced users to roll their own?

I think there are two question here. Custom model files, and unpacked input.

**Re custom model files**

I was under the impression we would eventually let users supply custom weight files in https://github.com/keras-team/keras-nlp/issues/314, so it seemed consistent to do the same here. I could see someone copying the vocab to a different cloud/local storage solution, but want the exact same logic as our gcp hosted pretrained tokenizer.

**Re unpacked input*

I'm ok to leave that out if we think it's unimportant (doesn't rule out adding it later).  Here's my pitch to keep it though...

1) I wanted a way to get the pretrained core elements for BERT, with as little assumptions as possible. Felt important for our "modular" identity or something. There should be a way to get just the bert tokenizer without packing. To have the recommend solution to that to be to look into the code and reverse engineer our gcp vocab urls felt slightly user hostile.
2) This seems useful just for demoing purposes! I could see a guide that looks like the following:

```python
# Bert uses a sub-word tokenizer for inputs.
# We will turn off the packing step for now to look at how bert tokenizes words.
tokenizer = keras_nlp.models.BertTokenizer("uncased_en", pack_inputs=False)
tokenizer("Some test input with extravagantly interminable words") -> ....

# Before feeding into a model, bert will need to "pack" inputs into a sequence of fixed length.
# We will also need to compute a mask and segment sequence blah blah...
# Let's take a look! 
tokenizer = keras_nlp.models.BertTokenizer("uncased_en", pack_inputs=True)
tokenizer("Some test input with extravagantly interminable words") -> ....
```
jbischof: Thanks for the clarifications @mattdangerw! This may warrant an offline discussion but I'm offering a different interpretation of our package identity inspired by "Simple use cases should be simple, advanced use cases should be possible". I think both #314 and the packing flag are examples are "framework" thinking where a monolithic API satisfies all use cases. Instead we want to offer the raw modules that can be composed to accomplish anything and a high level API to make the common use cases easy and light on boilerplate. 

If we want an example without packing, we should invoke the tokenizer directly and not provide a single API that provides every combination of functionality conceivable. #314 was requested by @fchollet to be consistent with a historical choice and now it sounds like a wedge to turn ourselves into a framework. 
mattdangerw: OK! After talking offline with Jonathan, I have reworked this to be a  `Layer` and not a `Tokenizer`, which wraps the tokenizer and packer into a single offering. This has the advantage of removing the `pack_inputs` argument, while still allowing direct access to the lower-level WordPiece tokenizer via `layer.tokenizer`. Makes the code very nice.

It does however, land us in an interesting naming conversion. It was a bit of a slight of hand to fold everything into the tokenizer interface, but now we have something explicitly is **not** a tokenizer class, it **has** a tokenizer. I have tentatively named this `BertPreprocessor`, but interested if there are other options we should explore.

 - `BertPreprocessor`: too general? what is preprocessing?
 - `BertTokenizer`: very inaccurate, this has a property called `tokenizer` and does not follow the interface.
 - `BertTokenizerAndPacker`: top marks for accuracy, but doesn't quite roll off the tongue.

Accepting suggestions :), also worth thinking about the idiomatic usage. E.g.

```
# This reads poorly.
preprocessor = BertPreprocessor()

def preprocess(x, y):
  preprocessor(x), y

ds = ds.map(preprocess)
```

@fchollet @jbischof 
mattdangerw: Thanks @jbischof!

Re `BertTokenizer` I feel like that would be a mistake with the current design. [Tokenizer](https://keras.io/api/keras_nlp/tokenizers/tokenizer/) is an interface, it means something specific in the library, (e.g. has `tokenize()`, `detokenize()`, `vocubulary_size()`, etc.).

If we want that, I think we should flip this back to a design where this actually extends the `tokenizer` interface. That might come with more complexity, but we preserve the ability to say, everything called a "tokenizer" in our library exposes the same consistent set of methods.

So probably just `BertPreprocessor` for now?
mattdangerw: @jbischof with this approach, we would also need to figure out how access the vocabulary size. That's something we should probably document, as it's a required arg on on `BertBase` when loading your own weights. Today that would be as follows.

```
preprocessor = keras_nlp.model.BertPreprocessor(
    vocabulary="uncase_en",
)
model = keras_nlp.models.BertBase(
    vocabulary_size=preprocessor.tokenizer.vocabulary_size(),
)
```

We could also just mirror `vocabulary_size()` on the layer, but then we start to get in a weird game of match the interface with actual the actual `Tokenizer` class. WDYT?
jbischof: For the case of pretrained checkpoints I think this should be handled automatically since the checkpoint should know its vocabulary. Left a comment in the code for that.

For the case of using a vocabulary with new model (including one trained by a user), the best experience will come with using a `Pipeline` as in #346.

Finally, for the case of users who want to control preprocessing directly, what if we made `vocabulary_size` part of the config of the preprocessor? Since the `vocabulary` is part of the `__init__` signature this trait should be immutable.
mattdangerw: @jbischof thanks! I'll just pull `vocabulary_size()` onto the layer. Want to avoid `preprocessor.vocabulary_size` vs `tokenizer.vocabulary_size()` inconsistencies in the library.

Patch:
diff --git a/keras_nlp/layers/multi_segment_packer.py b/keras_nlp/layers/multi_segment_packer.py
index c879f2d2..56ff2e15 100644
--- a/keras_nlp/layers/multi_segment_packer.py
+++ b/keras_nlp/layers/multi_segment_packer.py
@@ -30,7 +30,7 @@ class MultiSegmentPacker(keras.layers.Layer):
     Takes as input a list or tuple of token segments. The layer will process
     inputs as follows:
      - Truncate all input segments to fit within `sequence_length` according to
-       the `truncator` strategy.
+       the `truncate` strategy.
      - Concatenate all input segments, adding a single `start_value` at the
        start of the entire sequence, and multiple `end_value`s at the end of
        each segment.
@@ -55,7 +55,7 @@ class MultiSegmentPacker(keras.layers.Layer):
         pad_value: The id or token that is to be placed into the unused
             positions after the last segment in the sequence
             (called "[PAD]" for BERT).
-        truncator: The algorithm to truncate a list of batched segments to fit a
+        truncate: The algorithm to truncate a list of batched segments to fit a
             per-example length limit. The value can be either `round_robin` or
             `waterfall`:
                 - `"round_robin"`: Available space is assigned one token at a
@@ -104,17 +104,17 @@ class MultiSegmentPacker(keras.layers.Layer):
         start_value,
         end_value,
         pad_value=None,
-        truncator="round_robin",
+        truncate="round_robin",
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.sequence_length = sequence_length
-        if truncator not in ("round_robin", "waterfall"):
+        if truncate not in ("round_robin", "waterfall"):
             raise ValueError(
                 "Only 'round_robin' and 'waterfall' algorithms are "
-                "supported. Received %s" % truncator
+                "supported. Received %s" % truncate
             )
-        self.truncator = truncator
+        self.truncate = truncate
         self.start_value = start_value
         self.end_value = end_value
         self.pad_value = pad_value
@@ -127,7 +127,7 @@ class MultiSegmentPacker(keras.layers.Layer):
                 "start_value": self.start_value,
                 "end_value": self.end_value,
                 "pad_value": self.pad_value,
-                "truncator": self.truncator,
+                "truncate": self.truncate,
             }
         )
         return config
@@ -162,16 +162,16 @@ class MultiSegmentPacker(keras.layers.Layer):
     def _trim_inputs(self, inputs):
         """Trim inputs to desired length."""
         num_special_tokens = len(inputs) + 1
-        if self.truncator == "round_robin":
+        if self.truncate == "round_robin":
             return tf_text.RoundRobinTrimmer(
                 self.sequence_length - num_special_tokens
             ).trim(inputs)
-        elif self.truncator == "waterfall":
+        elif self.truncate == "waterfall":
             return tf_text.WaterfallTrimmer(
                 self.sequence_length - num_special_tokens
             ).trim(inputs)
         else:
-            raise ValueError("Unsupported truncator: %s" % self.truncator)
+            raise ValueError("Unsupported truncate: %s" % self.truncate)
 
     def _combine_inputs(self, segments):
         """Combine inputs with start and end values added."""
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index 10ee6a82..39cbf53a 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -17,6 +17,7 @@ from keras_nlp.models.bert import BertClassifier
 from keras_nlp.models.bert import BertCustom
 from keras_nlp.models.bert import BertLarge
 from keras_nlp.models.bert import BertMedium
+from keras_nlp.models.bert import BertPreprocessor
 from keras_nlp.models.bert import BertSmall
 from keras_nlp.models.bert import BertTiny
 from keras_nlp.models.roberta import RobertaBase
diff --git a/keras_nlp/models/bert.py b/keras_nlp/models/bert.py
index 6b9db9a5..b4f1a1ee 100644
--- a/keras_nlp/models/bert.py
+++ b/keras_nlp/models/bert.py
@@ -14,128 +14,213 @@
 
 """BERT model configurable class, preconfigured versions, and task heads."""
 
+import os
+
 import tensorflow as tf
 from tensorflow import keras
 
-from keras_nlp.layers import PositionEmbedding
-from keras_nlp.layers import TransformerEncoder
+from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
+from keras_nlp.layers.position_embedding import PositionEmbedding
+from keras_nlp.layers.transformer_encoder import TransformerEncoder
+from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
 
 
 def _bert_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
-def _handle_weights_and_vocab_size(bert_variant, weights, vocabulary_size):
-    """Look up pretrained defaults for `weights` and `vocabulary_size`.
-
-    This helper will validate the `weights` and `vocabulary_size` arguments, and
-    fully resolve them in the case we are loading pretrained weights.
-    """
-    if (vocabulary_size is None and weights is None) or (
-        vocabulary_size and weights
-    ):
-        raise ValueError(
-            "One of `vocabulary_size` or `weights` must be specified "
-            "(but not both). "
-            f"Received: weights={weights}, "
-            f"vocabulary_size={vocabulary_size}"
-        )
-
-    if weights:
-        if weights not in checkpoints[bert_variant]:
-            raise ValueError(
-                "`weights` must be one of "
-                f"""{", ".join(checkpoints[bert_variant])}. """
-                f"Received: {weights}"
-            )
-
-        vocabulary_size = checkpoints[bert_variant][weights]["vocabulary_size"]
-
-        # TODO(jbischof): consider changing format from `h5` to
-        # `tf.train.Checkpoint` once
-        # https://github.com/keras-team/keras/issues/16946 is resolved.
-        weights = keras.utils.get_file(
-            "model.h5",
-            BASE_PATH + f"{bert_variant}_{weights}/model.h5",
-            cache_subdir=f"models/{bert_variant}/{weights}/",
-            file_hash=checkpoints[bert_variant][weights]["md5"],
-        )
-
-    return weights, vocabulary_size
-
-
-# Pretrained models
-BASE_PATH = "https://storage.googleapis.com/keras-nlp/models/"
-
+# Metadata for loading pretrained model weights.
 checkpoints = {
     "bert_tiny": {
         "uncased_en": {
-            "md5": "c2b29fcbf8f814a0812e4ab89ef5c068",
-            "description": "Tiny size of BERT where all input is lowercased. "
-            "Trained on English Wikipedia + BooksCorpus.",
-            "vocabulary_size": 30522,
+            "description": (
+                "Tiny size of BERT where all input is lowercased. "
+                "Trained on English Wikipedia + BooksCorpus."
+            ),
+            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_tiny_uncased_en/model.h5",
+            "weights_hash": "c2b29fcbf8f814a0812e4ab89ef5c068",
         }
     },
     "bert_small": {
         "uncased_en": {
-            "md5": "08632c9479b034f342ba2c2b7afba5f7",
-            "description": "Small size of BERT where all input is lowercased. "
-            "Trained on English Wikipedia + BooksCorpus.",
-            "vocabulary_size": 30522,
+            "description": (
+                "Small size of BERT where all input is lowercased. "
+                "Trained on English Wikipedia + BooksCorpus."
+            ),
+            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_small_uncased_en/model.h5",
+            "weights_hash": "08632c9479b034f342ba2c2b7afba5f7",
         }
     },
     "bert_medium": {
         "uncased_en": {
-            "md5": "bb990e1184ec6b6185450c73833cd661",
-            "description": "Medium size of BERT where all input is lowercased. "
-            "Trained on English Wikipedia + BooksCorpus.",
-            "vocabulary_size": 30522,
+            "description": (
+                "Medium size of BERT where all input is lowercased. "
+                "Trained on English Wikipedia + BooksCorpus."
+            ),
+            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_medium_uncased_en/model.h5",
+            "weights_hash": "bb990e1184ec6b6185450c73833cd661",
         }
     },
     "bert_base": {
         "uncased_en": {
-            "md5": "9b2b2139f221988759ac9cdd17050b31",
-            "description": "Base size of BERT where all input is lowercased. "
-            "Trained on English Wikipedia + BooksCorpus.",
-            "vocabulary_size": 30522,
+            "description": (
+                "Base size of Bert where all input is lowercased. "
+                "Trained on English Wikipedia + BooksCorpus."
+            ),
+            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_uncased_en/model.h5",
+            "weights_hash": "9b2b2139f221988759ac9cdd17050b31",
         },
         "cased_en": {
-            "md5": "f94a6cb012e18f4fb8ec92abb91864e9",
-            "description": "Base size of BERT where case is maintained. "
-            "Trained on English Wikipedia + BooksCorpus.",
-            "vocabulary_size": 28996,
+            "description": (
+                "Base size of Bert where case is maintained. "
+                "Trained on English Wikipedia + BooksCorpus."
+            ),
+            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_cased_en/model.h5",
+            "weights_hash": "f94a6cb012e18f4fb8ec92abb91864e9",
         },
         "zh": {
-            "md5": "79afa421e386076e62ab42dad555ab0c",
-            "description": "Base size of BERT. Trained on Chinese Wikipedia.",
-            "vocabulary_size": 21128,
+            "description": ("Base size of BERT. Trained on Chinese Wikipedia."),
+            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_zh/model.h5",
+            "weights_hash": "79afa421e386076e62ab42dad555ab0c",
         },
         "multi_cased": {
-            "md5": "b0631cec0a1f2513c6cfd75ba29c33aa",
-            "description": "Base size of BERT. Trained on Wikipedias of "
-            "104 languages.",
-            "vocabulary_size": 119547,
+            "description": (
+                "Base size of BERT. Trained on Wikipedias of 104 languages."
+            ),
+            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_multi_cased/model.h5",
+            "weights_hash": "b0631cec0a1f2513c6cfd75ba29c33aa",
         },
     },
     "bert_large": {
         "uncased_en": {
-            "md5": "cc5cacc9565ef400ee4376105f40ddae",
-            "description": "Large size of BERT where all input is lowercased. "
-            "Trained on English Wikipedia + BooksCorpus.",
-            "vocabulary_size": 30522,
+            "description": (
+                "Large size of BERT where all input is lowercased. "
+                "Trained on English Wikipedia + BooksCorpus."
+            ),
+            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_large_uncased_en/model.h5",
+            "weights_hash": "cc5cacc9565ef400ee4376105f40ddae",
         },
         "cased_en": {
-            "md5": "8b8ab82290bbf4f8db87d4f100648890",
-            "description": "Large size of BERT where case is maintained. "
-            "Trained on English Wikipedia + BooksCorpus.",
-            "vocabulary_size": 28996,
+            "description": (
+                "Large size of BERT where case is maintained. "
+                "Trained on English Wikipedia + BooksCorpus."
+            ),
+            "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_large_cased_en/model.h5",
+            "weights_hash": "8b8ab82290bbf4f8db87d4f100648890",
         },
     },
 }
 
 
+# Metadata for loading pretrained tokenizer vocabularies.
+# We need the vocabulary_size hardcoded so we can instantiate a Bert network
+# with the right embedding size without downloading the matching vocabulary.
+# TODO(mattdangerw): Update our bucket structure so the vocabularies are
+# stored in an independent way, rather than reading from the base model.
+vocabularies = {
+    "uncased_en": {
+        "description": (
+            "The vocabulary for BERT models trained on "
+            "English Wikipedia + BooksCorpus where case is discarded."
+        ),
+        "vocabulary_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_uncased_en/vocab.txt",
+        "vocabulary_hash": "64800d5d8528ce344256daf115d4965e",
+        "vocabulary_size": 30522,
+        "lowercase": True,
+    },
+    "cased_en": {
+        "description": (
+            "The vocabulary for BERT models trained on "
+            "English Wikipedia + BooksCorpus where case is maintained."
+        ),
+        "vocabulary_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_cased_en/vocab.txt",
+        "vocabulary_hash": "bb6ca9b42e790e5cd986bbb16444d0e0",
+        "vocabulary_size": 28996,
+        "lowercase": False,
+    },
+    "zh": {
+        "description": (
+            "The vocabulary for BERT models trained on Chinese Wikipedia."
+        ),
+        "vocabulary_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_zh/vocab.txt",
+        "vocabulary_hash": "3b5b76c4aef48ecf8cb3abaafe960f09",
+        "vocabulary_size": 21128,
+        "lowercase": False,
+    },
+    "multi_cased": {
+        "description": (
+            "The vocabulary for BERT models trained on trained on Wikipedias "
+            "of 104 languages."
+        ),
+        "vocabulary_url": "https://storage.googleapis.com/keras-nlp/models/bert_base_multi_cased/vocab.txt",
+        "vocabulary_hash": "d9d865138d17f1958502ed060ecfeeb6",
+        "vocabulary_size": 119547,
+        "lowercase": False,
+    },
+}
+
+
+def _handle_pretrained_tokenizer_arguments(vocabulary, lowercase):
+    """Look up pretrained defaults for tokenizer arguments.
+
+    This helper will validate the `vocabulary` and `lowercase` arguments, and
+    fully resolve them in the case we are loading pretrained weights.
+    """
+
+    if isinstance(vocabulary, str) and vocabulary in vocabularies:
+        metadata = vocabularies[vocabulary]
+        vocabulary = keras.utils.get_file(
+            "vocab.txt",
+            metadata["vocabulary_url"],
+            cache_subdir=os.path.join("models", "bert", vocabulary),
+            file_hash=metadata["vocabulary_hash"],
+        )
+        lowercase = metadata["lowercase"]
+
+    return vocabulary, lowercase
+
+
+def _handle_pretrained_model_arguments(bert_variant, weights, vocabulary_size):
+    """Look up pretrained defaults for model arguments.
+
+    This helper will validate the `weights` and `vocabulary_size` arguments, and
+    fully resolve them in the case we are loading pretrained weights.
+    """
+    if (vocabulary_size is None and weights is None) or (
+        vocabulary_size and weights
+    ):
+        raise ValueError(
+            "One of `vocabulary_size` or `weights` must be specified "
+            "(but not both). "
+            f"Received: weights={weights}, "
+            f"vocabulary_size={vocabulary_size}"
+        )
+
+    if weights:
+        if weights not in checkpoints[bert_variant]:
+            raise ValueError(
+                "`weights` must be one of "
+                f"""{", ".join(checkpoints[bert_variant])}. """
+                f"Received: {weights}"
+            )
+        metadata = checkpoints[bert_variant][weights]
+        vocabulary_size = vocabularies[weights]["vocabulary_size"]
+
+        # TODO(jbischof): consider changing format from `h5` to
+        # `tf.train.Checkpoint` once
+        # https://github.com/keras-team/keras/issues/16946 is resolved.
+        weights = keras.utils.get_file(
+            "model.h5",
+            metadata["weights_url"],
+            cache_subdir=os.path.join("models", "bert", weights, bert_variant),
+            file_hash=metadata["weights_hash"],
+        )
+
+    return weights, vocabulary_size
+
+
 class BertCustom(keras.Model):
-    """Bi-directional Transformer-based encoder network.
+    """BERT encoder network with custom hyperparmeters.
 
     This network implements a bi-directional Transformer-based encoder as
     described in ["BERT: Pre-training of Deep Bidirectional Transformers for
@@ -144,7 +229,7 @@ class BertCustom(keras.Model):
     or classification task networks.
 
     This class gives a fully customizable BERT model with any number of layers,
-    heads, and embedding dimensions. For specific specific bert architectures
+    heads, and embedding dimensions. For specific specific BERT architectures
     defined in the paper, see for example `keras_nlp.models.BertBase`.
 
     Args:
@@ -166,7 +251,7 @@ class BertCustom(keras.Model):
         trainable: Boolean, optional. If the model's variables should be
             trainable.
 
-    Example usage:
+    Examples:
     ```python
     # Randomly initialized BERT encoder
     model = keras_nlp.models.BertCustom(
@@ -328,6 +413,154 @@ class BertCustom(keras.Model):
         return config
 
 
+PREPROCESSOR_DOCSTRING = """BERT 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"`, `"segment_ids"`,
+   `"padding_mask"`, that can be passed directly to a BERT 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. For a pretrained vocabulary,
+        `vocabulary` should be one of {names}, and should match the `weights`
+        parameter of any pretrained BERT model.
+    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.BertPreprocessor(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,
+)
+```
+"""
+
+
+class BertPreprocessor(keras.layers.Layer):
+    def __init__(
+        self,
+        vocabulary="uncased_en",
+        lowercase=False,
+        sequence_length=512,
+        truncate="round_robin",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        vocabulary, lowercase = _handle_pretrained_tokenizer_arguments(
+            vocabulary, lowercase
+        )
+
+        self.tokenizer = WordPieceTokenizer(
+            vocabulary=vocabulary,
+            lowercase=lowercase,
+        )
+
+        # Check for necessary special tokens.
+        cls_token = "[CLS]"
+        sep_token = "[SEP]"
+        pad_token = "[PAD]"
+        for token in [cls_token, pad_token, sep_token]:
+            if token not in self.tokenizer.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.packer = MultiSegmentPacker(
+            start_value=self.tokenizer.token_to_id(cls_token),
+            end_value=self.tokenizer.token_to_id(sep_token),
+            pad_value=self.tokenizer.token_to_id(pad_token),
+            truncate=truncate,
+            sequence_length=sequence_length,
+        )
+
+    def vocabulary_size(self):
+        """Returns the vocabulary size of the tokenizer."""
+        return self.tokenizer.vocabulary_size()
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary": self.tokenizer.vocabulary,
+                "lowercase": self.tokenizer.lowercase,
+                "sequence_length": self.packer.sequence_length,
+                "trucator": self.packer.trucator,
+            }
+        )
+        return config
+
+    def call(self, inputs):
+        if not isinstance(inputs, (list, tuple)):
+            inputs = [inputs]
+
+        inputs = [self.tokenizer(x) for x in inputs]
+        token_ids, segment_ids = self.packer(inputs)
+        return {
+            "token_ids": token_ids,
+            "segment_ids": segment_ids,
+            "padding_mask": token_ids != 0,
+        }
+
+
+setattr(
+    BertPreprocessor,
+    "__doc__",
+    PREPROCESSOR_DOCSTRING.format(names=", ".join(vocabularies)),
+)
+
+
 class BertClassifier(keras.Model):
     """BERT encoder model with a classification head.
 
@@ -338,7 +571,7 @@ class BertClassifier(keras.Model):
         trainable: Boolean, optional. If the model's variables should be
             trainable.
 
-    Example usage:
+    Examples:
     ```python
     # Randomly initialized BERT encoder
     model = keras_nlp.models.BertCustom(
@@ -390,8 +623,7 @@ class BertClassifier(keras.Model):
         self.num_classes = num_classes
 
 
-MODEL_DOCSTRING = """Bi-directional Transformer-based encoder network (BERT)
-    using "{type}" architecture.
+MODEL_DOCSTRING = """Bert "{type}" architecture.
 
     This network implements a bi-directional Transformer-based encoder as
     described in ["BERT: Pre-training of Deep Bidirectional Transformers for
@@ -410,7 +642,7 @@ MODEL_DOCSTRING = """Bi-directional Transformer-based encoder network (BERT)
         trainable: Boolean, optional. If the model's variables should be
             trainable.
 
-    Example usage:
+    Examples:
     ```python
     # Randomly initialized Bert{type} encoder
     model = keras_nlp.models.Bert{type}(vocabulary_size=10000)
@@ -434,7 +666,7 @@ MODEL_DOCSTRING = """Bi-directional Transformer-based encoder network (BERT)
 
 
 def BertTiny(weights=None, vocabulary_size=None, name=None, trainable=True):
-    weights, vocabulary_size = _handle_weights_and_vocab_size(
+    weights, vocabulary_size = _handle_pretrained_model_arguments(
         "bert_tiny", weights, vocabulary_size
     )
 
@@ -459,7 +691,7 @@ def BertTiny(weights=None, vocabulary_size=None, name=None, trainable=True):
 
 
 def BertSmall(weights=None, vocabulary_size=None, name=None, trainable=True):
-    weights, vocabulary_size = _handle_weights_and_vocab_size(
+    weights, vocabulary_size = _handle_pretrained_model_arguments(
         "bert_small", weights, vocabulary_size
     )
 
@@ -482,7 +714,7 @@ def BertSmall(weights=None, vocabulary_size=None, name=None, trainable=True):
 
 
 def BertMedium(weights=None, vocabulary_size=None, name=None, trainable=True):
-    weights, vocabulary_size = _handle_weights_and_vocab_size(
+    weights, vocabulary_size = _handle_pretrained_model_arguments(
         "bert_medium", weights, vocabulary_size
     )
 
@@ -505,7 +737,7 @@ def BertMedium(weights=None, vocabulary_size=None, name=None, trainable=True):
 
 
 def BertBase(weights=None, vocabulary_size=None, name=None, trainable=True):
-    weights, vocabulary_size = _handle_weights_and_vocab_size(
+    weights, vocabulary_size = _handle_pretrained_model_arguments(
         "bert_base", weights, vocabulary_size
     )
 
@@ -528,7 +760,7 @@ def BertBase(weights=None, vocabulary_size=None, name=None, trainable=True):
 
 
 def BertLarge(weights=None, vocabulary_size=None, name=None, trainable=True):
-    weights, vocabulary_size = _handle_weights_and_vocab_size(
+    weights, vocabulary_size = _handle_pretrained_model_arguments(
         "bert_large", weights, vocabulary_size
     )
 


Test patch:
diff --git a/keras_nlp/layers/multi_segment_packer_test.py b/keras_nlp/layers/multi_segment_packer_test.py
index 9a00d86b..dda505b3 100644
--- a/keras_nlp/layers/multi_segment_packer_test.py
+++ b/keras_nlp/layers/multi_segment_packer_test.py
@@ -43,7 +43,7 @@ class MultiSegmentPackerTest(tf.test.TestCase, parameterized.TestCase):
         seq1 = tf.constant(["a", "b", "c"])
         seq2 = tf.constant(["x", "y", "z"])
         packer = MultiSegmentPacker(
-            7, start_value="[CLS]", end_value="[SEP]", truncator="round_robin"
+            7, start_value="[CLS]", end_value="[SEP]", truncate="round_robin"
         )
         output = packer([seq1, seq2])
         self.assertAllEqual(
@@ -58,7 +58,7 @@ class MultiSegmentPackerTest(tf.test.TestCase, parameterized.TestCase):
         seq1 = tf.constant(["a", "b", "c"])
         seq2 = tf.constant(["x", "y", "z"])
         packer = MultiSegmentPacker(
-            7, start_value="[CLS]", end_value="[SEP]", truncator="waterfall"
+            7, start_value="[CLS]", end_value="[SEP]", truncate="waterfall"
         )
         output = packer([seq1, seq2])
         self.assertAllEqual(
@@ -73,7 +73,7 @@ class MultiSegmentPackerTest(tf.test.TestCase, parameterized.TestCase):
         seq1 = tf.constant([["a", "b", "c"], ["a", "b", "c"]])
         seq2 = tf.constant([["x", "y", "z"], ["x", "y", "z"]])
         packer = MultiSegmentPacker(
-            7, start_value="[CLS]", end_value="[SEP]", truncator="round_robin"
+            7, start_value="[CLS]", end_value="[SEP]", truncate="round_robin"
         )
         output = packer([seq1, seq2])
         self.assertAllEqual(
@@ -94,7 +94,7 @@ class MultiSegmentPackerTest(tf.test.TestCase, parameterized.TestCase):
         seq1 = tf.ragged.constant([["a", "b", "c"], ["a", "b"]])
         seq2 = tf.constant([["x", "y", "z"], ["x", "y", "z"]])
         packer = MultiSegmentPacker(
-            7, start_value="[CLS]", end_value="[SEP]", truncator="waterfall"
+            7, start_value="[CLS]", end_value="[SEP]", truncate="waterfall"
         )
         output = packer([seq1, seq2])
         self.assertAllEqual(
@@ -151,7 +151,7 @@ class MultiSegmentPackerTest(tf.test.TestCase, parameterized.TestCase):
         seq1 = tf.ragged.constant([["a", "b", "c"], ["a", "b"]])
         seq2 = tf.ragged.constant([["x", "y", "z"], ["x", "y", "z"]])
         original_packer = MultiSegmentPacker(
-            7, start_value="[CLS]", end_value="[SEP]", truncator="waterfall"
+            7, start_value="[CLS]", end_value="[SEP]", truncate="waterfall"
         )
         cloned_packer = MultiSegmentPacker.from_config(
             original_packer.get_config()
@@ -166,7 +166,7 @@ class MultiSegmentPackerTest(tf.test.TestCase, parameterized.TestCase):
         seq1 = tf.ragged.constant([["a", "b", "c"], ["a", "b"]])
         seq2 = tf.ragged.constant([["x", "y", "z"], ["x", "y", "z"]])
         packer = MultiSegmentPacker(
-            7, start_value="[CLS]", end_value="[SEP]", truncator="waterfall"
+            7, start_value="[CLS]", end_value="[SEP]", truncate="waterfall"
         )
         inputs = (
             keras.Input(dtype="string", ragged=True, shape=(None,)),
diff --git a/keras_nlp/models/bert_test.py b/keras_nlp/models/bert_test.py
index 0fd607e5..2c0694fc 100644
--- a/keras_nlp/models/bert_test.py
+++ b/keras_nlp/models/bert_test.py
@@ -22,6 +22,44 @@ from tensorflow import keras
 from keras_nlp.models import bert
 
 
+class BertPreprocessorTest(tf.test.TestCase):
+    def setUp(self):
+        self.vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
+        self.vocab += ["THE", "QUICK", "BROWN", "FOX"]
+        self.vocab += ["the", "quick", "brown", "fox"]
+
+    def test_tokenize(self):
+        input_data = ["THE QUICK BROWN FOX."]
+        preprocessor = bert.BertPreprocessor(
+            vocabulary=self.vocab,
+            sequence_length=8,
+        )
+        output = preprocessor(input_data)
+        self.assertAllEqual(output["token_ids"], [2, 5, 6, 7, 8, 1, 3, 0])
+        self.assertAllEqual(output["segment_ids"], [0, 0, 0, 0, 0, 0, 0, 0])
+        self.assertAllEqual(output["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 0])
+
+    def test_lowercase(self):
+        input_data = ["THE QUICK BROWN FOX."]
+        preprocessor = bert.BertPreprocessor(
+            vocabulary=self.vocab,
+            sequence_length=8,
+            lowercase=True,
+        )
+        output = preprocessor(input_data)
+        self.assertAllEqual(output["token_ids"], [2, 9, 10, 11, 12, 1, 3, 0])
+
+    def test_detokenize(self):
+        input_data = [[5, 6, 7, 8]]
+        preprocessor = bert.BertPreprocessor(vocabulary=self.vocab)
+        output = preprocessor.tokenizer.detokenize(input_data)
+        self.assertAllEqual(output, ["THE QUICK BROWN FOX"])
+
+    def test_vocabulary_size(self):
+        preprocessor = bert.BertPreprocessor(vocabulary=self.vocab)
+        self.assertEqual(preprocessor.vocabulary_size(), 13)
+
+
 class BertTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         self.model = bert.BertCustom(
diff --git a/keras_nlp/network_tests/test_load_ckpts.py b/keras_nlp/network_tests/test_load_ckpts.py
index 3f9a68af..796d0f66 100644
--- a/keras_nlp/network_tests/test_load_ckpts.py
+++ b/keras_nlp/network_tests/test_load_ckpts.py
@@ -43,3 +43,13 @@ class BertCkptTest(tf.test.TestCase, parameterized.TestCase):
             "padding_mask": tf.constant([1] * 512, shape=(1, 512)),
         }
         model(input_data)
+
+    @parameterized.named_parameters(
+        ("uncased_en", "uncased_en"),
+        ("cased_en", "cased_en"),
+        ("zh", "zh"),
+        ("multi_cased", "multi_cased"),
+    )
+    def test_load_vocabularies(self, vocabulary):
+        tokenizer = keras_nlp.models.BertPreprocessor(vocabulary=vocabulary)
+        tokenizer("The quick brown fox.")