Repository: keras-team/keras-nlp

Issue:
Add a BLEU metric
Splitting this issue out from https://github.com/keras-team/keras-nlp/issues/38.

We should add a BLEU metric as `keras_nlp.metrics.Bleu`.

Discussion:

Patch:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 5e1e4b91..bdf7689d 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -26,8 +26,8 @@ maintain two lists of github tags for contributors:
    a larger list of issues that may range in complexity.
 
 If you would like propose a new symbol or feature, please first review our
-design guide and roadmap linked above, and issue to discuss. If you have a
-specific design in mind, please include a Colab notebook showing the proposed
+design guide and roadmap linked above, and open an issue to discuss. If you have
+a specific design in mind, please include a Colab notebook showing the proposed
 design in a end-to-end example. Keep in mind that design for a new feature or
 use case may take longer than contributing to an open issue with a
 vetted-design.
diff --git a/README.md b/README.md
index c52da04c..61c5bb73 100644
--- a/README.md
+++ b/README.md
@@ -36,14 +36,23 @@ We are a new and growing project, and welcome [contributions](CONTRIBUTING.md).
 - [API Design Guide](API_DESIGN_GUIDE.md)
 - [Call for Contributions](https://github.com/keras-team/keras-nlp/issues?q=is%3Aissue+is%3Aopen+label%3A%22contributions+welcome%22)
 
-## Quick Start
+## Installation
 
-Install the latest release:
+To install the latest official release:
 
 ```
 pip install keras-nlp --upgrade
 ```
 
+To install the latest unreleased changes to the library, we recommend using
+pip to install directly from the master branch on github:
+
+```
+pip install git+https://github.com/keras-team/keras-nlp.git --upgrade
+```
+
+## Quickstart
+
 Tokenize text, build a tiny transformer, and train a single batch:
 
 ```python
@@ -80,6 +89,8 @@ model.compile(optimizer="adam", loss="binary_crossentropy", jit_compile=True)
 model.train_on_batch(x, y)
 ```
 
+For more in depth guides and examples, visit https://keras.io/keras_nlp/.
+
 ## Compatibility
 
 We follow [Semantic Versioning](https://semver.org/), and plan to
@@ -101,6 +112,8 @@ Here is the BibTeX entry:
 }
 ```
 
+## Acknowledgements
+
 Thank you to all of our wonderful contributors!
 
 <a href="https://github.com/keras-team/keras-nlp/graphs/contributors">
diff --git a/examples/bert/bert_finetune_glue.py b/examples/bert/bert_finetune_glue.py
index 598a09fb..b83623a7 100644
--- a/examples/bert/bert_finetune_glue.py
+++ b/examples/bert/bert_finetune_glue.py
@@ -156,7 +156,7 @@ class BertClassificationFinetuner(keras.Model):
             kernel_initializer=initializer,
             name="pooler",
         )
-        self._logit_layer = tf.keras.layers.Dense(
+        self._logit_layer = keras.layers.Dense(
             num_classes,
             kernel_initializer=initializer,
             name="logits",
diff --git a/examples/bert/bert_model.py b/examples/bert/bert_model.py
index 3027d8b3..faa150b3 100644
--- a/examples/bert/bert_model.py
+++ b/examples/bert/bert_model.py
@@ -47,7 +47,7 @@ def make_attention_mask(inputs, mask):
 
 
 # TODO(mattdangerw): This class is needed for TPU friendly embeddings, we should
-# remove it entirely and fix tf.keras.layers.Embedding as needed.
+# remove it entirely and fix keras.layers.Embedding as needed.
 class OnDeviceEmbedding(keras.layers.Layer):
     """Performs an embedding lookup suitable for TPU devices.
 
diff --git a/examples/bert/bert_train.py b/examples/bert/bert_train.py
index 27e234da..1a37536f 100644
--- a/examples/bert/bert_train.py
+++ b/examples/bert/bert_train.py
@@ -90,7 +90,7 @@ flags.DEFINE_integer(
 )
 
 
-class ClassificationHead(tf.keras.layers.Layer):
+class ClassificationHead(keras.layers.Layer):
     """Pooling head for sentence-level classification tasks.
 
     Args:
@@ -403,7 +403,7 @@ def get_checkpoint_callback():
             # Clear up the directory if users want to skip restoring.
             tf.io.gfile.rmtree(FLAGS.checkpoint_save_directory)
     checkpoint_path = FLAGS.checkpoint_save_directory
-    return tf.keras.callbacks.BackupAndRestore(
+    return keras.callbacks.BackupAndRestore(
         backup_dir=checkpoint_path,
     )
 
@@ -411,7 +411,7 @@ def get_checkpoint_callback():
 def get_tensorboard_callback():
     timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
     log_dir = FLAGS.tensorboard_log_path + timestamp
-    return tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)
+    return keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)
 
 
 def main(_):
@@ -419,7 +419,7 @@ def main(_):
         # If the job is on cloud, we will use cloud logging.
         import google.cloud.logging
 
-        tf.keras.utils.disable_interactive_logging()
+        keras.utils.disable_interactive_logging()
         client = google.cloud.logging.Client()
         client.setup_logging()
 
diff --git a/examples/machine_translation/inference.py b/examples/machine_translation/inference.py
index a2308e4b..77055eb8 100644
--- a/examples/machine_translation/inference.py
+++ b/examples/machine_translation/inference.py
@@ -16,6 +16,7 @@ import tensorflow as tf
 from absl import app
 from absl import flags
 from absl import logging
+from tensorflow import keras
 
 # Import data module to include the customized serializable, required for
 # loading tokenizer.
@@ -101,7 +102,7 @@ def decode_sequence(input_sentence, model, max_sequence_length, lookup_table):
 
 
 def main(_):
-    loaded_model = tf.keras.models.load_model(FLAGS.saved_model_path)
+    loaded_model = keras.models.load_model(FLAGS.saved_model_path)
 
     decoder_tokenizer = loaded_model.decoder_tokenizer
     vocab = decoder_tokenizer.get_vocabulary()
diff --git a/examples/machine_translation/train.py b/examples/machine_translation/train.py
index 648585bc..2e42b40a 100644
--- a/examples/machine_translation/train.py
+++ b/examples/machine_translation/train.py
@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import tensorflow as tf
 from absl import app
 from absl import flags
+from tensorflow import keras
 
 from examples.machine_translation.data import get_dataset_and_tokenizer
 from examples.machine_translation.model import TranslationModel
@@ -57,16 +57,16 @@ flags.DEFINE_string(
 
 
 def run_training(model, train_ds, val_ds):
-    learning_rate = tf.keras.optimizers.schedules.ExponentialDecay(
+    learning_rate = keras.optimizers.schedules.ExponentialDecay(
         initial_learning_rate=FLAGS.learning_rate,
         decay_steps=20,
         decay_rate=0.98,
     )
-    optimizer = tf.keras.optimizers.Adam(learning_rate)
-    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
-        reduction=tf.keras.losses.Reduction.NONE
+    optimizer = keras.optimizers.Adam(learning_rate)
+    loss_fn = keras.losses.SparseCategoricalCrossentropy(
+        reduction=keras.losses.Reduction.NONE
     )
-    metrics = tf.keras.metrics.SparseCategoricalAccuracy()
+    metrics = keras.metrics.SparseCategoricalAccuracy()
     model.compile(optimizer=optimizer, metrics=[metrics], loss=loss_fn)
     model.fit(
         train_ds,
diff --git a/keras_nlp/metrics/__init__.py b/keras_nlp/metrics/__init__.py
index 36c05c16..4c097878 100644
--- a/keras_nlp/metrics/__init__.py
+++ b/keras_nlp/metrics/__init__.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_nlp.metrics.bleu import Bleu
 from keras_nlp.metrics.edit_distance import EditDistance
 from keras_nlp.metrics.perplexity import Perplexity
 from keras_nlp.metrics.rouge_l import RougeL
diff --git a/keras_nlp/metrics/bleu.py b/keras_nlp/metrics/bleu.py
new file mode 100644
index 00000000..16267623
--- /dev/null
+++ b/keras_nlp/metrics/bleu.py
@@ -0,0 +1,386 @@
+# 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.
+
+"""BLEU metric implementation."""
+
+import collections
+import math
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.utils.tensor_utils import tensor_to_list
+
+REPLACE_SUBSTRINGS = [
+    ("<skipped>", ""),
+    ("-\n", ""),
+    ("\n", " "),
+    ("&quot;", '"'),
+    ("&amp;", "&"),
+    ("&lt;", "<"),
+    ("&gt;", ">"),
+]
+
+
+REGEX_PATTERNS = [
+    # language-dependent part (assuming Western languages)
+    (r"([\{-\~\[-\` -\&\(-\+\:-\@\/])", r" \1 "),
+    # tokenize period and comma unless preceded by a digit
+    (r"([^0-9])([\.,])", r"\1 \2 "),
+    # tokenize period and comma unless followed by a digit
+    (r"([\.,])([^0-9])", r" \1 \2"),
+    # tokenize dash when preceded by a digit
+    (r"([0-9])(-)", r"\1 \2 "),
+    # If last character is "." or ",", add space.
+    (r"[\.,]$", r" \0 \1"),
+    # one space only between words
+    (r"\s+", r" "),
+]
+
+
+class Bleu(keras.metrics.Metric):
+    """BLEU metric.
+
+    This class implements the BLEU metric. BLEU is generally used to evaluate
+    machine translation systems. By default, this implementation replicates
+    SacreBLEU, but user-defined tokenizers can be passed to deal with other
+    languages.
+
+    For BLEU score, we count the number of matching n-grams in the candidate
+    translation and the reference text. We find the "clipped count" of matching
+    n-grams so as to not give a high score to a (reference, prediction) pair
+    with redundant, repeated tokens. Secondly, BLEU score tends to reward
+    shorter predictions more, which is why a brevity penalty is applied to
+    penalise short predictions. For more details, see the following article:
+    https://cloud.google.com/translate/automl/docs/evaluate#bleu.
+
+    Note on input shapes:
+    For unbatched inputs, `y_pred` should be a tensor of shape `()`, and
+    `y_true` should be a tensor of shape `(num_references,)`. For batched
+    inputs, `y_pred` should be a tensor of shape `(batch_size,)`,
+    and `y_true` should be a tensor of shape `(batch_size, num_references)`. In
+    case of batched inputs, `y_true` can also be a ragged tensor of shape
+    `(batch_size, None)` if different samples have different number of
+    references.
+
+    Args:
+        tokenizer: callable. A function that takes a string `tf.RaggedTensor`
+            (of any shape), and tokenizes the strings in the tensor. If the
+            tokenizer is not specified, the default tokenizer is used. The
+            default tokenizer replicates the behaviour of SacreBLEU's
+            `"tokenizer_13a"` tokenizer
+            (https://github.com/mjpost/sacrebleu/blob/v2.1.0/sacrebleu/tokenizers/tokenizer_13a.py).
+        max_order: int. The maximum n-gram order to use. For example, if
+            `max_order` is set to 3, unigrams, bigrams, and trigrams will be
+            considered. Defaults to 4.
+        smooth: bool. Whether to apply Lin et al. 2004 smoothing to the BLEU
+            score. Adds 1 to the matched n-gram count (i.e., numerator) and 1
+            to the total n-gram count (i.e., denominator) for every order while
+            calculating precision. Defaults to False.
+        dtype: string or tf.dtypes.Dtype. Precision of metric computation. If
+               not specified, it defaults to tf.float32.
+        name: string. Name of the metric instance.
+        **kwargs: Other keyword arguments.
+
+    References:
+        - [Papineni et al., 2002](https://aclanthology.org/P02-1040/)
+        - [SacreBLEU](https://github.com/mjpost/sacrebleu)
+        - [Lin et al., 2004](https://aclanthology.org/P04-1077/)
+    """
+
+    def __init__(
+        self,
+        tokenizer=None,
+        max_order=4,
+        smooth=False,
+        dtype=None,
+        name="bleu",
+        **kwargs,
+    ):
+        super().__init__(name=name, dtype=dtype, **kwargs)
+
+        if not tf.as_dtype(self.dtype).is_floating:
+            raise ValueError(
+                "`dtype` must be a floating point type. "
+                f"Received: dtype={dtype}"
+            )
+
+        self.tokenizer = tokenizer
+        self.max_order = max_order
+        self.smooth = smooth
+
+        self._matches = self.add_weight(
+            shape=(self.max_order,),
+            name="bleu_matches",
+            initializer="zeros",
+            dtype=self.dtype,
+        )
+        self._possible_matches = self.add_weight(
+            shape=(self.max_order,),
+            name="bleu_possible_matches",
+            initializer="zeros",
+            dtype=self.dtype,
+        )
+        self._translation_length = self.add_weight(
+            name="bleu_translation_length",
+            initializer="zeros",
+            dtype=self.dtype,
+        )
+        self._reference_length = self.add_weight(
+            name="bleu_reference_length",
+            initializer="zeros",
+            dtype=self.dtype,
+        )
+        self._bleu = self.add_weight(
+            name="bleu",
+            initializer="zeros",
+            dtype=self.dtype,
+        )
+
+    def _tokenizer(self, inputs):
+        """
+        Tokenizes the input strings. By default, replicates the behaviour of
+        SacreBLEU's default tokenizer, namely, `tokenizer_13a`.
+        """
+        if self.tokenizer:
+            return self.tokenizer(inputs)
+
+        for pattern, replacement in REPLACE_SUBSTRINGS + REGEX_PATTERNS:
+            inputs = tf.strings.regex_replace(
+                input=inputs,
+                pattern=pattern,
+                rewrite=replacement,
+                replace_global=True,
+                name=None,
+            )
+        inputs = tf.strings.split(inputs)
+        return inputs
+
+    def _get_ngrams(self, segment, max_order):
+        """Extracts all n-grams up to a given maximum order from an input segment.
+
+        Uses Python ops. Inspired from
+        https://github.com/tensorflow/nmt/blob/master/nmt/scripts/bleu.py.
+
+        Args:
+            segment: list. Text segment from which n-grams will be
+                extracted.
+            max_order: int. Maximum length in tokens of the n-grams returned
+                by this method.
+        """
+        ngram_counts = collections.Counter()
+        for order in range(1, max_order + 1):
+            for i in range(0, len(segment) - order + 1):
+                ngram = tuple(segment[i : i + order])
+                ngram_counts[ngram] += 1
+        return ngram_counts
+
+    def _corpus_bleu(
+        self,
+        reference_corpus,
+        translation_corpus,
+        matches_by_order,
+        possible_matches_by_order,
+        translation_length,
+        reference_length,
+        max_order=4,
+        smooth=False,
+    ):
+        """Corpus BLEU implementation using Python ops.
+
+        Computes BLEU score of translated segments against one or more
+        references. Inspired from
+        https://github.com/tensorflow/nmt/blob/master/nmt/scripts/bleu.py.
+
+        Args:
+            reference_corpus: list of lists of references for each
+                translation. Each reference should be tokenized into a list
+                of tokens.
+            translation_corpus: list of translations to score. Each
+                translation should be tokenized into a list of tokens.
+            matches_by_order: list of floats containing the initial number
+                of matches for each order.
+            possible_matches_by_order: list of floats containing the initial
+                number of possible matches for each order.
+            translation_length: float. Initial number of tokens in all the
+                translations.
+            reference_length: float. Initial number of tokens in all the
+                references.
+            max_order: int. Maximum n-gram order to use when computing
+                BLEU score.
+            smooth: boolean. Whether or not to apply Lin et al. 2004
+                smoothing.
+        """
+        for (references, translation) in zip(
+            reference_corpus, translation_corpus
+        ):
+            reference_length += min(len(r) for r in references)
+            translation_length += len(translation)
+
+            merged_ref_ngram_counts = collections.Counter()
+            for reference in references:
+                merged_ref_ngram_counts |= self._get_ngrams(
+                    reference, max_order
+                )
+            translation_ngram_counts = self._get_ngrams(translation, max_order)
+            overlap = translation_ngram_counts & merged_ref_ngram_counts
+            for ngram in overlap:
+                matches_by_order[len(ngram) - 1] += overlap[ngram]
+            for order in range(1, max_order + 1):
+                possible_matches = len(translation) - order + 1
+                if possible_matches > 0:
+                    possible_matches_by_order[order - 1] += possible_matches
+
+        precisions = [0] * max_order
+        for i in range(0, max_order):
+            if smooth:
+                precisions[i] = (matches_by_order[i] + 1.0) / (
+                    possible_matches_by_order[i] + 1.0
+                )
+            else:
+                if possible_matches_by_order[i] > 0:
+                    precisions[i] = (
+                        float(matches_by_order[i])
+                        / possible_matches_by_order[i]
+                    )
+                else:
+                    precisions[i] = 0.0
+
+        if min(precisions) > 0:
+            p_log_sum = sum((1.0 / max_order) * math.log(p) for p in precisions)
+            geo_mean = math.exp(p_log_sum)
+        else:
+            geo_mean = 0
+
+        ratio = float(translation_length) / reference_length
+
+        if ratio > 1.0:
+            bp = 1.0
+        else:
+            bp = math.exp(1 - 1.0 / ratio)
+
+        bleu = geo_mean * bp
+
+        return (
+            bleu,
+            matches_by_order,
+            possible_matches_by_order,
+            translation_length,
+            reference_length,
+        )
+
+    def _calculate_bleu_score(self, references, translation):
+        references = tensor_to_list(references)
+        translation = tensor_to_list(translation)
+
+        matches = self._matches.numpy()
+        possible_matches = self._possible_matches.numpy()
+        translation_length = self._translation_length.numpy()
+        reference_length = self._reference_length.numpy()
+
+        (
+            bleu_score,
+            matches,
+            possible_matches,
+            translation_length,
+            reference_length,
+        ) = self._corpus_bleu(
+            reference_corpus=references,
+            translation_corpus=translation,
+            matches_by_order=matches,
+            possible_matches_by_order=possible_matches,
+            translation_length=translation_length,
+            reference_length=reference_length,
+            max_order=self.max_order,
+            smooth=self.smooth,
+        )
+        return (
+            tf.constant(bleu_score, dtype=self.dtype),
+            tf.constant(matches, dtype=self.dtype),
+            tf.constant(possible_matches, dtype=self.dtype),
+            tf.constant(translation_length, dtype=self.dtype),
+            tf.constant(reference_length, dtype=self.dtype),
+        )
+
+    def update_state(self, y_true, y_pred, sample_weight=None):
+        def validate_and_fix_rank(inputs, tensor_name, base_rank=0):
+            if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
+                inputs = tf.convert_to_tensor(inputs)
+
+            if inputs.shape.rank == base_rank:
+                return inputs[tf.newaxis]
+            elif inputs.shape.rank == base_rank + 1:
+                return inputs
+            elif inputs.shape.rank == base_rank + 2:
+                if tf.shape(inputs)[-1] != 1:
+                    raise ValueError(
+                        f"{tensor_name} is of rank {input.shape.rank}. The "
+                        f"last dimension must be of size 1."
+                    )
+                return tf.squeeze(inputs, axis=-1)
+            else:
+                raise ValueError(
+                    f"{tensor_name} must be of rank {base_rank}, {base_rank+1} "
+                    f"or {base_rank+2}. Found rank: {inputs.shape.rank}"
+                )
+
+        y_true = validate_and_fix_rank(y_true, "y_true", 1)
+        y_pred = validate_and_fix_rank(y_pred, "y_pred", 0)
+
+        # Tokenize the inputs.
+        y_true = self._tokenizer(y_true)
+        y_pred = self._tokenizer(y_pred)
+
+        (
+            bleu_score,
+            matches,
+            possible_matches,
+            translation_length,
+            reference_length,
+        ) = tf.py_function(
+            func=self._calculate_bleu_score,
+            inp=[y_true, y_pred],
+            Tout=[self.dtype, self.dtype, self.dtype, self.dtype, self.dtype],
+        )
+
+        self._matches.assign(matches)
+        self._possible_matches.assign(possible_matches)
+        self._translation_length.assign(translation_length)
+        self._reference_length.assign(reference_length)
+        self._bleu.assign(bleu_score)
+
+    def result(self):
+        return self._bleu
+
+    def reset_state(self):
+        self._matches.assign(
+            tf.zeros(shape=(self.max_order,), dtype=self.dtype)
+        )
+        self._possible_matches.assign(
+            tf.zeros(shape=(self.max_order,), dtype=self.dtype)
+        )
+        self._translation_length.assign(0.0)
+        self._reference_length.assign(0.0)
+        self._bleu.assign(0.0)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "tokenizer": self.tokenizer,
+                "max_order": self.max_order,
+                "smooth": self.smooth,
+            }
+        )
+        return config
diff --git a/keras_nlp/metrics/rouge_base.py b/keras_nlp/metrics/rouge_base.py
index bd9f6aa6..51e46486 100644
--- a/keras_nlp/metrics/rouge_base.py
+++ b/keras_nlp/metrics/rouge_base.py
@@ -48,6 +48,9 @@ class RougeBase(keras.metrics.Metric):
             not specified, it defaults to tf.float32.
         name: string. Name of the metric instance.
         **kwargs: Other keyword arguments.
+
+    References:
+        - [Lin et al., 2004](https://aclanthology.org/W04-1013/)
     """
 
     def __init__(
diff --git a/keras_nlp/metrics/rouge_l.py b/keras_nlp/metrics/rouge_l.py
index a4c9a258..a9d67633 100644
--- a/keras_nlp/metrics/rouge_l.py
+++ b/keras_nlp/metrics/rouge_l.py
@@ -38,6 +38,9 @@ class RougeL(RougeBase):
         name: string. Name of the metric instance.
         **kwargs: Other keyword arguments.
 
+    References:
+        - [Lin et al., 2004](https://aclanthology.org/W04-1013/)
+
     Examples:
 
     1. Various Input Types.
diff --git a/keras_nlp/metrics/rouge_n.py b/keras_nlp/metrics/rouge_n.py
index b7a1522a..68a9884d 100644
--- a/keras_nlp/metrics/rouge_n.py
+++ b/keras_nlp/metrics/rouge_n.py
@@ -40,6 +40,9 @@ class RougeN(RougeBase):
         name: string. Name of the metric instance.
         **kwargs: Other keyword arguments.
 
+    References:
+        - [Lin et al., 2004](https://aclanthology.org/W04-1013/)
+
     Examples:
 
     1. Various Input Types.
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index a7899463..75006ecb 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -109,14 +109,10 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
             tokenization.
         strip_accents: If true, all accent marks will be removed from text
             before tokenization.
-        split: If true, input will be split according to `split_pattern`
-            and `keep_pattern`. If false, input should be split before calling
-            the layer.
-        split_pattern: A regex pattern to match delimiters to split. By default,
-            all whitespace and punctuation marks will be split on.
-        keep_pattern: A regex pattern of delimiters contained in the
-            `split_pattern` of delimeters that should be kept as independent
-            tokens. By default, all punctuation marks will be kept as tokens.
+        split: If true, input will be split on whitespace and punctuation
+            marks, and all punctuation marks will be kept as tokens. If false,
+            input should be split ("pre-tokenized") before calling the
+            tokenizer, and passed as a dense or ragged tensor of whole words.
         suffix_indicator: The characters prepended to a wordpiece to indicate
             that it is a suffix to another subword.
         oov_token: The string value to substitute for an unknown token. It
@@ -160,14 +156,13 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
     'the quick brown fox .'
 
     Custom splitting.
-    >>> vocab = ["[UNK]", "fox", ","]
-    >>> inputs = ["fox,,fox,fox"]
-    >>> keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab,
-    ...     split_pattern=",", keep_pattern=",", dtype='string')(inputs)
-    <tf.RaggedTensor [[b'fox', b',', b',', b'fox', b',', b'fox']]>
-    >>> keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab,
-    ...     split_pattern=",", keep_pattern="", dtype='string')(inputs)
-    <tf.RaggedTensor [[b'fox', b'fox', b'fox']]>
+    >>> vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
+    >>> inputs = ["The$quick$brown$fox"]
+    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab,
+    ...     split=False, 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']]>
     """
 
     def __init__(
@@ -177,8 +172,6 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
         lowercase: bool = True,
         strip_accents: bool = True,
         split: bool = True,
-        split_pattern: str = None,
-        keep_pattern: str = None,
         suffix_indicator: str = "##",
         oov_token: str = "[UNK]",
         **kwargs,
@@ -211,18 +204,10 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
         if oov_token is None:
             raise ValueError("`oov_token` cannot be None.")
 
-        if split_pattern is None:
-            split_pattern = WHITESPACE_AND_PUNCTUATION_REGEX
-
-        if keep_pattern is None:
-            keep_pattern = PUNCTUATION_REGEX
-
         self.sequence_length = sequence_length
         self.lowercase = lowercase
         self.strip_accents = strip_accents
         self.split = split
-        self.split_pattern = split_pattern
-        self.keep_pattern = keep_pattern
         self.suffix_indicator = suffix_indicator
         self.oov_token = oov_token
 
@@ -275,8 +260,6 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
                 "lowercase": self.lowercase,
                 "strip_accents": self.strip_accents,
                 "split": self.split,
-                "split_pattern": self.split_pattern,
-                "keep_pattern": self.keep_pattern,
                 "suffix_indicator": self.suffix_indicator,
                 "oov_token": self.oov_token,
             }
@@ -284,12 +267,12 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
         return config
 
     def tokenize(self, inputs):
-        # Check if Input is Scalar or Not
         if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
             inputs = tf.convert_to_tensor(inputs)
-        scalar_input = tf.convert_to_tensor(inputs).shape.rank == 0
+        scalar_input = inputs.shape.rank == 0
         if scalar_input:
             inputs = tf.expand_dims(inputs, 0)
+
         # Optionally normalize and split inputs.
         if self.lowercase:
             inputs = tf_text.case_fold_utf8(inputs)
@@ -301,8 +284,8 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
         if self.split:
             inputs = tf_text.regex_split(
                 inputs,
-                delim_regex_pattern=self.split_pattern,
-                keep_delim_regex_pattern=self.keep_pattern,
+                delim_regex_pattern=WHITESPACE_AND_PUNCTUATION_REGEX,
+                keep_delim_regex_pattern=PUNCTUATION_REGEX,
             )
 
         # Apply word piece and coerce shape for outputs.
@@ -311,6 +294,7 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
         # split words and one for split subwords). We will collapse to a single
         # ragged dimension which is a better out of box default.
         tokens = tokens.merge_dims(-2, -1)
+
         # Convert to a dense output if `sequence_length` is set.
         if self.sequence_length:
             output_shape = tokens.shape.as_list()
diff --git a/keras_nlp/utils/tensor_utils.py b/keras_nlp/utils/tensor_utils.py
index 26fc815f..bcab0f24 100644
--- a/keras_nlp/utils/tensor_utils.py
+++ b/keras_nlp/utils/tensor_utils.py
@@ -25,16 +25,11 @@ def _decode_strings_to_utf8(inputs):
         return [_decode_strings_to_utf8(x) for x in inputs]
 
 
-def tensor_to_string_list(inputs):
-    """Detokenize and convert tensor to nested lists of python strings.
-
-    This is a convenience method which converts each byte string to a python
-    string.
+def tensor_to_list(inputs):
+    """Converts a tensor to nested lists.
 
     Args:
         inputs: Input tensor, or dict/list/tuple of input tensors.
-        *args: Additional positional arguments.
-        **kwargs: Additional keyword arguments.
     """
     if not isinstance(inputs, (tf.RaggedTensor, tf.Tensor)):
         inputs = tf.convert_to_tensor(inputs)
@@ -44,4 +39,17 @@ def tensor_to_string_list(inputs):
         list_outputs = inputs.numpy()
         if inputs.shape.rank != 0:
             list_outputs = list_outputs.tolist()
+    return list_outputs
+
+
+def tensor_to_string_list(inputs):
+    """Detokenize and convert tensor to nested lists of python strings.
+
+    This is a convenience method which converts each byte string to a python
+    string.
+
+    Args:
+        inputs: Input tensor, or dict/list/tuple of input tensors.
+    """
+    list_outputs = tensor_to_list(inputs)
     return _decode_strings_to_utf8(list_outputs)
diff --git a/keras_nlp/utils/text_generation.py b/keras_nlp/utils/text_generation.py
index fbae45ce..3b9a088b 100644
--- a/keras_nlp/utils/text_generation.py
+++ b/keras_nlp/utils/text_generation.py
@@ -15,6 +15,8 @@
 """Text generation utilities."""
 
 import tensorflow as tf
+from absl import logging
+from tensorflow import keras
 
 
 def validate_prompt(prompt):
@@ -98,14 +100,14 @@ def greedy_search(
     END_ID = 2
 
     # Create a dummy model to predict the next token.
-    model = tf.keras.Sequential(
+    model = keras.Sequential(
         [
-            tf.keras.Input(shape=[None]),
-            tf.keras.layers.Embedding(
+            keras.Input(shape=[None]),
+            keras.layers.Embedding(
                 input_dim=VOCAB_SIZE,
                 output_dim=FEATURE_SIZE,
             ),
-            tf.keras.layers.Dense(VOCAB_SIZE, activation="softmax"),
+            keras.layers.Dense(VOCAB_SIZE, activation="softmax"),
         ]
     )
 
@@ -155,9 +157,7 @@ def greedy_search(
             prompt, max_length, end_token_id, pad_token_id
         )
 
-    if input_is_1d:
-        return tf.squeeze(prompt)
-    return prompt
+    return tf.squeeze(prompt) if input_is_1d else prompt
 
 
 def beam_search(
@@ -211,14 +211,14 @@ def beam_search(
     END_ID = 2
 
     # Create a dummy model to predict the next token.
-    model = tf.keras.Sequential(
+    model = keras.Sequential(
         [
-            tf.keras.Input(shape=[None]),
-            tf.keras.layers.Embedding(
+            keras.Input(shape=[None]),
+            keras.layers.Embedding(
                 input_dim=VOCAB_SIZE,
                 output_dim=FEATURE_SIZE,
             ),
-            tf.keras.layers.Dense(VOCAB_SIZE, activation="softmax"),
+            keras.layers.Dense(VOCAB_SIZE, activation="softmax"),
         ]
     )
 
@@ -260,54 +260,50 @@ def beam_search(
     validate_token_probability_fn(token_probability_fn, prompt)
 
     batch_size, length = prompt.shape
-    if length < max_length:
-        # Initialize beam.
-        beams = tf.expand_dims(prompt, 1)
-        beams_prob = tf.zeros([batch_size, 1])
-        i = length
-        while i < max_length:
-            beam_size = beams.shape[1]
-            beam_preds = []
-            for j in range(beam_size):
-                preds = token_probability_fn(beams[:, j, :])
-                if from_logits:
-                    preds = tf.keras.activations.softmax(preds, axis=-1)
-                beam_preds.append(preds)
-            stacked_preds = tf.stack(beam_preds, axis=1)
-            vocab_size = stacked_preds.shape[2]
-            logits = tf.reshape(
-                stacked_preds, [batch_size, beam_size * vocab_size]
-            )
-            probs = tf.math.log(logits) + tf.repeat(
-                beams_prob, repeats=vocab_size, axis=1
-            )
-            num_beams = min(beam_size * vocab_size, num_beams)
-            candidate_prob, candidate_indexes = tf.math.top_k(
-                probs, k=num_beams, sorted=False
-            )
-            candidate_beam_indexes = candidate_indexes // vocab_size
-            next_token = candidate_indexes % vocab_size
-
-            beams = tf.gather(
-                beams, candidate_beam_indexes, axis=1, batch_dims=1
-            )
-            beams = tf.concat([beams, next_token[..., tf.newaxis]], axis=-1)
-            beams_prob = candidate_prob
-            i += 1
-        # Get the beam with the maximum probability.
-        max_indexes = tf.math.argmax(beams_prob, axis=-1)
-        max_beams = tf.gather(
-            beams, max_indexes[:, tf.newaxis], axis=1, batch_dims=1
+    if length >= max_length:
+        return tf.squeeze(prompt) if input_is_1d else prompt
+
+    # Initialize beam.
+    beams = tf.expand_dims(prompt, 1)
+    beams_prob = tf.zeros([batch_size, 1])
+    i = length
+    while i < max_length:
+        beam_size = beams.shape[1]
+        beam_preds = []
+        for j in range(beam_size):
+            preds = token_probability_fn(beams[:, j, :])
+            if from_logits:
+                preds = keras.activations.softmax(preds, axis=-1)
+            beam_preds.append(preds)
+        stacked_preds = tf.stack(beam_preds, axis=1)
+        vocab_size = stacked_preds.shape[2]
+        logits = tf.reshape(stacked_preds, [batch_size, beam_size * vocab_size])
+        probs = tf.math.log(logits) + tf.repeat(
+            beams_prob, repeats=vocab_size, axis=1
+        )
+        num_beams = min(beam_size * vocab_size, num_beams)
+        candidate_prob, candidate_indexes = tf.math.top_k(
+            probs, k=num_beams, sorted=False
         )
-        prompt = tf.squeeze(max_beams)
+        candidate_beam_indexes = candidate_indexes // vocab_size
+        next_token = candidate_indexes % vocab_size
+
+        beams = tf.gather(beams, candidate_beam_indexes, axis=1, batch_dims=1)
+        beams = tf.concat([beams, next_token[..., tf.newaxis]], axis=-1)
+        beams_prob = candidate_prob
+        i += 1
+    # Get the beam with the maximum probability.
+    max_indexes = tf.math.argmax(beams_prob, axis=-1)
+    max_beams = tf.gather(
+        beams, max_indexes[:, tf.newaxis], axis=1, batch_dims=1
+    )
+    prompt = tf.squeeze(max_beams)
 
     if end_token_id is not None:
         prompt = mask_tokens_after_end_token(
             prompt, max_length, end_token_id, pad_token_id
         )
-    if input_is_1d:
-        return tf.squeeze(prompt)
-    return prompt
+    return tf.squeeze(prompt) if input_is_1d else prompt
 
 
 def random_search(
@@ -358,14 +354,14 @@ def random_search(
     END_ID = 2
 
     # Create a dummy model to predict the next token.
-    model = tf.keras.Sequential(
+    model = keras.Sequential(
         [
-            tf.keras.Input(shape=[None]),
-            tf.keras.layers.Embedding(
+            keras.Input(shape=[None]),
+            keras.layers.Embedding(
                 input_dim=VOCAB_SIZE,
                 output_dim=FEATURE_SIZE,
             ),
-            tf.keras.layers.Dense(VOCAB_SIZE, activation="softmax"),
+            keras.layers.Dense(VOCAB_SIZE, activation="softmax"),
         ]
     )
 
@@ -405,7 +401,7 @@ def random_search(
         # If the prompt has reached our desired length, exit while loop.
         pred = token_probability_fn(prompt)
         if from_logits:
-            pred = tf.keras.activations.softmax(pred, axis=-1)
+            pred = keras.activations.softmax(pred, axis=-1)
         next_token = tf.cast(
             tf.random.categorical(tf.math.log(pred), 1, seed=seed),
             dtype=prompt.dtype,
@@ -418,9 +414,7 @@ def random_search(
         prompt = mask_tokens_after_end_token(
             prompt, max_length, end_token_id, pad_token_id
         )
-    if input_is_1d:
-        return tf.squeeze(prompt)
-    return prompt
+    return tf.squeeze(prompt) if input_is_1d else prompt
 
 
 def top_k_search(
@@ -474,14 +468,14 @@ def top_k_search(
     END_ID = 2
 
     # Create a dummy model to predict the next token.
-    model = tf.keras.Sequential(
+    model = keras.Sequential(
         [
-            tf.keras.Input(shape=[None]),
-            tf.keras.layers.Embedding(
+            keras.Input(shape=[None]),
+            keras.layers.Embedding(
                 input_dim=VOCAB_SIZE,
                 output_dim=FEATURE_SIZE,
             ),
-            tf.keras.layers.Dense(VOCAB_SIZE, activation="softmax"),
+            keras.layers.Dense(VOCAB_SIZE, activation="softmax"),
         ]
     )
 
@@ -510,8 +504,6 @@ def top_k_search(
             "tf.function or run `tf.config.run_functions_eagerly(True)` to run "
             "tf.function in eager mode."
         )
-    if k <= 0:
-        raise ValueError(f"`k` should be strictly positive. Received: `k={k}`.")
 
     prompt = validate_prompt(prompt)
     input_is_1d = prompt.shape.rank == 1
@@ -519,14 +511,23 @@ def top_k_search(
         prompt = prompt[tf.newaxis, :]
     validate_token_probability_fn(token_probability_fn, prompt)
 
+    if k <= 0:
+        raise ValueError(f"`k` should be strictly positive. Received: `k={k}`.")
+    # If k is greater than the vocabulary size, use the entire vocabulary.
+    pred = token_probability_fn(prompt)
+    if k > pred.shape[1]:
+        logging.warning(
+            f"`k` larger than vocabulary size={pred.shape[1]}."
+            f"Setting `k` to vocabulary size. Received: `k={k}`."
+        )
+        k = pred.shape[1]
+
     i = prompt.shape[1]
     while i < max_length:
-        # If the prompt has reached our desired length, exit while loop.
         pred = token_probability_fn(prompt)
         if from_logits:
-            pred = tf.keras.activations.softmax(pred, axis=-1)
-        # If k is greater than the vocabulary size, use the entire vocabulary.
-        k = min(k, pred.shape[1])
+            pred = keras.activations.softmax(pred, axis=-1)
+
         # Filter out top-k tokens.
         top_k_pred, top_k_indices = tf.math.top_k(pred, k=k, sorted=False)
         # Sample the next token from the probability distribution.
@@ -544,9 +545,7 @@ def top_k_search(
         prompt = mask_tokens_after_end_token(
             prompt, max_length, end_token_id, pad_token_id
         )
-    if input_is_1d:
-        return tf.squeeze(prompt)
-    return prompt
+    return tf.squeeze(prompt) if input_is_1d else prompt
 
 
 def top_p_search(
@@ -602,14 +601,14 @@ def top_p_search(
     END_ID = 2
 
     # Create a dummy model to predict the next token.
-    model = tf.keras.Sequential(
+    model = keras.Sequential(
         [
-            tf.keras.Input(shape=[None]),
-            tf.keras.layers.Embedding(
+            keras.Input(shape=[None]),
+            keras.layers.Embedding(
                 input_dim=VOCAB_SIZE,
                 output_dim=FEATURE_SIZE,
             ),
-            tf.keras.layers.Dense(VOCAB_SIZE, activation="softmax"),
+            keras.layers.Dense(VOCAB_SIZE, activation="softmax"),
         ]
     )
 
@@ -654,7 +653,7 @@ def top_p_search(
         # If the prompt has reached our desired length, exit while loop.
         pred = token_probability_fn(prompt)
         if from_logits:
-            pred = tf.keras.activations.softmax(pred, axis=-1)
+            pred = keras.activations.softmax(pred, axis=-1)
         # Sort preds in descending order.
         sorted_preds, sorted_indices = tf.math.top_k(
             pred, k=pred.shape[1], sorted=True
@@ -688,6 +687,4 @@ def top_p_search(
         prompt = mask_tokens_after_end_token(
             prompt, max_length, end_token_id, pad_token_id
         )
-    if input_is_1d:
-        return tf.squeeze(prompt)
-    return prompt
+    return tf.squeeze(prompt) if input_is_1d else prompt


Test patch:
diff --git a/keras_nlp/layers/mlm_head_test.py b/keras_nlp/layers/mlm_head_test.py
index 0271d915..699eda07 100644
--- a/keras_nlp/layers/mlm_head_test.py
+++ b/keras_nlp/layers/mlm_head_test.py
@@ -117,9 +117,7 @@ class MLMHeadTest(tf.test.TestCase):
         )
         label_data = tf.random.uniform(shape=(4, 5), maxval=100, dtype="int32")
 
-        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
-            from_logits=False
-        )
+        loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=False)
         optimizer = keras.optimizers.Adam()
         with tf.GradientTape() as tape:
             model((token_data, position_data))
diff --git a/keras_nlp/layers/position_embedding_test.py b/keras_nlp/layers/position_embedding_test.py
index 7d2741be..bd07e9b2 100644
--- a/keras_nlp/layers/position_embedding_test.py
+++ b/keras_nlp/layers/position_embedding_test.py
@@ -36,7 +36,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         test_layer = position_embedding.PositionEmbedding(
             sequence_length=sequence_length
         )
-        input_tensor = tf.keras.Input(shape=(sequence_length, feature_size))
+        input_tensor = keras.Input(shape=(sequence_length, feature_size))
         output_tensor = test_layer(input_tensor)
 
         # When using static position embedding shapes, the output is expected
@@ -53,7 +53,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         test_layer = position_embedding.PositionEmbedding(
             sequence_length=sequence_length
         )
-        input_tensor = tf.keras.Input(
+        input_tensor = keras.Input(
             shape=(feature_size, sequence_length, feature_size)
         )
         output_tensor = test_layer(input_tensor)
@@ -77,7 +77,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         test_layer = position_embedding.PositionEmbedding(
             sequence_length=sequence_length, dtype="float16"
         )
-        input_tensor = tf.keras.Input(shape=(sequence_length, feature_size))
+        input_tensor = keras.Input(shape=(sequence_length, feature_size))
         output_tensor = test_layer(input_tensor)
 
         # When using static position embedding shapes, the output is expected
@@ -94,7 +94,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
             sequence_length=max_sequence_length
         )
         # Create a 3-dimensional input (the first dimension is implicit).
-        input_tensor = tf.keras.Input(shape=(None, feature_size))
+        input_tensor = keras.Input(shape=(None, feature_size))
         output_tensor = test_layer(input_tensor)
 
         # When using dynamic position embedding shapes, the output is expected
@@ -110,7 +110,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
             sequence_length=max_sequence_length
         )
         # Create a 4-dimensional input (the first dimension is implicit).
-        input_tensor = tf.keras.Input(shape=(None, None, feature_size))
+        input_tensor = keras.Input(shape=(None, None, feature_size))
         output_tensor = test_layer(input_tensor)
 
         # When using dynamic position embedding shapes, the output is expected
@@ -126,10 +126,10 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
             sequence_length=max_sequence_length
         )
         # Create a 3-dimensional input (the first dimension is implicit).
-        input_tensor = tf.keras.Input(shape=(None, feature_size))
+        input_tensor = keras.Input(shape=(None, feature_size))
         output_tensor = test_layer(input_tensor)
 
-        model = tf.keras.Model(input_tensor, output_tensor)
+        model = keras.Model(input_tensor, output_tensor)
 
         # Create input data that is shorter than max_sequence_length, which
         # should trigger a down-slice.
@@ -151,7 +151,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
             sequence_length=max_sequence_length,
             initializer=custom_init,
         )
-        inputs = tf.keras.Input(shape=(max_sequence_length, feature_size))
+        inputs = keras.Input(shape=(max_sequence_length, feature_size))
         outputs = test_layer(inputs)
         model = keras.Model(inputs=inputs, outputs=outputs)
 
@@ -176,11 +176,11 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
             initializer=custom_init,
         )
         # Create a 3-dimensional ragged input (the first dimension is implicit).
-        input_tensor = tf.keras.Input(
+        input_tensor = keras.Input(
             shape=(None, feature_size), dtype=tf.float32, ragged=True
         )
         output_tensor = test_layer(input_tensor)
-        model = tf.keras.Model(input_tensor, output_tensor)
+        model = keras.Model(input_tensor, output_tensor)
 
         input_data = tf.ragged.constant(
             [
@@ -213,11 +213,11 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
             initializer=custom_init,
         )
         # Create a 4-dimensional ragged input (the first dimension is implicit).
-        input_tensor = tf.keras.Input(
+        input_tensor = keras.Input(
             shape=(None, None, feature_size), dtype=tf.float32, ragged=True
         )
         output_tensor = test_layer(input_tensor)
-        model = tf.keras.Model(input_tensor, output_tensor)
+        model = keras.Model(input_tensor, output_tensor)
 
         input_data = tf.ragged.constant(
             [
@@ -256,7 +256,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         test_layer = position_embedding.PositionEmbedding(
             sequence_length=max_sequence_length
         )
-        inputs = tf.keras.Input(shape=(max_sequence_length, feature_size))
+        inputs = keras.Input(shape=(max_sequence_length, feature_size))
         outputs = test_layer(inputs)
         model = keras.Model(inputs=inputs, outputs=outputs)
 
@@ -308,7 +308,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         test_layer = position_embedding.PositionEmbedding(
             sequence_length=max_sequence_length
         )
-        inputs = tf.keras.Input(shape=(max_sequence_length, feature_size))
+        inputs = keras.Input(shape=(max_sequence_length, feature_size))
         outputs = test_layer(inputs)
         model = keras.Model(inputs=inputs, outputs=outputs)
 
diff --git a/keras_nlp/layers/token_and_position_embedding_test.py b/keras_nlp/layers/token_and_position_embedding_test.py
index cfaadef6..2639d60b 100644
--- a/keras_nlp/layers/token_and_position_embedding_test.py
+++ b/keras_nlp/layers/token_and_position_embedding_test.py
@@ -68,11 +68,11 @@ class TokenAndPositionEmbeddingTest(tf.test.TestCase):
         )
         # Create a 2-dimensional ragged input
         # (the first dimension is implicit).
-        input_tensor = tf.keras.Input(
+        input_tensor = keras.Input(
             shape=(sequence_length,), dtype=tf.float32, ragged=True
         )
         output_tensor = test_layer(input_tensor)
-        model = tf.keras.Model(input_tensor, output_tensor)
+        model = keras.Model(input_tensor, output_tensor)
 
         input_data = tf.ragged.constant(
             [
@@ -112,9 +112,9 @@ class TokenAndPositionEmbeddingTest(tf.test.TestCase):
         )
         # Create a 2-dimensional input
         # (the first dimension is implicit).
-        inputs = tf.keras.Input(shape=(sequence_length,), dtype="int32")
+        inputs = keras.Input(shape=(sequence_length,), dtype="int32")
         outputs = test_layer(inputs)
-        model = tf.keras.Model(inputs, outputs)
+        model = keras.Model(inputs, outputs)
 
         input_data = tf.ones((2, sequence_length), dtype="int32")
         expected_output_data = tf.ones((2, sequence_length, embedding_dim)) * 2
diff --git a/keras_nlp/metrics/bleu_test.py b/keras_nlp/metrics/bleu_test.py
new file mode 100644
index 00000000..a2d314c6
--- /dev/null
+++ b/keras_nlp/metrics/bleu_test.py
@@ -0,0 +1,268 @@
+# 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 Bleu."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.metrics import Bleu
+from keras_nlp.tokenizers import ByteTokenizer
+
+
+class BleuTest(tf.test.TestCase):
+    def test_initialization(self):
+        bleu = Bleu()
+        result = bleu.result()
+
+        self.assertEqual(result, 0.0)
+
+    def test_scalar_input(self):
+        bleu = Bleu(smooth=True)
+        y_true = [
+            "He eats a sweet apple.",
+            "He is eating a tasty apple, isn't he?",
+        ]
+        y_pred = "He He He eats sweet apple which is a fruit."
+
+        bleu_val = bleu(y_true, y_pred)
+        self.assertAlmostEqual(bleu_val.numpy(), 0.212, delta=1e-3)
+
+    def test_1d_list_input(self):
+        bleu = Bleu()
+        y_true = [
+            ["He eats a sweet apple."],
+            ["Silicon Valley is one of my favourite shows!"],
+        ]
+        y_pred = [
+            "He He He eats sweet apple which is a fruit.",
+            "I love Silicon Valley, it's one of my favourite shows.",
+        ]
+
+        bleu_val = bleu(y_true, y_pred)
+        self.assertAlmostEqual(bleu_val.numpy(), 0.243, delta=1e-3)
+
+    def test_2d_list_input(self):
+        bleu = Bleu()
+        y_true = [
+            [["He eats a sweet apple."]],
+            [["Silicon Valley is one of my favourite shows!"]],
+        ]
+        y_pred = [
+            ["He He He eats sweet apple which is a fruit."],
+            ["I love Silicon Valley, it's one of my favourite shows."],
+        ]
+
+        bleu_val = bleu(y_true, y_pred)
+        self.assertAlmostEqual(bleu_val.numpy(), 0.243, delta=1e-3)
+
+    def test_1d_tensor_input(self):
+        bleu = Bleu()
+        y_true = tf.ragged.constant(
+            [
+                ["He eats a sweet apple."],
+                ["Silicon Valley is one of my favourite shows!"],
+            ]
+        )
+        y_pred = tf.constant(
+            [
+                "He He He eats sweet apple which is a fruit.",
+                "I love Silicon Valley, it's one of my favourite shows.",
+            ]
+        )
+
+        bleu_val = bleu(y_true, y_pred)
+        self.assertAlmostEqual(bleu_val.numpy(), 0.243, delta=1e-3)
+
+    def test_2d_tensor_input(self):
+        bleu = Bleu()
+        y_true = tf.constant(
+            [
+                [["He eats a sweet apple."]],
+                [["Silicon Valley is one of my favourite shows!"]],
+            ]
+        )
+        y_pred = tf.constant(
+            [
+                ["He He He eats sweet apple which is a fruit."],
+                ["I love Silicon Valley, it's one of my favourite shows."],
+            ]
+        )
+
+        bleu_val = bleu(y_true, y_pred)
+        self.assertAlmostEqual(bleu_val.numpy(), 0.243, delta=1e-3)
+
+    def test_custom_tokenizer(self):
+        byte_tokenizer = ByteTokenizer()
+        bleu = Bleu(tokenizer=byte_tokenizer)
+        y_true = tf.ragged.constant(
+            [
+                ["He eats a sweet apple."],
+                ["Silicon Valley is one of my favourite shows!"],
+            ]
+        )
+        y_pred = tf.constant(
+            [
+                "He He He eats sweet apple which is a fruit.",
+                "I love Silicon Valley, it's one of my favourite shows.",
+            ]
+        )
+
+        bleu_val = bleu(y_true, y_pred)
+        self.assertAlmostEqual(bleu_val.numpy(), 0.609, delta=1e-3)
+
+    def test_different_order(self):
+        bleu = Bleu(max_order=5)
+        y_true = tf.ragged.constant(
+            [
+                ["He eats a sweet apple."],
+                ["Silicon Valley is one of my favourite shows!"],
+            ]
+        )
+        y_pred = tf.constant(
+            [
+                "He He He eats sweet apple which is a fruit.",
+                "I love Silicon Valley, it's one of my favourite shows.",
+            ]
+        )
+
+        bleu_val = bleu(y_true, y_pred)
+        self.assertAlmostEqual(bleu_val.numpy(), 0.188, delta=1e-3)
+
+    def test_model_compile(self):
+        inputs = keras.Input(shape=(), dtype="string")
+        outputs = tf.identity(inputs)
+        model = keras.Model(inputs, outputs)
+
+        model.compile(metrics=[Bleu()])
+
+        x = tf.constant(
+            [
+                "He He He eats sweet apple which is a fruit.",
+                "I love Silicon Valley, it's one of my favourite shows.",
+            ]
+        )
+        y = tf.constant(
+            [
+                ["He eats a sweet apple."],
+                ["Silicon Valley is one of my favourite shows!"],
+            ]
+        )
+
+        output = model.evaluate(x, y, return_dict=True)
+        self.assertAlmostEqual(output["bleu"], 0.243, delta=1e-3)
+
+    def test_reset_state(self):
+        bleu = Bleu()
+        y_true = tf.ragged.constant(
+            [
+                ["He eats a sweet apple."],
+                ["Silicon Valley is one of my favourite shows!"],
+            ]
+        )
+        y_pred = tf.constant(
+            [
+                "He He He eats sweet apple which is a fruit.",
+                "I love Silicon Valley, it's one of my favourite shows.",
+            ]
+        )
+
+        bleu.update_state(y_true, y_pred)
+        bleu_val = bleu.result()
+        self.assertNotEqual(bleu_val.numpy(), 0.0)
+
+        bleu.reset_state()
+        bleu_val = bleu.result()
+        self.assertEqual(bleu_val, 0.0)
+
+    def test_update_state(self):
+        bleu = Bleu()
+        y_true_1 = tf.ragged.constant(
+            [
+                ["He eats a sweet apple."],
+                ["Silicon Valley is one of my favourite shows!"],
+            ]
+        )
+        y_pred_1 = tf.constant(
+            [
+                "He He He eats sweet apple which is a fruit.",
+                "I love Silicon Valley, it's one of my favourite shows.",
+            ]
+        )
+
+        bleu.update_state(y_true_1, y_pred_1)
+        bleu_val = bleu.result()
+        self.assertAlmostEqual(bleu_val.numpy(), 0.243, delta=1e-3)
+
+        y_true_2 = tf.constant(["Virat Kohli is the GOAT."])
+        y_pred_2 = tf.constant("Virat Kohli is the greatest of all time!")
+
+        bleu.update_state(y_true_2, y_pred_2)
+        bleu_val = bleu.result()
+        self.assertAlmostEqual(bleu_val.numpy(), 0.26, delta=1e-3)
+
+    def test_merge_state_normalize(self):
+        bleu_1 = Bleu(smooth=True)
+        bleu_2 = Bleu(smooth=True)
+
+        y_true_1 = tf.ragged.constant(
+            [
+                ["He eats a sweet apple."],
+                ["Silicon Valley is one of my favourite shows!"],
+            ]
+        )
+        y_pred_1 = tf.constant(
+            [
+                "He He He eats sweet apple which is a fruit.",
+                "I love Silicon Valley, it's one of my favourite shows.",
+            ]
+        )
+
+        y_true_2 = tf.constant(["Virat Kohli is the GOAT."])
+        y_pred_2 = tf.constant("Virat Kohli is the greatest of all time!")
+
+        y_true_3 = tf.constant([["Watching Test cricket is so much fun."]])
+        y_pred_3 = tf.constant(["Test is the best format in cricket."])
+
+        bleu_1.update_state(y_true_1, y_pred_1)
+        bleu_1.update_state(y_true_2, y_pred_2)
+        bleu_val = bleu_1.result()
+        self.assertAlmostEqual(bleu_val.numpy(), 0.293, delta=1e-3)
+
+        bleu_2.update_state(y_true_3, y_pred_3)
+        bleu_val = bleu_2.result()
+        self.assertAlmostEqual(bleu_val.numpy(), 0.202, delta=1e-3)
+
+        merged_bleu = Bleu(smooth=True)
+        merged_bleu.merge_state([bleu_1, bleu_2])
+        bleu_val = merged_bleu.result()
+        self.assertAlmostEqual(bleu_val.numpy(), 0.495, delta=1e-3)
+
+    def test_get_config(self):
+        byte_tokenizer = ByteTokenizer()
+        bleu = Bleu(
+            tokenizer=byte_tokenizer,
+            max_order=8,
+            smooth=True,
+            dtype=tf.float64,
+            name="bleu_test",
+        )
+
+        config = bleu.get_config()
+        expected_config_subset = {
+            "tokenizer": byte_tokenizer,
+            "max_order": 8,
+            "smooth": True,
+        }
+        self.assertEqual(config, {**config, **expected_config_subset})
diff --git a/keras_nlp/tokenizers/unicode_character_tokenizer_test.py b/keras_nlp/tokenizers/unicode_character_tokenizer_test.py
index 334a8139..098b5fdc 100644
--- a/keras_nlp/tokenizers/unicode_character_tokenizer_test.py
+++ b/keras_nlp/tokenizers/unicode_character_tokenizer_test.py
@@ -249,9 +249,9 @@ class UnicodeCharacterTokenizerTest(tf.test.TestCase):
             ["ninja", "samurai", "▀▁▂▃", "keras", "tensorflow"]
         )
         tokenizer = UnicodeCharacterTokenizer()
-        inputs = tf.keras.Input(dtype="string", shape=())
+        inputs = keras.Input(dtype="string", shape=())
         outputs = tokenizer.detokenize(tokenizer.tokenize(inputs))
-        model = tf.keras.Model(inputs, outputs)
+        model = keras.Model(inputs, outputs)
         model_output = model(input_data)
         self.assertAllEqual(
             model_output,
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
index fc3e6237..1d5a62e8 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer_test.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
@@ -116,22 +116,6 @@ class WordPieceTokenizerTest(tf.test.TestCase):
         call_output = tokenizer(input_data)
         self.assertAllEqual(call_output, [[1, 2, 3, 4, 5]])
 
-    def test_custom_spliting(self):
-        input_data = ["a,e,i,o,u"]
-        vocab_data = ["[UNK]", "a", "e", "i", "o", "u", ","]
-        # Don't keep commas.
-        tokenizer = WordPieceTokenizer(
-            vocabulary=vocab_data, split_pattern=",", keep_pattern=""
-        )
-        call_output = tokenizer(input_data)
-        self.assertAllEqual(call_output, [[1, 2, 3, 4, 5]])
-        # Keep commas.
-        tokenizer = WordPieceTokenizer(
-            vocabulary=vocab_data, split_pattern=",", keep_pattern=","
-        )
-        call_output = tokenizer(input_data)
-        self.assertAllEqual(call_output, [[1, 6, 2, 6, 3, 6, 4, 6, 5]])
-
     def test_no_spliting(self):
         input_data = ["t o k e n", "m i s s i n g", "t o k e n"]
         vocab_data = ["[UNK]", "t o k e n"]
diff --git a/keras_nlp/utils/tensor_utils_test.py b/keras_nlp/utils/tensor_utils_test.py
index d9941f75..bdb9d728 100644
--- a/keras_nlp/utils/tensor_utils_test.py
+++ b/keras_nlp/utils/tensor_utils_test.py
@@ -14,9 +14,27 @@
 
 import tensorflow as tf
 
+from keras_nlp.utils.tensor_utils import tensor_to_list
 from keras_nlp.utils.tensor_utils import tensor_to_string_list
 
 
+class TensorToListTest(tf.test.TestCase):
+    def test_ragged_input(self):
+        input_data = tf.ragged.constant([[1, 2], [4, 5, 6]])
+        list_output = tensor_to_list(input_data)
+        self.assertAllEqual(list_output, [[1, 2], [4, 5, 6]])
+
+    def test_dense_input(self):
+        input_data = tf.constant([[1, 2], [3, 4]])
+        list_output = tensor_to_list(input_data)
+        self.assertAllEqual(list_output, [[1, 2], [3, 4]])
+
+    def test_scalar_input(self):
+        input_data = tf.constant(1)
+        list_output = tensor_to_list(input_data)
+        self.assertEqual(list_output, 1)
+
+
 class TensorToStringListTest(tf.test.TestCase):
     def test_detokenize_to_strings_for_ragged(self):
         input_data = tf.ragged.constant([["▀▁▂▃", "samurai"]])
diff --git a/keras_nlp/utils/text_generation_test.py b/keras_nlp/utils/text_generation_test.py
index 5e6f6a2c..1731fb4e 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 tensorflow import keras
 
 from keras_nlp.utils.text_generation import beam_search
 from keras_nlp.utils.text_generation import greedy_search
@@ -28,19 +29,19 @@ from keras_nlp.utils.text_generation import top_p_search
 class GreedySearchTextGenerationTest(tf.test.TestCase):
     def setUp(self):
         super().setUp()
-        vocab_size = 10
-        feature_size = 16
+        self.vocab_size = 10
+        self.feature_size = 16
 
         # Create a dummy model to predict the next token.
-        model = tf.keras.Sequential(
+        model = keras.Sequential(
             [
-                tf.keras.Input(shape=[None]),
-                tf.keras.layers.Embedding(
-                    input_dim=vocab_size,
-                    output_dim=feature_size,
+                keras.Input(shape=[None]),
+                keras.layers.Embedding(
+                    input_dim=self.vocab_size,
+                    output_dim=self.feature_size,
                 ),
-                tf.keras.layers.Dense(vocab_size),
-                tf.keras.layers.Softmax(),
+                keras.layers.Dense(self.vocab_size),
+                keras.layers.Softmax(),
             ]
         )
 
@@ -116,19 +117,19 @@ class GreedySearchTextGenerationTest(tf.test.TestCase):
 class BeamSearchTextGenerationTest(tf.test.TestCase):
     def setUp(self):
         super().setUp()
-        vocab_size = 10
-        feature_size = 16
+        self.vocab_size = 10
+        self.feature_size = 16
 
         # Create a dummy model to predict the next token.
-        model = tf.keras.Sequential(
+        model = keras.Sequential(
             [
-                tf.keras.Input(shape=[None]),
-                tf.keras.layers.Embedding(
-                    input_dim=vocab_size,
-                    output_dim=feature_size,
+                keras.Input(shape=[None]),
+                keras.layers.Embedding(
+                    input_dim=self.vocab_size,
+                    output_dim=self.feature_size,
                 ),
-                tf.keras.layers.Dense(vocab_size),
-                tf.keras.layers.Softmax(),
+                keras.layers.Dense(self.vocab_size),
+                keras.layers.Softmax(),
             ]
         )
 
@@ -159,6 +160,25 @@ class BeamSearchTextGenerationTest(tf.test.TestCase):
         )
         self.assertEquals(outputs.shape, [5])
 
+    def test_generate_with_max_length_prompt(self):
+        inputs = tf.ones(shape=(5,))
+        outputs = beam_search(
+            self.token_probability_fn,
+            inputs,
+            max_length=5,
+            num_beams=5,
+        )
+        self.assertEquals(outputs.shape, [5])
+
+        inputs = tf.ones(shape=(6,))
+        outputs = beam_search(
+            self.token_probability_fn,
+            inputs,
+            max_length=5,
+            num_beams=5,
+        )
+        self.assertEquals(outputs.shape, [6])
+
     def test_generate_with_2d_prompt(self):
         inputs = tf.constant([[1], [1]])
         outputs = beam_search(
@@ -270,19 +290,19 @@ class BeamSearchTextGenerationTest(tf.test.TestCase):
 class RandomSearchTextGenerationTest(tf.test.TestCase):
     def setUp(self):
         super().setUp()
-        vocab_size = 10
-        feature_size = 16
+        self.vocab_size = 10
+        self.feature_size = 16
 
         # Create a dummy model to predict the next token.
-        model = tf.keras.Sequential(
+        model = keras.Sequential(
             [
-                tf.keras.Input(shape=[None]),
-                tf.keras.layers.Embedding(
-                    input_dim=vocab_size,
-                    output_dim=feature_size,
+                keras.Input(shape=[None]),
+                keras.layers.Embedding(
+                    input_dim=self.vocab_size,
+                    output_dim=self.feature_size,
                 ),
-                tf.keras.layers.Dense(vocab_size),
-                tf.keras.layers.Softmax(),
+                keras.layers.Dense(self.vocab_size),
+                keras.layers.Softmax(),
             ]
         )
 
@@ -391,7 +411,7 @@ class RandomSearchTextGenerationTest(tf.test.TestCase):
 
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.keras.activations.softmax(
+            prob = keras.activations.softmax(
                 tf.constant([[1.0, 2.0, 3, 0, 4.0]])
             )
             return tf.repeat(prob, batch_size, axis=0)
@@ -420,19 +440,19 @@ class RandomSearchTextGenerationTest(tf.test.TestCase):
 class TopKSearchTextGenerationTest(tf.test.TestCase):
     def setUp(self):
         super().setUp()
-        vocab_size = 10
-        feature_size = 16
+        self.vocab_size = 10
+        self.feature_size = 16
 
         # Create a dummy model to predict the next token.
-        model = tf.keras.Sequential(
+        model = keras.Sequential(
             [
-                tf.keras.Input(shape=[None]),
-                tf.keras.layers.Embedding(
-                    input_dim=vocab_size,
-                    output_dim=feature_size,
+                keras.Input(shape=[None]),
+                keras.layers.Embedding(
+                    input_dim=self.vocab_size,
+                    output_dim=self.feature_size,
                 ),
-                tf.keras.layers.Dense(vocab_size),
-                tf.keras.layers.Softmax(),
+                keras.layers.Dense(self.vocab_size),
+                keras.layers.Softmax(),
             ]
         )
 
@@ -456,6 +476,16 @@ class TopKSearchTextGenerationTest(tf.test.TestCase):
         )
         self.assertEquals(outputs.shape, [5])
 
+    def test_generate_with_k_too_big(self):
+        inputs = tf.constant([1])
+        outputs = top_k_search(
+            self.token_probability_fn,
+            inputs,
+            max_length=5,
+            k=self.vocab_size + 1,
+        )
+        self.assertEquals(outputs.shape, [5])
+
     def test_generate_with_2d_prompt(self):
         inputs = tf.constant([[1], [1]])
         outputs = top_k_search(
@@ -571,7 +601,7 @@ class TopKSearchTextGenerationTest(tf.test.TestCase):
 
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.keras.activations.softmax(
+            prob = keras.activations.softmax(
                 tf.constant([[1.0, 2.0, 3.0, 4.0]])
             )
             return tf.repeat(prob, batch_size, axis=0)
@@ -602,19 +632,19 @@ class TopKSearchTextGenerationTest(tf.test.TestCase):
 class TopPSearchTextGenerationTest(tf.test.TestCase):
     def setUp(self):
         super().setUp()
-        vocab_size = 10
-        feature_size = 16
+        self.vocab_size = 10
+        self.feature_size = 16
 
         # Create a dummy model to predict the next token.
-        model = tf.keras.Sequential(
+        model = keras.Sequential(
             [
-                tf.keras.Input(shape=[None]),
-                tf.keras.layers.Embedding(
-                    input_dim=vocab_size,
-                    output_dim=feature_size,
+                keras.Input(shape=[None]),
+                keras.layers.Embedding(
+                    input_dim=self.vocab_size,
+                    output_dim=self.feature_size,
                 ),
-                tf.keras.layers.Dense(vocab_size),
-                tf.keras.layers.Softmax(),
+                keras.layers.Dense(self.vocab_size),
+                keras.layers.Softmax(),
             ]
         )
 
@@ -751,7 +781,7 @@ class TopPSearchTextGenerationTest(tf.test.TestCase):
 
         def token_probability_fn(inputs):
             batch_size = inputs.shape[0]
-            prob = tf.keras.activations.softmax(
+            prob = keras.activations.softmax(
                 tf.constant([[1.0, 2.0, 3.0, 4.0]])
             )
             return tf.repeat(prob, batch_size, axis=0)