Repository: keras-team/keras-nlp
Pull request #342 (2022-09-06 06:16:34)

Text Generation Functions: Add Benchmark Script
Resolves https://github.com/keras-team/keras-nlp/issues/335

Tested it here: https://colab.research.google.com/drive/1eaYRSGYY6qu6KhfUJHNVd22TUol_WkPK?usp=sharing

abheesht17: > This looks great! Just one last comment. Also, can you rebase on the latest changes (we recently fixed our testing).

Done!
mattdangerw: This is good to go, we are just waiting for https://github.com/keras-team/keras-nlp/pull/341 to land, so we can land this PR with beam search.
mattdangerw: @abheesht17 looks like this needs a reformat!
abheesht17: > @abheesht17 looks like this needs a reformat!

Yep, give me 15 minutes, I'm running the beam search test run, will paste the output in the README
abheesht17: @mattdangerw, this PR ready to be merged.

Patch:
diff --git a/keras_nlp/benchmarks/README.md b/keras_nlp/benchmarks/README.md
new file mode 100644
index 00000000..40ffa59d
--- /dev/null
+++ b/keras_nlp/benchmarks/README.md
@@ -0,0 +1,27 @@
+# KerasNLP Benchmarks
+
+This directory houses a collection of scripts for benchmarking APIs and utility
+functions which KerasNLP provides.
+
+## Text Generation
+For benchmarking text generation functions, the following command can be run
+from the root of the repository:
+
+```sh
+python3 ./keras_nlp/benchmarks/text_generation.py
+```
+
+On running this script on Google Colab (with Tesla T4 GPU, and TensorFlow 2.10.0),
+the following results were obtained:
+
+| **Decoding Strategy** 	| **Graph Mode (sec)** 	| **Graph Mode with XLA (sec)** 	|
+|:---------------------:	|:--------------------:	|:-----------------------------:	|
+|     Greedy Search     	|        495.78        	|             293.77            	|
+|      Beam Search      	|        564.23        	|             615.17            	|
+|     Random Search     	|        446.55        	|             296.21            	|
+|      Top-k Search     	|        458.68        	|             302.66            	|
+|      Top-p Search     	|        468.63        	|             565.50             	|
+
+To change the configuration, say, for example, number of layers in the transformer
+model used for inference, the user can modify the config dictionaries given at
+the top of the script.
diff --git a/keras_nlp/benchmarks/__init__.py b/keras_nlp/benchmarks/__init__.py
new file mode 100644
index 00000000..6e4df4e7
--- /dev/null
+++ b/keras_nlp/benchmarks/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/keras_nlp/benchmarks/text_generation.py b/keras_nlp/benchmarks/text_generation.py
new file mode 100644
index 00000000..59bbd182
--- /dev/null
+++ b/keras_nlp/benchmarks/text_generation.py
@@ -0,0 +1,218 @@
+# 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.
+
+"""Benchmark for text generation."""
+
+import time
+
+import tensorflow as tf
+from tensorflow import keras
+
+import keras_nlp
+from keras_nlp.utils import beam_search
+from keras_nlp.utils import greedy_search
+from keras_nlp.utils import random_search
+from keras_nlp.utils import top_k_search
+from keras_nlp.utils import top_p_search
+
+SEED = 42
+
+DATASET_ARGS = {
+    "vocab_size": 40000,
+    "num_samples": 1000,
+    "batch_size": 2,
+}
+
+TEXT_GEN_ARGS = {
+    "max_length": 64,
+    "end_token_id": 2,
+    "pad_token_id": 0,
+}
+
+MODEL_ARGS = {
+    "max_length": 300,
+    "embed_dim": 768,
+    "num_layers": 8,
+    "num_heads": 8,
+    "ff_dim": 3072,
+}
+
+TEST_RUNS = [
+    {
+        "decoding_fn": greedy_search,
+        "execution_methods": ["xla", "graph"],
+        "args": TEXT_GEN_ARGS,
+    },
+    {
+        "decoding_fn": beam_search,
+        "execution_methods": ["xla", "graph"],
+        "args": {
+            "num_beams": 2,
+            "from_logits": True,
+            **TEXT_GEN_ARGS,
+        },
+    },
+    {
+        "decoding_fn": random_search,
+        "execution_methods": ["xla", "graph"],
+        "args": {
+            "seed": SEED,
+            "from_logits": True,
+            **TEXT_GEN_ARGS,
+        },
+    },
+    {
+        "decoding_fn": top_k_search,
+        "execution_methods": ["xla", "graph"],
+        "args": {
+            "k": 5,
+            "seed": SEED,
+            "from_logits": True,
+            **TEXT_GEN_ARGS,
+        },
+    },
+    {
+        "decoding_fn": top_p_search,
+        "execution_methods": ["xla", "graph"],
+        "args": {
+            "p": 0.9,
+            "seed": SEED,
+            "from_logits": True,
+            **TEXT_GEN_ARGS,
+        },
+    },
+]
+
+
+def generate_random_ds(vocab_size, num_samples, batch_size, seed):
+    prompt_length = 2
+    inputs = tf.random.uniform(
+        shape=(num_samples, prompt_length),
+        minval=0,
+        maxval=vocab_size - 1,
+        dtype=tf.dtypes.int32,
+        seed=seed,
+    )
+
+    ds = tf.data.Dataset.from_tensor_slices(inputs)
+    ds = ds.batch(batch_size)
+    return ds
+
+
+def build_model(
+    vocab_size, max_length, embed_dim, num_layers, num_heads, ff_dim
+):
+    inputs = keras.layers.Input(shape=(None,), dtype=tf.int32)
+    # Embedding.
+    x = keras_nlp.layers.TokenAndPositionEmbedding(
+        vocabulary_size=vocab_size,
+        sequence_length=max_length,
+        embedding_dim=embed_dim,
+        mask_zero=True,
+    )(inputs)
+    # Transformer decoders.
+    for _ in range(num_layers):
+        x = keras_nlp.layers.TransformerDecoder(
+            num_heads=num_heads,
+            intermediate_dim=ff_dim,
+        )(x)
+    # Output.
+    outputs = keras.layers.Dense(vocab_size)(x)
+    model = keras.Model(inputs=inputs, outputs=outputs)
+    return model
+
+
+def generate_text(
+    decoding_fn,
+    token_probability_fn,
+    prompt,
+    text_gen_args,
+    jit_compile,
+):
+    class TestModel(tf.keras.Model):
+        def call(self, inputs):
+            generated = decoding_fn(
+                token_probability_fn=token_probability_fn,
+                prompt=inputs,
+                **text_gen_args,
+            )
+            return generated
+
+    test_model = TestModel()
+    test_model.compile(jit_compile=jit_compile)
+
+    t0 = time.time()
+    _ = test_model.predict(prompt)
+    return time.time() - t0
+
+
+def main():
+    keras.utils.set_random_seed(SEED)
+    csv_path = time.strftime("text_gen_%Y-%m-%d_%H-%M-%S.csv")
+
+    ds = generate_random_ds(
+        vocab_size=DATASET_ARGS["vocab_size"],
+        num_samples=DATASET_ARGS["num_samples"],
+        batch_size=DATASET_ARGS["batch_size"],
+        seed=SEED,
+    )
+
+    model = build_model(
+        vocab_size=DATASET_ARGS["vocab_size"],
+        max_length=MODEL_ARGS["max_length"],
+        embed_dim=MODEL_ARGS["embed_dim"],
+        num_layers=MODEL_ARGS["num_layers"],
+        num_heads=MODEL_ARGS["num_heads"],
+        ff_dim=MODEL_ARGS["ff_dim"],
+    )
+
+    def token_logits_fn(inputs):
+        output = model(inputs)
+        return output[:, -1, :]
+
+    print("*************************************\n")
+
+    with open(csv_path, "w") as res_handler:
+        res_handler.write("decoding_strategy,execution_method,time\n")
+        for test_run in TEST_RUNS:
+            decoding_fn = test_run["decoding_fn"]
+            decoding_strategy = decoding_fn.__name__
+
+            for execution_method in test_run["execution_methods"]:
+                print(f"Running {decoding_strategy} in {execution_method} mode")
+
+                if execution_method == "graph":
+                    jit_compile = False
+                elif execution_method == "xla":
+                    jit_compile = True
+
+                time_taken = generate_text(
+                    decoding_fn=decoding_fn,
+                    token_probability_fn=token_logits_fn,
+                    prompt=ds,
+                    text_gen_args=test_run["args"],
+                    jit_compile=jit_compile,
+                )
+                print("Time taken: ", time_taken)
+                res_handler.write(
+                    f"{decoding_strategy},{execution_method}," f"{time_taken}\n"
+                )
+                print()
+            print("*************************************")
+
+    print(f"Writing results to {csv_path}")
+
+
+if __name__ == "__main__":
+    main()
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
     )
 
diff --git a/keras_nlp/utils/text_generation.py b/keras_nlp/utils/text_generation.py
index aaabe632..0e07806c 100644
--- a/keras_nlp/utils/text_generation.py
+++ b/keras_nlp/utils/text_generation.py
@@ -802,7 +802,7 @@ def top_p_search(
         probs = tf.where(
             shifted_keep_mask,
             sorted_preds,
-            tf.zeros(pred.shape, dtype=sorted_preds.dtype),
+            tf.zeros(tf.shape(pred), dtype=sorted_preds.dtype),
         )
         sorted_next_token = tf.random.categorical(
             tf.math.log(probs), 1, seed=seed


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.")