Repository: keras-team/keras-nlp

Issue:
Add a model proto training utility for sentence piece
Along with https://github.com/keras-team/keras-nlp/issues/248, we should add a utility to train a model proto for sentence piece. This can leverage the sentencepiece pip package to run their trainer.

We should not expose most of the option the sentencepiece trainer has, the goal will be for this to cover the majority of use cases, with advanced users till using the sentencepiece pip package directly.

Discussion:
chenmoneygithub: @aflah02 Thanks for the PR! Could you share a colab on how to use this trainer? and btw, what's the time cost for training wikipedia-en?

Patch:
diff --git a/examples/bert/bert_config.py b/examples/bert/bert_config.py
index e5febb92..6a0b8902 100644
--- a/examples/bert/bert_config.py
+++ b/examples/bert/bert_config.py
@@ -12,60 +12,49 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# TODO(jbischof): remove in favor of BertBase, BertSmall, etc
 MODEL_CONFIGS = {
     "tiny": {
         "num_layers": 2,
-        "hidden_size": 128,
+        "hidden_dim": 128,
         "dropout": 0.1,
-        "num_attention_heads": 2,
-        "inner_size": 512,
-        "inner_activation": "gelu",
-        "initializer_range": 0.02,
+        "num_heads": 2,
+        "intermediate_dim": 512,
     },
     "mini": {
         "num_layers": 4,
-        "hidden_size": 256,
+        "hidden_dim": 256,
         "dropout": 0.1,
-        "num_attention_heads": 4,
-        "inner_size": 1024,
-        "inner_activation": "gelu",
-        "initializer_range": 0.02,
+        "num_heads": 4,
+        "intermediate_dim": 1024,
     },
     "small": {
         "num_layers": 4,
-        "hidden_size": 512,
+        "hidden_dim": 512,
         "dropout": 0.1,
-        "num_attention_heads": 8,
-        "inner_size": 2048,
-        "inner_activation": "gelu",
-        "initializer_range": 0.02,
+        "num_heads": 8,
+        "intermediate_dim": 2048,
     },
     "medium": {
         "num_layers": 8,
-        "hidden_size": 512,
+        "hidden_dim": 512,
         "dropout": 0.1,
-        "num_attention_heads": 8,
-        "inner_size": 2048,
-        "inner_activation": "gelu",
-        "initializer_range": 0.02,
+        "num_heads": 8,
+        "intermediate_dim": 2048,
     },
     "base": {
         "num_layers": 12,
-        "hidden_size": 768,
+        "hidden_dim": 768,
         "dropout": 0.1,
-        "num_attention_heads": 12,
-        "inner_size": 3072,
-        "inner_activation": "gelu",
-        "initializer_range": 0.02,
+        "num_heads": 12,
+        "intermediate_dim": 3072,
     },
     "large": {
         "num_layers": 24,
-        "hidden_size": 1024,
+        "hidden_dim": 1024,
         "dropout": 0.1,
-        "num_attention_heads": 16,
-        "inner_size": 4096,
-        "inner_activation": "gelu",
-        "initializer_range": 0.02,
+        "num_heads": 16,
+        "intermediate_dim": 4096,
     },
 }
 
diff --git a/examples/bert/bert_finetune_glue.py b/examples/bert/bert_finetune_glue.py
index 55e010d3..6b1beec6 100644
--- a/examples/bert/bert_finetune_glue.py
+++ b/examples/bert/bert_finetune_glue.py
@@ -18,23 +18,16 @@ import tempfile
 import datasets
 import keras_tuner
 import tensorflow as tf
-import tensorflow_text as tf_text
 from absl import app
 from absl import flags
 from tensorflow import keras
 
+import keras_nlp
 from examples.bert.bert_config import FINETUNING_CONFIG
-from examples.bert.bert_config import MODEL_CONFIGS
 from examples.bert.bert_config import PREPROCESSING_CONFIG
 
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string(
-    "model_size",
-    "tiny",
-    "One of: tiny, mini, small, medium, base, or large.",
-)
-
 flags.DEFINE_string(
     "vocab_file",
     None,
@@ -72,39 +65,6 @@ flags.DEFINE_bool(
 )
 
 
-def pack_inputs(
-    inputs,
-    seq_length,
-    start_of_sequence_id,
-    end_of_segment_id,
-    padding_id,
-):
-    # In case inputs weren't truncated (as they should have been),
-    # fall back to some ad-hoc truncation.
-    trimmed_segments = tf_text.RoundRobinTrimmer(
-        seq_length - len(inputs) - 1
-    ).trim(inputs)
-    # Combine segments.
-    segments_combined, segment_ids = tf_text.combine_segments(
-        trimmed_segments,
-        start_of_sequence_id=start_of_sequence_id,
-        end_of_segment_id=end_of_segment_id,
-    )
-    # Pad to dense Tensors.
-    input_word_ids, _ = tf_text.pad_model_inputs(
-        segments_combined, seq_length, pad_value=padding_id
-    )
-    input_type_ids, input_mask = tf_text.pad_model_inputs(
-        segment_ids, seq_length, pad_value=0
-    )
-    # Assemble nest of input tensors as expected by BERT model.
-    return {
-        "input_ids": input_word_ids,
-        "input_mask": input_mask,
-        "segment_ids": input_type_ids,
-    }
-
-
 def load_data(task_name):
     if task_name in ("cola", "sst2"):
         feature_names = ("sentence",)
@@ -142,39 +102,14 @@ def load_data(task_name):
     return train_ds, test_ds, validation_ds
 
 
-class BertClassificationFinetuner(keras.Model):
-    """Adds a classification head to a pre-trained BERT model for finetuning"""
-
-    def __init__(self, bert_model, num_classes, initializer, **kwargs):
-        super().__init__(**kwargs)
-        self.bert_model = bert_model
-        self._logit_layer = keras.layers.Dense(
-            num_classes,
-            kernel_initializer=initializer,
-            name="logits",
-        )
-
-    def call(self, inputs):
-        # Ignore the sequence output, use the pooled output.
-        _, pooled_output = self.bert_model(inputs)
-        return self._logit_layer(pooled_output)
-
-
 class BertHyperModel(keras_tuner.HyperModel):
     """Creates a hypermodel to help with the search space for finetuning."""
 
-    def __init__(self, model_config):
-        self.model_config = model_config
-
     def build(self, hp):
         model = keras.models.load_model(FLAGS.saved_model_input, compile=False)
-        model_config = self.model_config
-        finetuning_model = BertClassificationFinetuner(
-            bert_model=model,
+        finetuning_model = keras_nlp.models.BertClassifier(
+            base_model=model,
             num_classes=3 if FLAGS.task_name in ("mnli", "ax") else 2,
-            initializer=keras.initializers.TruncatedNormal(
-                stddev=model_config["initializer_range"]
-            ),
         )
         finetuning_model.compile(
             optimizer=keras.optimizers.Adam(
@@ -191,31 +126,24 @@ class BertHyperModel(keras_tuner.HyperModel):
 def main(_):
     print(f"Reading input model from {FLAGS.saved_model_input}")
 
-    vocab = []
-    with open(FLAGS.vocab_file, "r") as vocab_file:
-        for line in vocab_file:
-            vocab.append(line.strip())
-    tokenizer = tf_text.BertTokenizer(
-        FLAGS.vocab_file,
-        lower_case=FLAGS.do_lower_case,
-        token_out_type=tf.int32,
+    tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
+        vocabulary=FLAGS.vocab_file,
+        lowercase=FLAGS.do_lower_case,
+    )
+    packer = keras_nlp.layers.MultiSegmentPacker(
+        sequence_length=PREPROCESSING_CONFIG["max_seq_length"],
+        start_value=tokenizer.token_to_id("[CLS]"),
+        end_value=tokenizer.token_to_id("[SEP]"),
     )
-    start_id = vocab.index("[CLS]")
-    end_id = vocab.index("[SEP]")
-    pad_id = vocab.index("[PAD]")
-
-    model_config = MODEL_CONFIGS[FLAGS.model_size]
 
     def preprocess_data(inputs, labels):
-        inputs = [tokenizer.tokenize(x).merge_dims(1, -1) for x in inputs]
-        inputs = pack_inputs(
-            inputs,
-            PREPROCESSING_CONFIG["max_seq_length"],
-            start_of_sequence_id=start_id,
-            end_of_segment_id=end_id,
-            padding_id=pad_id,
-        )
-        return inputs, labels
+        inputs = [tokenizer(x) for x in inputs]
+        token_ids, segment_ids = packer(inputs)
+        return {
+            "input_ids": token_ids,
+            "input_mask": tf.cast(token_ids != 0, "int32"),
+            "segment_ids": segment_ids,
+        }, labels
 
     # Read and preprocess GLUE task data.
     train_ds, test_ds, validation_ds = load_data(FLAGS.task_name)
@@ -232,7 +160,7 @@ def main(_):
     )
 
     # Create a hypermodel object for a RandomSearch.
-    hypermodel = BertHyperModel(model_config)
+    hypermodel = BertHyperModel()
 
     # Initialize the random search over the 4 learning rate parameters, for 4
     # trials and 3 epochs for each trial.
@@ -263,8 +191,6 @@ def main(_):
         print("Evaluating on test set.")
         finetuning_model.evaluate(test_ds)
 
-    # TODO(mattdangerw): After incorporating keras_nlp tokenization, save an
-    # end-to-end model includeing preprocessing that operates on raw strings.
     if FLAGS.saved_model_output:
         print(f"Saving to {FLAGS.saved_model_output}")
         finetuning_model.save(FLAGS.saved_model_output)
diff --git a/examples/bert/bert_model.py b/examples/bert/bert_model.py
deleted file mode 100644
index d6dd1256..00000000
--- a/examples/bert/bert_model.py
+++ /dev/null
@@ -1,200 +0,0 @@
-# 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.
-"""Bert model and layer implementations.
-
-These components come from the tensorflow official model repository for BERT:
-https://github.com/tensorflow/models/tree/master/official/nlp/modeling
-
-This is to get us into a testable state. We should work to replace all of these
-components with components from the keras-nlp library.
-"""
-
-import tensorflow as tf
-from tensorflow import keras
-
-import keras_nlp
-
-
-class BertModel(keras.Model):
-    """Bi-directional Transformer-based encoder network.
-
-    This network implements a bi-directional Transformer-based encoder as
-    described in "BERT: Pre-training of Deep Bidirectional Transformers for
-    Language Understanding" (https://arxiv.org/abs/1810.04805). It includes the
-    embedding lookups and transformer layers, but not the masked language model
-    or classification task networks.
-
-    The default values for this object are taken from the BERT-Base
-    implementation in "BERT: Pre-training of Deep Bidirectional Transformers for
-    Language Understanding".
-
-    Args:
-        vocab_size: The size of the token vocabulary.
-        num_layers: The number of transformer layers.
-        hidden_size: The size of the transformer hidden layers.
-        dropout: Dropout probability for the Transformer encoder.
-        num_attention_heads: The number of attention heads for each transformer.
-            The hidden size must be divisible by the number of attention heads.
-        inner_size: The output dimension of the first Dense layer in a two-layer
-            feedforward network for each transformer.
-        inner_activation: The activation for the first Dense layer in a
-            two-layer feedforward network for each transformer.
-        initializer_range: The initialzer range to use for a truncated normal
-            initializer.
-        max_sequence_length: The maximum sequence length that this encoder can
-            consume. If None, max_sequence_length uses the value from sequence
-            length. This determines the variable shape for positional
-            embeddings.
-        type_vocab_size: The number of types that the 'segment_ids' input can
-            take.
-        norm_first: Whether to normalize inputs to attention and intermediate
-            dense layers. If set False, output of attention and intermediate
-            dense layers is normalized.
-    """
-
-    def __init__(
-        self,
-        vocab_size,
-        num_layers=12,
-        hidden_size=768,
-        dropout=0.1,
-        num_attention_heads=12,
-        inner_size=3072,
-        inner_activation="gelu",
-        initializer_range=0.02,
-        max_sequence_length=512,
-        type_vocab_size=2,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.num_layers = num_layers
-        self.num_attention_heads = num_attention_heads
-        self.max_sequence_length = max_sequence_length
-        self.type_vocab_size = type_vocab_size
-        self.inner_size = inner_size
-        self.inner_activation = keras.activations.get(inner_activation)
-        self.initializer_range = initializer_range
-        self.initializer = keras.initializers.TruncatedNormal(
-            stddev=initializer_range
-        )
-        self.dropout = dropout
-
-        self._embedding_layer = keras.layers.Embedding(
-            input_dim=vocab_size,
-            output_dim=hidden_size,
-            embeddings_initializer=self.initializer,
-            name="word_embeddings",
-        )
-
-        self._position_embedding_layer = keras_nlp.layers.PositionEmbedding(
-            initializer=self.initializer,
-            sequence_length=max_sequence_length,
-            name="position_embedding",
-        )
-
-        self._type_embedding_layer = keras.layers.Embedding(
-            input_dim=type_vocab_size,
-            output_dim=hidden_size,
-            embeddings_initializer=self.initializer,
-            name="type_embeddings",
-        )
-
-        self._embedding_norm_layer = keras.layers.LayerNormalization(
-            name="embeddings/layer_norm",
-            axis=-1,
-            epsilon=1e-12,
-            dtype=tf.float32,
-        )
-
-        self._embedding_dropout = keras.layers.Dropout(
-            rate=dropout, name="embedding_dropout"
-        )
-
-        self._transformer_layers = []
-        for i in range(num_layers):
-            layer = keras_nlp.layers.TransformerEncoder(
-                num_heads=num_attention_heads,
-                intermediate_dim=inner_size,
-                activation=self.inner_activation,
-                dropout=dropout,
-                kernel_initializer=self.initializer,
-                name="transformer/layer_%d" % i,
-            )
-            self._transformer_layers.append(layer)
-
-        # This is used as the intermediate output for the NSP prediction head.
-        # It is important we include this in the mode, as we want to preserve
-        # these weights for fine-tuning tasks.
-        self._pooler_layer = keras.layers.Dense(
-            units=hidden_size,
-            activation="tanh",
-            kernel_initializer=self.initializer,
-            name="pooler_dense",
-        )
-
-        self.inputs = dict(
-            input_ids=keras.Input(shape=(None,), dtype=tf.int32),
-            input_mask=keras.Input(shape=(None,), dtype=tf.int32),
-            segment_ids=keras.Input(shape=(None,), dtype=tf.int32),
-        )
-
-    def call(self, inputs):
-        if isinstance(inputs, dict):
-            input_ids = inputs.get("input_ids")
-            input_mask = inputs.get("input_mask")
-            segment_ids = inputs.get("segment_ids")
-        else:
-            raise ValueError(f"Inputs should be a dict. Received: {inputs}.")
-
-        word_embeddings = None
-        word_embeddings = self._embedding_layer(input_ids)
-        position_embeddings = self._position_embedding_layer(word_embeddings)
-        type_embeddings = self._type_embedding_layer(segment_ids)
-
-        embeddings = word_embeddings + position_embeddings + type_embeddings
-        embeddings = self._embedding_norm_layer(embeddings)
-        embeddings = self._embedding_dropout(embeddings)
-
-        x = embeddings
-        for layer in self._transformer_layers:
-            x = layer(x, padding_mask=input_mask)
-        sequence_output = x
-        pooled_output = self._pooler_layer(x[:, 0, :])  # 0 is the [CLS] token.
-        return sequence_output, pooled_output
-
-    def get_embedding_table(self):
-        return self._embedding_layer.embeddings
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "vocab_size": self.vocab_size,
-                "hidden_size": self.hidden_size,
-                "num_layers": self.num_layers,
-                "num_attention_heads": self.num_attention_heads,
-                "max_sequence_length": self.max_sequence_length,
-                "type_vocab_size": self.type_vocab_size,
-                "inner_size": self.inner_size,
-                "inner_activation": keras.activations.serialize(
-                    self.inner_activation
-                ),
-                "dropout": self.dropout,
-                "initializer_range": self.initializer_range,
-            }
-        )
-        return config
diff --git a/examples/bert/bert_preprocess.py b/examples/bert/bert_preprocess.py
index 36b6bc0e..00e0e8cb 100644
--- a/examples/bert/bert_preprocess.py
+++ b/examples/bert/bert_preprocess.py
@@ -375,6 +375,9 @@ def create_masked_lm_predictions(
 ):
     """Creates the predictions for the masked LM objective."""
 
+    # TODO(jbischof): replace with keras_nlp.layers.MLMMaskGenerator
+    # (Issue #166)
+
     cand_indexes = []
     for (i, token) in enumerate(tokens):
         if token == "[CLS]" or token == "[SEP]":
diff --git a/examples/bert/bert_train.py b/examples/bert/bert_train.py
index fc787cfd..ff5480f3 100644
--- a/examples/bert/bert_train.py
+++ b/examples/bert/bert_train.py
@@ -21,10 +21,10 @@ from absl import flags
 from absl import logging
 from tensorflow import keras
 
+import keras_nlp
 from examples.bert.bert_config import MODEL_CONFIGS
 from examples.bert.bert_config import PREPROCESSING_CONFIG
 from examples.bert.bert_config import TRAINING_CONFIG
-from examples.bert.bert_model import BertModel
 
 FLAGS = flags.FLAGS
 
@@ -91,7 +91,7 @@ flags.DEFINE_integer(
 
 
 class MaskedLMHead(keras.layers.Layer):
-    """Masked language model network head for BERT.
+    """Masked language model network head for Bert.
 
     This layer implements a masked language model based on the provided
     transformer based encoder. It assumes that the encoder network being passed
@@ -99,13 +99,14 @@ class MaskedLMHead(keras.layers.Layer):
 
     Example:
     ```python
-    encoder=modeling.networks.BertEncoder(...)
-    lm_layer=MaskedLMHead(embedding_table=encoder.get_embedding_table())
+    encoder = keras_nlp.models.BertBase()
+    lm_layer = MaskedLMHead(embedding_table=encoder.get_embedding_table())
     ```
 
     Args:
         embedding_table: The embedding table from encoder network.
-        inner_activation: The activation, if any, for the inner dense layer.
+        intermediate_activation: The activation, if any, for the inner dense
+            layer.
         initializer: The initializer for the dense layer. Defaults to a Glorot
             uniform initializer.
         output: The output style for this layer. Can be either 'logits' or
@@ -115,20 +116,22 @@ class MaskedLMHead(keras.layers.Layer):
     def __init__(
         self,
         embedding_table,
-        inner_activation="gelu",
+        intermediate_activation="gelu",
         initializer="glorot_uniform",
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.embedding_table = embedding_table
-        self.inner_activation = keras.activations.get(inner_activation)
+        self.intermediate_activation = keras.activations.get(
+            intermediate_activation
+        )
         self.initializer = initializer
 
     def build(self, input_shape):
-        self._vocab_size, hidden_size = self.embedding_table.shape
+        self._vocab_size, hidden_dim = self.embedding_table.shape
         self.dense = keras.layers.Dense(
-            hidden_size,
-            activation=self.inner_activation,
+            hidden_dim,
+            activation=self.intermediate_activation,
             kernel_initializer=self.initializer,
             name="transform/dense",
         )
@@ -162,7 +165,7 @@ class MaskedLMHead(keras.layers.Layer):
 
         Args:
             sequence_tensor: Sequence output of shape
-                (`batch_size`, `seq_length`, `hidden_size`) where `hidden_size`
+                (`batch_size`, `seq_length`, `hidden_dim`) where `hidden_dim`
                 is number of hidden units.
             positions: Positions ids of tokens in sequence to mask for
                 pretraining of with dimension (batch_size, num_predictions)
@@ -171,7 +174,7 @@ class MaskedLMHead(keras.layers.Layer):
 
         Returns:
             Masked out sequence tensor of shape (batch_size * num_predictions,
-            `hidden_size`).
+            `hidden_dim`).
         """
         sequence_shape = tf.shape(sequence_tensor)
         batch_size, seq_length = sequence_shape[0], sequence_shape[1]
@@ -189,17 +192,20 @@ class MaskedLMHead(keras.layers.Layer):
         return output_tensor
 
 
-class BertPretrainer(keras.Model):
-    def __init__(self, bert_model, **kwargs):
+class BertPretrainingModel(keras.Model):
+    """MLM + NSP model with Bert encoder."""
+
+    def __init__(self, encoder, **kwargs):
         super().__init__(**kwargs)
-        self.bert_model = bert_model
+        self.encoder = encoder
+        # TODO(jbischof): replace with keras_nlp.layers.MLMHead (Issue #166)
         self.masked_lm_head = MaskedLMHead(
-            bert_model.get_embedding_table(),
-            initializer=bert_model.initializer,
+            embedding_table=encoder.token_embedding.embeddings,
+            initializer=keras.initializers.TruncatedNormal(stddev=0.02),
         )
         self.next_sentence_head = keras.layers.Dense(
-            2,
-            kernel_initializer=bert_model.initializer,
+            encoder.num_segments,
+            kernel_initializer=keras.initializers.TruncatedNormal(stddev=0.02),
         )
         self.loss_tracker = keras.metrics.Mean(name="loss")
         self.lm_loss_tracker = keras.metrics.Mean(name="lm_loss")
@@ -212,13 +218,17 @@ class BertPretrainer(keras.Model):
         )
 
     def call(self, data):
-        sequence_output, pooled_output = self.bert_model(
+        encoder_output = self.encoder(
             {
                 "input_ids": data["input_ids"],
                 "input_mask": data["input_mask"],
                 "segment_ids": data["segment_ids"],
             }
         )
+        sequence_output, pooled_output = (
+            encoder_output["sequence_output"],
+            encoder_output["pooled_output"],
+        )
         lm_preds = self.masked_lm_head(
             sequence_output, data["masked_lm_positions"]
         )
@@ -396,14 +406,13 @@ def main(_):
     dataset = dataset.repeat()
 
     with strategy.scope():
-        # Create a BERT model the input config.
-        model = BertModel(
-            vocab_size=len(vocab),
-            **model_config,
+        # Create a Bert model the input config.
+        encoder = keras_nlp.models.Bert(
+            vocabulary_size=len(vocab), **model_config
         )
         # Make sure model has been called.
-        model(model.inputs)
-        model.summary()
+        encoder(encoder.inputs)
+        encoder.summary()
 
         # Allow overriding train steps from the command line for quick testing.
         if FLAGS.num_train_steps is not None:
@@ -420,7 +429,7 @@ def main(_):
         )
         optimizer = keras.optimizers.Adam(learning_rate=learning_rate_schedule)
 
-        pretraining_model = BertPretrainer(model)
+        pretraining_model = BertPretrainingModel(encoder)
         pretraining_model.compile(
             optimizer=optimizer,
         )
@@ -443,7 +452,7 @@ def main(_):
 
     model_path = FLAGS.saved_model_output
     logging.info(f"Saving to {FLAGS.saved_model_output}")
-    model.save(model_path)
+    encoder.save(model_path)
 
 
 if __name__ == "__main__":
diff --git a/keras_nlp/__init__.py b/keras_nlp/__init__.py
index 506d071c..67ab65ba 100644
--- a/keras_nlp/__init__.py
+++ b/keras_nlp/__init__.py
@@ -14,6 +14,7 @@
 
 from keras_nlp import layers
 from keras_nlp import metrics
+from keras_nlp import models
 from keras_nlp import tokenizers
 from keras_nlp import utils
 
diff --git a/keras_nlp/layers/random_deletion.py b/keras_nlp/layers/random_deletion.py
index 582f3918..d25ae723 100644
--- a/keras_nlp/layers/random_deletion.py
+++ b/keras_nlp/layers/random_deletion.py
@@ -46,7 +46,7 @@ class RandomDeletion(keras.layers.Layer):
             indicates that should not be considered a candidate for deletion.
             Unlike the `skip_fn` argument, this argument need not be
             tracable--it can be any python function.
-        seed: A seed for the rng.
+        seed: A seed for the random number generator.
 
     Examples:
 
diff --git a/keras_nlp/layers/sine_position_encoding.py b/keras_nlp/layers/sine_position_encoding.py
index 0ce8ec3c..e3aca378 100644
--- a/keras_nlp/layers/sine_position_encoding.py
+++ b/keras_nlp/layers/sine_position_encoding.py
@@ -62,6 +62,8 @@ class SinePositionEncoding(keras.layers.Layer):
         self.max_wavelength = max_wavelength
 
     def call(self, inputs):
+        # TODO(jbischof): replace `hidden_size` with`hidden_dim` for consistency
+        # with other layers.
         input_shape = tf.shape(inputs)
         # length of sequence is the second last dimension of the inputs
         seq_length = input_shape[-2]
diff --git a/keras_nlp/layers/transformer_encoder.py b/keras_nlp/layers/transformer_encoder.py
index c67409d9..08503447 100644
--- a/keras_nlp/layers/transformer_encoder.py
+++ b/keras_nlp/layers/transformer_encoder.py
@@ -68,7 +68,6 @@ class TransformerEncoder(keras.layers.Layer):
     # Call encoder on the inputs.
     input_data = tf.random.uniform(shape=[2, 10, 64])
     output = model(input_data)
-
     ```
 
     References:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
new file mode 100644
index 00000000..261a563b
--- /dev/null
+++ b/keras_nlp/models/__init__.py
@@ -0,0 +1,17 @@
+# 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.
+
+from keras_nlp.models.bert import Bert
+from keras_nlp.models.bert import BertBase
+from keras_nlp.models.bert import BertClassifier
diff --git a/keras_nlp/models/bert.py b/keras_nlp/models/bert.py
new file mode 100644
index 00000000..adaf3d93
--- /dev/null
+++ b/keras_nlp/models/bert.py
@@ -0,0 +1,324 @@
+# 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.
+
+"""Bert model configurable class, preconfigured versions, and task heads."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers import PositionEmbedding
+from keras_nlp.layers import TransformerEncoder
+
+
+def _bert_kernel_initializer(stddev=0.02):
+    return keras.initializers.TruncatedNormal(stddev=stddev)
+
+
+class Bert(keras.Model):
+    """Bi-directional Transformer-based encoder network.
+
+    This network implements a bi-directional Transformer-based encoder as
+    described in ["BERT: Pre-training of Deep Bidirectional Transformers for
+    Language Understanding"](https://arxiv.org/abs/1810.04805). It includes the
+    embedding lookups and transformer layers, but not the masked language model
+    or classification task networks.
+
+    This class gives a fully configurable Bert model with any number of layers,
+    heads, and embedding dimensions. For specific specific bert architectures
+    defined in the paper, see for example `keras_nlp.models.BertBase`.
+
+    Args:
+        vocabulary_size: Int. The size of the token vocabulary.
+        num_layers: Int. The number of transformer layers.
+        num_heads: Int. The number of attention heads for each transformer.
+            The hidden size must be divisible by the number of attention heads.
+        hidden_dim: Int. The size of the transformer encoding and pooler layers.
+        intermediate_dim: Int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        dropout: Float. Dropout probability for the Transformer encoder.
+        max_sequence_length: Int. The maximum sequence length that this encoder
+            can consume. If None, `max_sequence_length` uses the value from
+            sequence length. This determines the variable shape for positional
+            embeddings.
+        num_segments: Int. The number of types that the 'segment_ids' input can
+            take.
+        cls_token_index: Int. Index of [CLS] token in the vocabulary.
+        name: String, optional. Name of the model.
+        trainable: Boolean, optional. If the model's variables should be
+            trainable.
+
+    Example usage:
+    ```python
+    # Randomly initialized Bert encoder
+    encoder = keras_nlp.models.Bert(
+        vocabulary_size=30522,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=12
+    )
+
+     # Call encoder on the inputs.
+    input_data = {
+        "input_ids": tf.random.uniform(
+            shape=(1, 12), dtype=tf.int64, maxval=30522),
+        "segment_ids": tf.constant(
+            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+        "input_mask": tf.constant(
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+    }
+    output = encoder(input_data)
+    ```
+    """
+
+    # TODO(jbischof): consider changing `intermediate_dim` and `hidden_dim` to
+    # less confusing name here and in TransformerEncoder (`feed_forward_dim`?)
+
+    def __init__(
+        self,
+        vocabulary_size,
+        num_layers,
+        num_heads,
+        hidden_dim,
+        intermediate_dim,
+        dropout=0.1,
+        max_sequence_length=512,
+        num_segments=2,
+        cls_token_index=0,
+        name=None,
+        trainable=True,
+    ):
+
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="input_ids"
+        )
+        segment_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="segment_ids"
+        )
+        input_mask = keras.Input(
+            shape=(None,), dtype="int32", name="input_mask"
+        )
+
+        # Embed tokens, positions, and segment ids.
+        token_embedding_layer = keras.layers.Embedding(
+            input_dim=vocabulary_size,
+            output_dim=hidden_dim,
+            embeddings_initializer=_bert_kernel_initializer(),
+            name="token_embedding",
+        )
+        token_embedding = token_embedding_layer(token_id_input)
+        position_embedding = PositionEmbedding(
+            initializer=_bert_kernel_initializer(),
+            sequence_length=max_sequence_length,
+            name="position_embedding",
+        )(token_embedding)
+        segment_embedding = keras.layers.Embedding(
+            input_dim=num_segments,
+            output_dim=hidden_dim,
+            embeddings_initializer=_bert_kernel_initializer(),
+            name="segment_embedding",
+        )(segment_id_input)
+
+        # Sum, normailze and apply dropout to embeddings.
+        x = keras.layers.Add()(
+            (token_embedding, position_embedding, segment_embedding)
+        )
+        x = keras.layers.LayerNormalization(
+            name="embeddings_layer_norm",
+            axis=-1,
+            epsilon=1e-12,
+            dtype=tf.float32,
+        )(x)
+        x = keras.layers.Dropout(
+            dropout,
+            name="embeddings_dropout",
+        )(x)
+
+        # Apply successive transformer encoder blocks.
+        for i in range(num_layers):
+            x = TransformerEncoder(
+                num_heads=num_heads,
+                intermediate_dim=intermediate_dim,
+                activation=lambda x: keras.activations.gelu(
+                    x, approximate=True
+                ),
+                dropout=dropout,
+                kernel_initializer=_bert_kernel_initializer(),
+                name=f"""transformer_layer_{i}""",
+            )(x, padding_mask=input_mask)
+
+        # Construct the two Bert outputs. The pooled output is a dense layer on
+        # top of the [CLS] token.
+        sequence_output = x
+        pooled_output = keras.layers.Dense(
+            hidden_dim,
+            kernel_initializer=_bert_kernel_initializer(),
+            activation="tanh",
+            name="pooled_dense",
+        )(x[:, cls_token_index, :])
+
+        # Instantiate using Functional API Model constructor
+        super().__init__(
+            inputs={
+                "input_ids": token_id_input,
+                "segment_ids": segment_id_input,
+                "input_mask": input_mask,
+            },
+            outputs={
+                "sequence_output": sequence_output,
+                "pooled_output": pooled_output,
+            },
+            name=name,
+            trainable=trainable,
+        )
+        # All references to `self` below this line
+        self.token_embedding = token_embedding_layer
+        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.num_segments = num_segments
+        self.intermediate_dim = intermediate_dim
+        self.dropout = dropout
+        self.cls_token_index = cls_token_index
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "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,
+                "num_segments": self.num_segments,
+                "dropout": self.dropout,
+                "cls_token_index": self.cls_token_index,
+            }
+        )
+        return config
+
+
+class BertClassifier(keras.Model):
+    """Bert encoder model with a classification head.
+
+    Args:
+        base_model: A `keras_nlp.models.Bert` to encode inputs.
+        num_classes: Int. Number of classes to predict.
+        name: String, optional. Name of the model.
+        trainable: Boolean, optional. If the model's variables should be
+            trainable.
+
+    Example usage:
+    ```
+    python
+    # Randomly initialized Bert encoder
+    encoder = keras_nlp.models.Bert(
+        vocabulary_size=30522,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=12
+    )
+
+     # Call classifier on the inputs.
+    input_data = {
+        "input_ids": tf.random.uniform(
+            shape=(1, 12), dtype=tf.int64, maxval=30522),
+        "segment_ids": tf.constant(
+            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+        "input_mask": tf.constant(
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+    }
+    classifier = keras_nlp.models.BertClassifier(encoder, 4)
+    logits = classifier(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        base_model,
+        num_classes,
+        name=None,
+        trainable=True,
+    ):
+        inputs = base_model.input
+        pooled = base_model(inputs)["pooled_output"]
+        outputs = keras.layers.Dense(
+            num_classes,
+            kernel_initializer=_bert_kernel_initializer(),
+            name="logits",
+        )(pooled)
+        # Instantiate using Functional API Model constructor
+        super().__init__(
+            inputs=inputs, outputs=outputs, name=name, trainable=trainable
+        )
+        # All references to `self` below this line
+        self.base_model = base_model
+        self.num_classes = num_classes
+
+
+def BertBase(name=None, trainable=True):
+    """Bi-directional Transformer-based encoder network (Bert) using "Base"
+    architecture.
+
+    This network implements a bi-directional Transformer-based encoder as
+    described in ["BERT: Pre-training of Deep Bidirectional Transformers for
+    Language Understanding"](https://arxiv.org/abs/1810.04805). It includes the
+    embedding lookups and transformer layers, but not the masked language model
+    or classification task networks.
+
+    Args:
+        name: String, optional. Name of the model.
+        trainable: Boolean, optional. If the model's variables should be
+            trainable.
+
+    Example usage:
+    ```python
+    # Randomly initialized BertBase encoder
+    encoder = keras_nlp.models.BertBase()
+
+     # Call encoder on the inputs.
+    input_data = {
+        "input_ids": tf.random.uniform(
+            shape=(1, 512), dtype=tf.int64, maxval=encoder.vocabulary_size),
+        "segment_ids": tf.constant(
+            [0] * 200 + [1] * 312, shape=(1, 512)),
+        "input_mask": tf.constant(
+            [1] * 512, shape=(1, 512)),
+    }
+    output = encoder(input_data)
+    ```
+    """
+
+    model = Bert(
+        vocabulary_size=30522,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        dropout=0.1,
+        max_sequence_length=512,
+        name=name,
+        trainable=trainable,
+    )
+
+    # TODO(jbischof): add some documentation or magic to load our checkpoints
+    # TODO(jbischof): attach the tokenizer
+    return model
diff --git a/keras_nlp/tokenizers/__init__.py b/keras_nlp/tokenizers/__init__.py
index 05817187..048ea15d 100644
--- a/keras_nlp/tokenizers/__init__.py
+++ b/keras_nlp/tokenizers/__init__.py
@@ -14,6 +14,9 @@
 
 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 (
+    compute_sentence_piece_proto,
+)
 from keras_nlp.tokenizers.tokenizer import Tokenizer
 from keras_nlp.tokenizers.unicode_codepoint_tokenizer import (
     UnicodeCodepointTokenizer,
diff --git a/keras_nlp/tokenizers/sentence_piece_tokenizer_trainer.py b/keras_nlp/tokenizers/sentence_piece_tokenizer_trainer.py
new file mode 100644
index 00000000..66189901
--- /dev/null
+++ b/keras_nlp/tokenizers/sentence_piece_tokenizer_trainer.py
@@ -0,0 +1,142 @@
+# 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.
+"""Trainer for SentencePiece Tokenizer."""
+import io
+
+import tensorflow as tf
+
+try:
+    import sentencepiece as spm
+except ImportError:
+    spm = None
+
+
+def compute_sentence_piece_proto(
+    data,
+    vocabulary_size,
+    model_type="unigram",
+    proto_output_file=None,
+    lowercase=False,
+):
+    r"""A utility to train a SentencePiece vocabulary.
+
+    Trains a SentencePiece vocabulary from an input dataset or a list of
+    filenames.
+
+    If `data` is a list of filenames, the file format is required to be plain
+    text files, and the text will be read in line by line during training.
+
+    Args:
+        data: A `tf.data.Dataset`, or a list of filenames.
+        vocabulary_size: int. The maximum size of a vocabulary to be trained.
+        model_type: str, defaults to `"unigram"`. The model algorithm must be one
+            of `"unigram"`, `"bpe"`, `"word"` or `"char"`.
+        proto_output_file: str, defaults to `None`. If provided it will be used
+            as model_file which is passed to model_writer.
+            If `None`, the model_file will be `io.BytesIO` object.
+        lowercase: bool, defaults to `False`. If true, the input text will be
+            lowercased before tokenization.
+
+    Returns:
+        A `bytes` object with a serialized SentencePiece proto or
+        `None` if proto_output_file if provided.
+
+    Examples:
+
+    Basic Usage (from Dataset).
+    >>> inputs = tf.data.Dataset.from_tensor_slices(["Drifting Along"])
+    >>> proto = keras_nlp.tokenizers.compute_sentence_piece_proto(inputs, vocabulary_size=15)
+    >>> tokenizer = keras_nlp.tokenizers.SentencePieceTokenizer(proto=proto)
+    >>> outputs = inputs.map(tokenizer)
+    >>> for output in outputs:
+    ...     print(output)
+    tf.Tensor([ 4  8 12  5  9 14  5  6 13  4  7 10 11  6 13],
+    shape=(15,), dtype=int32)
+
+    Basic Usage (with files).
+    ``` python
+    with open("test.txt", "w+") as f: f.write("Drifting Along\n")
+    inputs = ["test.txt"]
+    proto = keras_nlp.tokenizers.compute_sentence_piece_proto(
+         inputs, vocabulary_size=15, proto_output_file="model.spm")
+    tokenizer = keras_nlp.tokenizers.SentencePieceTokenizer(proto="model.spm")
+    ds = tf.data.Dataset.from_tensor_slices(["the quick brown fox."])
+    ds = ds.map(tokenizer)
+    for output in ds:
+         print(output)
+    tf.Tensor([ 4 14  1  4  1  5  1  4  1 12 11  1  6  4  9 11  1],
+    shape=(17,), dtype=int32)
+    ```
+
+    Usage with lowercase
+    >>> inputs = tf.data.Dataset.from_tensor_slices(["Drifting Along"])
+    >>> proto = keras_nlp.tokenizers.compute_sentence_piece_proto(
+    ...     inputs, vocabulary_size=15, lowercase=True)
+    >>> tokenizer = keras_nlp.tokenizers.SentencePieceTokenizer(proto=proto)
+    >>> outputs = inputs.map(tokenizer)
+    >>> for output in outputs:
+    ...     print(output)
+    tf.Tensor([ 4  8 12  5  9 14  5  6 13  4  7 10 11  6 13],
+    shape=(15,), dtype=int32)
+    """
+
+    if spm is None:
+        raise ImportError(
+            "sentencepiece is not installed. Please install it via `pip install sentencepiece`."
+        )
+
+    if not isinstance(data, (list, tuple, tf.data.Dataset)):
+        raise ValueError(
+            "The `data` argument must be either `tf.data.Dataset` or `tuple` or `list`. "
+            f"Received: type(data)={type(data)}."
+        )
+
+    if model_type not in ["unigram", "bpe", "word", "char"]:
+        raise ValueError(
+            "The `model_type` argument must be one of `unigram`, `bpe`, `word`"
+            f"or `char`. Received: model_type={model_type}."
+        )
+
+    model_writer = (
+        open(proto_output_file, "wb") if proto_output_file else io.BytesIO()
+    )
+    is_dataset = isinstance(data, tf.data.Dataset)
+    if is_dataset:
+        spm.SentencePieceTrainer.train(
+            sentence_iterator=data.as_numpy_iterator(),
+            model_writer=model_writer,
+            vocab_size=vocabulary_size,
+            model_type=model_type,
+            normalization_rule_name="nmt_nfkc_cf" if lowercase else "nmt_nfkc",
+            pad_id=0,
+            unk_id=1,
+            bos_id=2,
+            eos_id=3,
+        )
+    else:
+        spm.SentencePieceTrainer.train(
+            input=data,
+            model_writer=model_writer,
+            vocab_size=vocabulary_size,
+            model_type=model_type,
+            normalization_rule_name="nmt_nfkc_cf" if lowercase else "nmt_nfkc",
+            pad_id=0,
+            unk_id=1,
+            bos_id=2,
+            eos_id=3,
+        )
+    if proto_output_file:
+        model_writer.close()
+    else:
+        return model_writer.getvalue()
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index b078f1ef..e17c4bdd 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -148,9 +148,9 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
             plain text file containing a single WordPiece token per line.
         sequence_length: int. If set, the output will be converted to a dense
             tensor and padded/trimmed so all outputs are of sequence_length.
-        lowercase: bool, defaults to `True`. If true, the input text will be
+        lowercase: bool, defaults to `False`. If true, the input text will be
             lowercased before tokenization.
-        strip_accents: bool, defaults to `True`. If true, all accent marks will
+        strip_accents: bool, defaults to `False`. If true, all accent marks will
             be removed from text before tokenization.
         split: bool, defaults to `True`. If true, input will be split on
             whitespace and punctuation marks, and all punctuation marks will be
@@ -172,7 +172,10 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
     Ragged outputs.
     >>> vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
     >>> inputs = ["The quick brown fox."]
-    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab)
+    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
+    ...     vocabulary=vocab,
+    ...     lowercase=True,
+    ... )
     >>> tokenizer(inputs)
     <tf.RaggedTensor [[1, 2, 3, 4, 5, 6, 7]]>
 
@@ -180,7 +183,10 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
     >>> vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
     >>> inputs = ["The quick brown fox."]
     >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
-    ...     vocabulary=vocab, sequence_length=10)
+    ...     vocabulary=vocab,
+    ...     sequence_length=10,
+    ...     lowercase=True,
+    ... )
     >>> tokenizer(inputs)
     <tf.Tensor: shape=(1, 10), dtype=int32, numpy=
     array([[1, 2, 3, 4, 5, 6, 7, 0, 0, 0]], dtype=int32)>
@@ -189,22 +195,32 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
     >>> vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
     >>> inputs = ["The quick brown fox."]
     >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
-    ...     vocabulary=vocab, dtype="string")
+    ...     vocabulary=vocab,
+    ...     lowercase=True,
+    ...     dtype="string",
+    ... )
     >>> tokenizer(inputs)
     <tf.RaggedTensor [[b'the', b'qu', b'##ick', b'br', b'##own', b'fox', b'.']]>
 
     Detokenization.
     >>> vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
     >>> inputs = "The quick brown fox."
-    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab)
+    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
+    ...     vocabulary=vocab,
+    ...     lowercase=True,
+    ... )
     >>> tokenizer.detokenize(tokenizer.tokenize(inputs)).numpy().decode('utf-8')
     'the quick brown fox .'
 
     Custom splitting.
     >>> vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
     >>> inputs = ["The$quick$brown$fox"]
-    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab,
-    ...     split=False, dtype='string')
+    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
+    ...     vocabulary=vocab,
+    ...     split=False,
+    ...     lowercase=True,
+    ...     dtype='string',
+    ... )
     >>> split_inputs = tf.strings.split(inputs, sep="$")
     >>> tokenizer(split_inputs)
     <tf.RaggedTensor [[b'the', b'qu', b'##ick', b'br', b'##own', b'fox']]>
@@ -214,8 +230,8 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
         self,
         vocabulary=None,
         sequence_length: int = None,
-        lowercase: bool = True,
-        strip_accents: bool = True,
+        lowercase: bool = False,
+        strip_accents: bool = False,
         split: bool = True,
         suffix_indicator: str = "##",
         oov_token: str = "[UNK]",
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_trainer.py b/keras_nlp/tokenizers/word_piece_tokenizer_trainer.py
index eda6c989..bfdd258e 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer_trainer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_trainer.py
@@ -25,8 +25,8 @@ def compute_word_piece_vocabulary(
     data,
     vocabulary_size,
     vocabulary_output_file=None,
-    lowercase=True,
-    strip_accents=True,
+    lowercase=False,
+    strip_accents=False,
     split=True,
     suffix_indicator="##",
     reserved_tokens=["[PAD]", "[CLS]", "[SEP]", "[UNK]", "[MASK]"],
@@ -45,9 +45,9 @@ def compute_word_piece_vocabulary(
         vocabulary_size: int. The maximum size of a vocabulary to be trained.
         vocabulary_output_file: str, defaults to `None`. The location to write a
             vocabulary file.
-        lowercase: bool, defaults to `True`. If true, the input text will be
+        lowercase: bool, defaults to `False`. If true, the input text will be
             lowercased before tokenization.
-        strip_accents: bool, defaults to `True`. If true, all accent marks will
+        strip_accents: bool, defaults to `False`. If true, all accent marks will
             be removed from text before tokenization.
         split: bool, defaults to `True`. If true, input will be split on
             whitespace and punctuation marks, and all punctuation marks will be
@@ -77,9 +77,12 @@ def compute_word_piece_vocabulary(
     tf.Tensor([ 6 12 10 12  8 12  7 12  9 12], shape=(10,), dtype=int32)
 
     Basic Usage (from filenames).
-    >>> with open("test.txt", "w+") as f: f.write("bat sat pat mat rat\n")
-    >>> inputs = ["test.txt"]
-    >>> vocab = compute_word_piece_vocabulary(inputs, 13)
+    ```python
+    with open("test.txt", "w+") as f:
+        f.write("bat sat pat mat rat\n")
+    inputs = ["test.txt"]
+    vocab = compute_word_piece_vocabulary(inputs, 13)
+    ```
 
     Custom Split Usage (from Dataset).
     >>> def normalize_and_split(x):
@@ -89,7 +92,7 @@ def compute_word_piece_vocabulary(
     >>> inputs = tf.data.Dataset.from_tensor_slices(["bat sat: pat mat rat.\n"])
     >>> split_inputs = inputs.map(normalize_and_split)
     >>> vocab = compute_word_piece_vocabulary(
-    ...     split_inputs, 13, split=False
+    ...     split_inputs, 13, split=False,
     ... )
     >>> vocab
     ['[PAD]', '[CLS]', '[SEP]', '[UNK]', '[MASK]', 'a', 'b', 'm', 'p', 'r', 's', 't', '##at']
@@ -97,18 +100,19 @@ def compute_word_piece_vocabulary(
     >>> inputs.map(tokenizer.tokenize)
 
     Custom Split Usage (from filenames).
-    >>> def normalize_and_split(x):
-    ...     "Strip punctuation and split on whitespace."
-    ...     x = tf.strings.regex_replace(x, r"\p{P}", "")
-    ...     return tf.strings.split(x)
-    >>> with open("test.txt", "w+") as f: f.write("bat sat: pat mat rat.\n")
-    >>> inputs = tf.data.TextLineDataset(["test.txt"])
-    >>> split_inputs = inputs.map(normalize_and_split)
-    >>> vocab = compute_word_piece_vocabulary(split_inputs, 13, split=False)
-    >>> vocab
-    ['[PAD]', '[CLS]', '[SEP]', '[UNK]', '[MASK]', 'a', 'b', 'm', 'p', 'r', 's', 't', '##at']
-    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab)
-    >>> inputs.map(tokenizer.tokenize)
+    ```python
+    def normalize_and_split(x):
+        "Strip punctuation and split on whitespace."
+        x = tf.strings.regex_replace(x, r"\p{P}", "")
+        return tf.strings.split(x)
+    with open("test.txt", "w+") as f:
+        f.write("bat sat: pat mat rat.\n")
+    inputs = tf.data.TextLineDataset(["test.txt"])
+    split_inputs = inputs.map(normalize_and_split)
+    vocab = compute_word_piece_vocabulary(split_inputs, 13, split=False)
+    tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab)
+    inputs.map(tokenizer.tokenize)
+    ```
     """
     # Read data files.
 
diff --git a/keras_nlp/utils/text_generation.py b/keras_nlp/utils/text_generation.py
index 3b9a088b..afc20b84 100644
--- a/keras_nlp/utils/text_generation.py
+++ b/keras_nlp/utils/text_generation.py
@@ -18,6 +18,8 @@ import tensorflow as tf
 from absl import logging
 from tensorflow import keras
 
+# TODO (@chenmoneygithub): Refactor code to reuse snippets.
+
 
 def validate_prompt(prompt):
     """Helper function to validate input to text_generation utils."""
@@ -52,7 +54,11 @@ def mask_tokens_after_end_token(prompt, max_length, end_token_id, pad_token_id):
     # Find index of first end_token_id.
     end_indices = tf.math.argmax(prompt == end_token_id, -1)
     # Use max_length if no `end_token_id` is found.
-    end_indices = tf.where(end_indices == 0, max_length, end_indices)
+    end_indices = tf.where(
+        end_indices == 0,
+        tf.cast(max_length, dtype=end_indices.dtype),
+        end_indices,
+    )
     # Build a mask including end_token and replace tokens after end_token
     # with `pad_token_id`.
     valid_indices = tf.sequence_mask(end_indices + 1, maxlen=max_length)
@@ -128,14 +134,6 @@ def greedy_search(
     ```
 
     """
-    if not tf.executing_eagerly():
-        raise RuntimeError(
-            "`keras_nlp.utils.greedy_search` currently requires an eager "
-            "execution context. Please call `greedy_search` outside "
-            "tf.function or run `tf.config.run_functions_eagerly(True)` to run "
-            "tf.function in eager mode."
-        )
-
     prompt = validate_prompt(prompt)
 
     input_is_1d = prompt.shape.rank == 1
@@ -143,14 +141,41 @@ def greedy_search(
         prompt = prompt[tf.newaxis, :]
     validate_token_probability_fn(token_probability_fn, prompt)
 
-    i = prompt.shape[1]
-    while i < max_length:
-        # If the prompt has reached our desired length, exit while loop.
-        pred = token_probability_fn(prompt)
+    shape = tf.shape(prompt)
+    batch_size = shape[0]
+    length = shape[1]
+
+    # Pad the prompt with `pad_token_id` to `max_length`.
+    padding = tf.fill((batch_size, max_length - length), pad_token_id)
+    prompt = tf.concat((prompt, padding), axis=1)
+
+    def one_step(length, prompt):
+        pred = token_probability_fn(prompt[:, :length])
         next_token = tf.cast(tf.argmax(pred, axis=-1), dtype=prompt.dtype)
+
         # Append the next token to current sequence.
-        prompt = tf.concat([prompt, next_token[:, tf.newaxis]], axis=-1)
-        i += 1
+        def add_token(args):
+            sequence, token = args
+            return tf.tensor_scatter_nd_update(
+                tensor=sequence, indices=[[length]], updates=[token]
+            )
+
+        prompt = tf.map_fn(
+            fn=add_token,
+            elems=(prompt, next_token),
+            fn_output_signature=tf.TensorSpec(
+                shape=(max_length), dtype=prompt.dtype
+            ),
+        )
+        length += 1
+        return (length, prompt)
+
+    # Run a while loop till text of length `max_length` has been generated.
+    length, prompt = tf.while_loop(
+        cond=lambda length, _: tf.less(length, max_length),
+        body=one_step,
+        loop_vars=(length, prompt),
+    )
 
     if end_token_id is not None:
         prompt = mask_tokens_after_end_token(
@@ -382,33 +407,55 @@ def random_search(
     ```
 
     """
-    if not tf.executing_eagerly():
-        raise RuntimeError(
-            "`keras_nlp.utils.random_sampling` currently requires an eager "
-            "execution context. Please call `random_sampling` outside "
-            "tf.function or run `tf.config.run_functions_eagerly(True)` to run "
-            "tf.function in eager mode."
-        )
-
     prompt = validate_prompt(prompt)
     input_is_1d = prompt.shape.rank == 1
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
     validate_token_probability_fn(token_probability_fn, prompt)
 
-    i = prompt.shape[1]
-    while i < max_length:
-        # If the prompt has reached our desired length, exit while loop.
+    shape = tf.shape(prompt)
+    batch_size = shape[0]
+    length = shape[1]
+
+    # Pad the prompt with `pad_token_id` to `max_length`.
+    padding = tf.fill((batch_size, max_length - length), pad_token_id)
+    prompt = tf.concat((prompt, padding), axis=1)
+
+    def one_step(length, prompt):
         pred = token_probability_fn(prompt)
         if from_logits:
             pred = keras.activations.softmax(pred, axis=-1)
-        next_token = tf.cast(
-            tf.random.categorical(tf.math.log(pred), 1, seed=seed),
-            dtype=prompt.dtype,
+        next_token = tf.squeeze(
+            tf.cast(
+                tf.random.categorical(tf.math.log(pred), 1, seed=seed),
+                dtype=prompt.dtype,
+            ),
+            axis=1,
         )
+
         # Append the next token to current sequence.
-        prompt = tf.concat([prompt, next_token], axis=-1)
-        i += 1
+        def add_token(args):
+            sequence, token = args
+            return tf.tensor_scatter_nd_update(
+                tensor=sequence, indices=[[length]], updates=[token]
+            )
+
+        prompt = tf.map_fn(
+            fn=add_token,
+            elems=(prompt, next_token),
+            fn_output_signature=tf.TensorSpec(
+                shape=(max_length), dtype=prompt.dtype
+            ),
+        )
+        length += 1
+        return (length, prompt)
+
+    # Run a while loop till text of length `max_length` has been generated.
+    length, prompt = tf.while_loop(
+        cond=lambda length, _: tf.less(length, max_length),
+        body=one_step,
+        loop_vars=(length, prompt),
+    )
 
     if end_token_id is not None:
         prompt = mask_tokens_after_end_token(
@@ -497,14 +544,6 @@ def top_k_search(
     ```
 
     """
-    if not tf.executing_eagerly():
-        raise RuntimeError(
-            "`keras_nlp.utils.top_k_search` currently requires an eager "
-            "execution context. Please call `top_k_search` outside "
-            "tf.function or run `tf.config.run_functions_eagerly(True)` to run "
-            "tf.function in eager mode."
-        )
-
     prompt = validate_prompt(prompt)
     input_is_1d = prompt.shape.rank == 1
     if input_is_1d:
@@ -522,8 +561,15 @@ def top_k_search(
         )
         k = pred.shape[1]
 
-    i = prompt.shape[1]
-    while i < max_length:
+    shape = tf.shape(prompt)
+    batch_size = shape[0]
+    length = shape[1]
+
+    # Pad the prompt with `pad_token_id` to `max_length`.
+    padding = tf.fill((batch_size, max_length - length), pad_token_id)
+    prompt = tf.concat((prompt, padding), axis=1)
+
+    def one_step(length, prompt):
         pred = token_probability_fn(prompt)
         if from_logits:
             pred = keras.activations.softmax(pred, axis=-1)
@@ -534,12 +580,34 @@ def top_k_search(
         next_token = tf.random.categorical(
             tf.math.log(top_k_pred), 1, seed=seed
         )
+
         # Rearrange to get the next token idx from the original order.
         next_token = tf.gather_nd(top_k_indices, next_token, batch_dims=1)
         next_token = tf.cast(next_token, dtype=prompt.dtype)
+
         # Append the next token to current sequence.
-        prompt = tf.concat([prompt, next_token[:, tf.newaxis]], axis=-1)
-        i += 1
+        def add_token(args):
+            sequence, token = args
+            return tf.tensor_scatter_nd_update(
+                tensor=sequence, indices=[[length]], updates=[token]
+            )
+
+        prompt = tf.map_fn(
+            fn=add_token,
+            elems=(prompt, next_token),
+            fn_output_signature=tf.TensorSpec(
+                shape=(max_length), dtype=prompt.dtype
+            ),
+        )
+        length += 1
+        return (length, prompt)
+
+    # Run a while loop till text of length `max_length` has been generated.
+    length, prompt = tf.while_loop(
+        cond=lambda length, _: tf.less(length, max_length),
+        body=one_step,
+        loop_vars=(length, prompt),
+    )
 
     if end_token_id is not None:
         prompt = mask_tokens_after_end_token(
@@ -630,13 +698,6 @@ def top_p_search(
     ```
 
     """
-    if not tf.executing_eagerly():
-        raise RuntimeError(
-            "`keras_nlp.utils.top_p_search` currently requires an eager "
-            "execution context. Please call `top_p_search` outside "
-            "tf.function or run `tf.config.run_functions_eagerly(True)` to run "
-            "tf.function in eager mode."
-        )
     if p <= 0 or p >= 1:
         raise ValueError(
             f"`p` should be in the range (0, 1). Received: `p={p}`."
@@ -648,9 +709,15 @@ def top_p_search(
         prompt = prompt[tf.newaxis, :]
     validate_token_probability_fn(token_probability_fn, prompt)
 
-    i = prompt.shape[1]
-    while i < max_length:
-        # If the prompt has reached our desired length, exit while loop.
+    shape = tf.shape(prompt)
+    batch_size = shape[0]
+    length = shape[1]
+
+    # Pad the prompt with `pad_token_id` to `max_length`.
+    padding = tf.fill((batch_size, max_length - length), pad_token_id)
+    prompt = tf.concat((prompt, padding), axis=1)
+
+    def one_step(length, prompt):
         pred = token_probability_fn(prompt)
         if from_logits:
             pred = keras.activations.softmax(pred, axis=-1)
@@ -679,9 +746,30 @@ def top_p_search(
             sorted_indices, sorted_next_token, batch_dims=1
         )
         next_token = tf.cast(next_token, dtype=prompt.dtype)
+
         # Append the next token to current sequence.
-        prompt = tf.concat([prompt, next_token[:, tf.newaxis]], axis=-1)
-        i += 1
+        def add_token(args):
+            sequence, token = args
+            return tf.tensor_scatter_nd_update(
+                tensor=sequence, indices=[[length]], updates=[token]
+            )
+
+        prompt = tf.map_fn(
+            fn=add_token,
+            elems=(prompt, next_token),
+            fn_output_signature=tf.TensorSpec(
+                shape=(max_length), dtype=prompt.dtype
+            ),
+        )
+        length += 1
+        return (length, prompt)
+
+    # Run a while loop till text of length `max_length` has been generated.
+    length, prompt = tf.while_loop(
+        cond=lambda length, _: tf.less(length, max_length),
+        body=one_step,
+        loop_vars=(length, prompt),
+    )
 
     if end_token_id is not None:
         prompt = mask_tokens_after_end_token(


Test patch:
diff --git a/keras_nlp/layers/random_deletion_test.py b/keras_nlp/layers/random_deletion_test.py
index edacd3f3..b345400a 100644
--- a/keras_nlp/layers/random_deletion_test.py
+++ b/keras_nlp/layers/random_deletion_test.py
@@ -105,6 +105,8 @@ class RandomDeletionTest(tf.test.TestCase):
         output = tf.strings.reduce_join(augmented, separator=" ", axis=-1)
         self.assertAllEqual(output.shape, tf.convert_to_tensor(inputs).shape)
         exp_output = [b"Hey like", b"Keras Tensorflow"]
+        for i in range(output.shape[0]):
+            self.assertAllEqual(output[i], exp_output[i])
 
     def test_get_config_and_from_config(self):
         augmenter = random_deletion.RandomDeletion(
diff --git a/keras_nlp/models/bert_test.py b/keras_nlp/models/bert_test.py
new file mode 100644
index 00000000..a329e0f0
--- /dev/null
+++ b/keras_nlp/models/bert_test.py
@@ -0,0 +1,112 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for Bert model."""
+
+import os
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.models import bert
+
+
+class BertTest(tf.test.TestCase):
+    def test_valid_call_bert(self):
+        model = bert.Bert(
+            vocabulary_size=30522,
+            num_layers=12,
+            num_heads=12,
+            hidden_dim=768,
+            intermediate_dim=3072,
+            max_sequence_length=12,
+            name="encoder",
+        )
+        input_data = {
+            "input_ids": tf.random.uniform(
+                shape=(1, 12), dtype=tf.int64, maxval=30522
+            ),
+            "segment_ids": tf.constant(
+                [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+            ),
+            "input_mask": tf.constant(
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+            ),
+        }
+        model(input_data)
+
+    def test_valid_call_classifier(self):
+        model = bert.Bert(
+            vocabulary_size=30522,
+            num_layers=12,
+            num_heads=12,
+            hidden_dim=768,
+            intermediate_dim=3072,
+            max_sequence_length=12,
+            name="encoder",
+        )
+        input_data = {
+            "input_ids": tf.random.uniform(
+                shape=(1, 12), dtype=tf.int64, maxval=30522
+            ),
+            "segment_ids": tf.constant(
+                [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+            ),
+            "input_mask": tf.constant(
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+            ),
+        }
+        classifier = bert.BertClassifier(model, 4, name="classifier")
+        classifier(input_data)
+
+    def test_valid_call_bert_base(self):
+        model = bert.BertBase(name="encoder")
+        input_data = {
+            "input_ids": tf.random.uniform(
+                shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
+            ),
+            "segment_ids": tf.constant([0] * 200 + [1] * 312, shape=(1, 512)),
+            "input_mask": tf.constant([1] * 512, shape=(1, 512)),
+        }
+        model(input_data)
+
+    def test_saving_model(self):
+        model = bert.Bert(
+            vocabulary_size=30522,
+            num_layers=12,
+            num_heads=12,
+            hidden_dim=768,
+            intermediate_dim=3072,
+            max_sequence_length=12,
+            name="encoder",
+        )
+        input_data = {
+            "input_ids": tf.random.uniform(
+                shape=(1, 12), dtype=tf.int64, maxval=30522
+            ),
+            "segment_ids": tf.constant(
+                [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+            ),
+            "input_mask": tf.constant(
+                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+            ),
+        }
+        model_output = model(input_data)
+        save_path = os.path.join(self.get_temp_dir(), "model")
+        model.save(save_path)
+        restored_model = keras.models.load_model(save_path)
+
+        restored_output = restored_model(input_data)
+        self.assertAllClose(
+            model_output["pooled_output"], restored_output["pooled_output"]
+        )
diff --git a/keras_nlp/tokenizers/sentence_piece_tokenizer_trainer_test.py b/keras_nlp/tokenizers/sentence_piece_tokenizer_trainer_test.py
new file mode 100644
index 00000000..9dd732a1
--- /dev/null
+++ b/keras_nlp/tokenizers/sentence_piece_tokenizer_trainer_test.py
@@ -0,0 +1,98 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for SentencePiece Tokenizer Trainer."""
+
+import os
+import re
+
+import tensorflow as tf
+
+from keras_nlp.tokenizers.sentence_piece_tokenizer import SentencePieceTokenizer
+from keras_nlp.tokenizers.sentence_piece_tokenizer_trainer import (
+    compute_sentence_piece_proto,
+)
+
+
+class SentencePieceTokenizerTrainerTest(tf.test.TestCase):
+    def test_dataset_input(self):
+        test_text = ["Ninjas and Samurais"]
+        expected_output = [
+            [5, 9, 6, 7, 11, 4, 8, 5, 4, 7, 10, 5, 12, 4, 13, 15, 14, 4, 6, 8]
+        ]
+        data = tf.data.Dataset.from_tensor_slices(test_text)
+        proto = compute_sentence_piece_proto(data, 16)
+        tokenizer = SentencePieceTokenizer(proto=proto)
+        test_output = tokenizer(test_text)
+        self.assertAllEqual(expected_output, test_output)
+
+    def test_file_input(self):
+        test_text = "Ninja Land"
+        with open(os.path.join(self.get_temp_dir(), "test.txt"), "w+") as f:
+            f.write(test_text + "\n")
+        expected_output = [6, 8, 9, 5, 11, 4, 6, 7, 4, 5, 10]
+        proto = compute_sentence_piece_proto(
+            [os.path.join(self.get_temp_dir(), "test.txt")],
+            12,
+        )
+        tokenizer = SentencePieceTokenizer(proto=proto)
+        test_output = tokenizer(test_text)
+        self.assertAllEqual(expected_output, test_output)
+
+    def test_multiple_file_input(self):
+        with open(os.path.join(self.get_temp_dir(), "test1.txt"), "w+") as f:
+            f.write("Drifting Along\n")
+        with open(os.path.join(self.get_temp_dir(), "test2.txt"), "w+") as f:
+            f.write("Woah look there\n")
+        inputs = [
+            os.path.join(self.get_temp_dir(), "test1.txt"),
+            os.path.join(self.get_temp_dir(), "test2.txt"),
+        ]
+        proto = compute_sentence_piece_proto(inputs, 20)
+        tokenizer = SentencePieceTokenizer(proto=proto)
+        test_text = "Woah Along"
+        test_output = tokenizer(test_text)
+        expected_output = [4, 16, 5, 17, 9, 4, 15, 12, 5, 11, 18]
+        self.assertAllEqual(expected_output, test_output)
+
+    def test_invalid_input(self):
+        test_text_invalid = {"file1": "test.txt"}
+        with self.assertRaisesRegex(
+            ValueError,
+            re.escape(f"Received: type(data)={type(test_text_invalid)}."),
+        ):
+            compute_sentence_piece_proto(test_text_invalid, 10)
+
+    def test_lowercase(self):
+        inputs = tf.data.Dataset.from_tensor_slices(["Drifting Along"])
+        proto = compute_sentence_piece_proto(
+            inputs, vocabulary_size=15, lowercase=True
+        )
+        tokenizer = SentencePieceTokenizer(proto=proto)
+        output = inputs.map(tokenizer).take(1).get_single_element()
+        expected_output = [4, 8, 12, 5, 9, 14, 5, 6, 13, 4, 7, 10, 11, 6, 13]
+        self.assertAllEqual(expected_output, output)
+
+    def test_proto_output_file(self):
+        inputs = tf.data.Dataset.from_tensor_slices(["Drifting Along"])
+        compute_sentence_piece_proto(
+            inputs,
+            vocabulary_size=15,
+            proto_output_file=os.path.join(self.get_temp_dir(), "model.spm"),
+        )
+        tokenizer = SentencePieceTokenizer(
+            proto=os.path.join(self.get_temp_dir(), "model.spm")
+        )
+        output = inputs.map(tokenizer).take(1).get_single_element()
+        expected_output = [4, 8, 12, 5, 9, 14, 5, 6, 13, 4, 7, 10, 11, 6, 13]
+        self.assertAllEqual(expected_output, output)
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
index 1d5a62e8..cde365a4 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer_test.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
@@ -89,7 +89,7 @@ class WordPieceTokenizerTest(tf.test.TestCase):
     def test_lowercase(self):
         input_data = ["the QUicK brOWN FOX"]
         vocab_data = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox"]
-        tokenizer = WordPieceTokenizer(vocabulary=vocab_data)
+        tokenizer = WordPieceTokenizer(vocabulary=vocab_data, lowercase=True)
         call_output = tokenizer(input_data)
         self.assertAllEqual(call_output, [[1, 2, 3, 4, 5, 6]])
 
@@ -103,7 +103,9 @@ class WordPieceTokenizerTest(tf.test.TestCase):
     def test_strip_accents(self):
         input_data = ["á é í ó ú"]
         vocab_data = ["[UNK]", "a", "e", "i", "o", "u"]
-        tokenizer = WordPieceTokenizer(vocabulary=vocab_data)
+        tokenizer = WordPieceTokenizer(
+            vocabulary=vocab_data, strip_accents=True
+        )
         call_output = tokenizer(input_data)
         self.assertAllEqual(call_output, [[1, 2, 3, 4, 5]])
 
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_trainer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_trainer_test.py
index e048966e..85b11be5 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer_trainer_test.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_trainer_test.py
@@ -32,18 +32,21 @@ class WordPieceTokenizerTrainerTest(tf.test.TestCase):
 
     def test_filenames_input(self):
         test_text = "baa maa caa saa aaa"
-        with open("test.txt", "w+") as f:
+        input_file = os.path.join(self.get_temp_dir(), "test.txt")
+        with open(input_file, "w+") as f:
             f.write(test_text + "\n")
         test_output = ["a", "b", "c", "m", "s", "##aa", "##a", "##b"]
         vocab = compute_word_piece_vocabulary(
-            ["test.txt"], 8, reserved_tokens=[]
+            [input_file],
+            8,
+            reserved_tokens=[],
         )
         self.assertAllEqual(vocab, test_output)
-        os.remove("test.txt")
 
     def test_filenames_without_split(self):
         test_text = "baa maa caa saa aaa"
-        with open("test.txt", "w+") as f:
+        input_file = os.path.join(self.get_temp_dir(), "test.txt")
+        with open(input_file, "w+") as f:
             f.write(test_text + "\n")
 
         with self.assertRaisesRegex(
@@ -54,7 +57,6 @@ class WordPieceTokenizerTrainerTest(tf.test.TestCase):
             r"`compute_word_piece_vocabulary\(\)` with split=False.",
         ):
             compute_word_piece_vocabulary(["test.txt"], 10, split=False)
-        os.remove("test.txt")
 
     def test_invalid_input(self):
         test_text_invalid = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4])
@@ -143,16 +145,19 @@ class WordPieceTokenizerTrainerTest(tf.test.TestCase):
     def test_output_file(self):
         test_text = tf.data.Dataset.from_tensor_slices(["BaA Maa Caa Saa AAa"])
         test_output = ["a", "b", "c", "m", "s", "##aa", "##a", "##b"]
-
+        vocab_file = os.path.join(self.get_temp_dir(), "test.txt")
         compute_word_piece_vocabulary(
-            test_text, 8, vocabulary_output_file="test.txt", reserved_tokens=[]
+            test_text,
+            8,
+            vocab_file,
+            lowercase=True,
+            reserved_tokens=[],
         )
         vocab_from_file = []
-        with open("test.txt", "r") as f:
+        with open(vocab_file, "r") as f:
             for line in f:
                 vocab_from_file.append(line.strip())
         self.assertAllEqual(vocab_from_file, test_output)
-        os.remove("test.txt")
 
     def test_reserved_tokens(self):
         # This dummy text/token would be removed for being too long.
diff --git a/keras_nlp/utils/text_generation_test.py b/keras_nlp/utils/text_generation_test.py
index 6030ce44..9fb3f084 100644
--- a/keras_nlp/utils/text_generation_test.py
+++ b/keras_nlp/utils/text_generation_test.py
@@ -17,6 +17,7 @@ import random
 
 import numpy as np
 import tensorflow as tf
+from absl.testing import parameterized
 from tensorflow import keras
 
 from keras_nlp.utils.text_generation import beam_search
@@ -26,7 +27,7 @@ from keras_nlp.utils.text_generation import top_k_search
 from keras_nlp.utils.text_generation import top_p_search
 
 
-class GreedySearchTextGenerationTest(tf.test.TestCase):
+class GreedySearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         super().setUp()
         self.vocab_size = 10
@@ -81,7 +82,7 @@ class GreedySearchTextGenerationTest(tf.test.TestCase):
     def test_assert_generation_is_correct(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         batch_size = 10
@@ -97,7 +98,7 @@ class GreedySearchTextGenerationTest(tf.test.TestCase):
     def test_end_token_id(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         max_length = 5
@@ -113,6 +114,69 @@ class GreedySearchTextGenerationTest(tf.test.TestCase):
         expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
         self.assertAllEqual(outputs, expected_outputs)
 
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = greedy_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(inputs)
+        self.assertAllEqual(outputs, expected_outputs)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile_batched_ds(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = greedy_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        ds = tf.data.Dataset.from_tensor_slices(inputs).batch(2)
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(ds)
+        self.assertAllEqual(outputs, expected_outputs)
+
 
 class BeamSearchTextGenerationTest(tf.test.TestCase):
     def setUp(self):
@@ -249,7 +313,7 @@ class BeamSearchTextGenerationTest(tf.test.TestCase):
     def test_assert_generation_is_correct(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         batch_size = 10
@@ -269,7 +333,7 @@ class BeamSearchTextGenerationTest(tf.test.TestCase):
     def test_end_token_id(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         max_length = 5
@@ -287,7 +351,7 @@ class BeamSearchTextGenerationTest(tf.test.TestCase):
         self.assertAllEqual(outputs, expected_outputs)
 
 
-class RandomSearchTextGenerationTest(tf.test.TestCase):
+class RandomSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         super().setUp()
         self.vocab_size = 10
@@ -342,7 +406,7 @@ class RandomSearchTextGenerationTest(tf.test.TestCase):
     def test_assert_probability_distribution_generation_is_correct(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         batch_size = 10
@@ -360,14 +424,14 @@ class RandomSearchTextGenerationTest(tf.test.TestCase):
                 outputs_count[pred] += 1
         self.assertAllClose(
             outputs_count / np.sum(outputs_count),
-            [0.01, 0.01, 0.08, 0.9],
+            [0.0, 0.0, 0.0, 1.0],
             rtol=0.2,
         )
 
     def test_end_token_id(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         max_length = 5
@@ -386,6 +450,71 @@ class RandomSearchTextGenerationTest(tf.test.TestCase):
         expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
         self.assertAllEqual(outputs, expected_outputs)
 
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = random_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    seed=42,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model(inputs)
+        self.assertAllEqual(outputs, expected_outputs)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile_batched_ds(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = random_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    seed=42,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        ds = tf.data.Dataset.from_tensor_slices(inputs).batch(2)
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(ds)
+        self.assertAllEqual(outputs, expected_outputs)
+
     def test_from_logits(self):
         def token_logits_fn(inputs):
             batch_size = inputs.shape[0]
@@ -420,7 +549,7 @@ class RandomSearchTextGenerationTest(tf.test.TestCase):
         self.assertAllEqual(output_logit, output_probs)
 
 
-class TopKSearchTextGenerationTest(tf.test.TestCase):
+class TopKSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         super().setUp()
         self.vocab_size = 10
@@ -491,7 +620,7 @@ class TopKSearchTextGenerationTest(tf.test.TestCase):
     def test_assert_probability_distribution_generation_is_correct(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.1, 0.2, 0.3, 0.4]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         batch_size = 10
@@ -513,7 +642,7 @@ class TopKSearchTextGenerationTest(tf.test.TestCase):
                 outputs_count[pred] += 1
         self.assertAllClose(
             outputs_count / np.sum(outputs_count),
-            [0.0, 0.0, 0.429, 0.571],
+            [0.0, 0.0, 0.0, 1.0],
             rtol=0.2,
         )
 
@@ -539,7 +668,7 @@ class TopKSearchTextGenerationTest(tf.test.TestCase):
     def test_end_token_id(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         max_length = 5
@@ -559,6 +688,71 @@ class TopKSearchTextGenerationTest(tf.test.TestCase):
         expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
         self.assertAllEqual(outputs, expected_outputs)
 
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = random_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    seed=42,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(inputs)
+        self.assertAllEqual(outputs, expected_outputs)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile_batched_ds(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = random_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    seed=42,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        ds = tf.data.Dataset.from_tensor_slices(inputs).batch(2)
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(ds)
+        self.assertAllEqual(outputs, expected_outputs)
+
     def test_from_logits(self):
         def token_logits_fn(inputs):
             batch_size = inputs.shape[0]
@@ -595,7 +789,7 @@ class TopKSearchTextGenerationTest(tf.test.TestCase):
         self.assertAllEqual(output_logit, output_probs)
 
 
-class TopPSearchTextGenerationTest(tf.test.TestCase):
+class TopPSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         super().setUp()
         self.vocab_size = 10
@@ -656,7 +850,7 @@ class TopPSearchTextGenerationTest(tf.test.TestCase):
     def test_assert_probability_distribution_generation_is_correct(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.1, 0.2, 0.3, 0.4]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         batch_size = 10
@@ -678,7 +872,7 @@ class TopPSearchTextGenerationTest(tf.test.TestCase):
                 outputs_count[pred] += 1
         self.assertAllClose(
             outputs_count / np.sum(outputs_count),
-            [0.0, 0.0, 0.429, 0.571],
+            [0.0, 0.0, 0.0, 1.0],
             rtol=0.2,
         )
 
@@ -700,7 +894,7 @@ class TopPSearchTextGenerationTest(tf.test.TestCase):
     def test_end_token_id(self):
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
             return tf.repeat(prob, batch_size, axis=0)
 
         max_length = 5
@@ -720,6 +914,73 @@ class TopPSearchTextGenerationTest(tf.test.TestCase):
         expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
         self.assertAllEqual(outputs, expected_outputs)
 
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = top_p_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    p=0.92,
+                    seed=1,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(inputs)
+        self.assertAllEqual(outputs, expected_outputs)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_model_compile_batched_ds(self, jit_compile):
+        def token_probability_fn(inputs):
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])
+            return tf.repeat(prob, 2, axis=0)
+
+        max_length = 5
+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs):
+                generated = top_p_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=max_length,
+                    p=0.92,
+                    seed=1,
+                    end_token_id=2,
+                    pad_token_id=0,
+                )
+                return generated
+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        ds = tf.data.Dataset.from_tensor_slices(inputs).batch(2)
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+
+        model = TestModel()
+        model.compile(jit_compile=jit_compile)
+
+        outputs = model.predict(ds)
+        self.assertAllEqual(outputs, expected_outputs)
+
     def test_from_logits(self):
         def token_logits_fn(inputs):
             batch_size = inputs.shape[0]