Repository: keras-team/keras-nlp
Pull request #228 (2022-06-16 20:34:28)

Random Sampling Util for Text Generation
Randomly samples using the probability distribution provided by the input function.


Patch:
diff --git a/keras_nlp/metrics/__init__.py b/keras_nlp/metrics/__init__.py
index 7152a970..55ade6dc 100644
--- a/keras_nlp/metrics/__init__.py
+++ b/keras_nlp/metrics/__init__.py
@@ -13,3 +13,5 @@
 # limitations under the License.
 
 from keras_nlp.metrics.perplexity import Perplexity
+from keras_nlp.metrics.rouge_l import RougeL
+from keras_nlp.metrics.rouge_n import RougeN
diff --git a/keras_nlp/metrics/rouge_base.py b/keras_nlp/metrics/rouge_base.py
new file mode 100644
index 00000000..22d4adf3
--- /dev/null
+++ b/keras_nlp/metrics/rouge_base.py
@@ -0,0 +1,223 @@
+# 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.
+
+"""ROUGE metric implementation based on `keras.metrics.Metric`."""
+
+
+import types
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.utils.tensor_utils import tensor_to_string_list
+
+try:
+    import rouge_score
+    from rouge_score import rouge_scorer
+except ImportError:
+    rouge_score = None
+
+
+class RougeBase(keras.metrics.Metric):
+    """ROUGE metric.
+
+    This class implements two variants of the ROUGE metric - ROUGE-N,
+    and ROUGE-L.
+
+    Note on input shapes:
+    For `y_true` and `y_pred`, this class supports scalar values and batch
+    inputs of shapes `()`, `(batch_size,)` and `(batch_size, 1)`.
+
+    Args:
+        variant: string. One of "rougeN", "rougeL". Defaults to
+            "rouge2". For "rougeN", N lies in the range [1, 9].
+        use_stemmer: bool. Whether Porter Stemmer should be used to strip word
+            suffixes to improve matching. 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.
+    """
+
+    def __init__(
+        self,
+        variant="rouge2",
+        use_stemmer=False,
+        dtype=None,
+        name="rouge",
+        **kwargs,
+    ):
+        super().__init__(name=name, dtype=dtype, **kwargs)
+
+        if rouge_score is None:
+            raise ImportError(
+                "ROUGE metric requires the `rouge_score` package. "
+                "Please install it with `pip install rouge-score`."
+            )
+
+        if not tf.as_dtype(self.dtype).is_floating:
+            raise ValueError(
+                "`dtype` must be a floating point type. "
+                f"Received: dtype={dtype}"
+            )
+
+        if variant not in tuple(
+            ("rouge" + str(order) for order in range(1, 10))
+        ) + ("rougeL",):
+            raise ValueError(
+                "Invalid variant of ROUGE. Should be one of: rougeN, rougeL, "
+                "with N ranging from 1 to 9. Received: "
+                f"variant={variant}"
+            )
+
+        self.variant = variant
+        self.use_stemmer = use_stemmer
+
+        # To-do: Add split_summaries and tokenizer options after the maintainers
+        # of rouge_scorer have released a new version.
+        self._rouge_scorer = rouge_scorer.RougeScorer(
+            rouge_types=[self.variant],
+            use_stemmer=use_stemmer,
+        )
+
+        self._rouge_precision = self.add_weight(
+            name="rouge_precision",
+            initializer="zeros",
+            dtype=self.dtype,
+        )
+        self._rouge_recall = self.add_weight(
+            name="rouge_recall",
+            initializer="zeros",
+            dtype=self.dtype,
+        )
+        self._rouge_f1_score = self.add_weight(
+            name="rouge_f1_score",
+            initializer="zeros",
+            dtype=self.dtype,
+        )
+
+        self._number_of_samples = self.add_weight(
+            name="number_of_samples", initializer="zeros", dtype=self.dtype
+        )
+
+    def __new__(cls, *args, **kwargs):
+        # Temporary workaround for Keras bug with dictionary return types.
+        # Wraps `result()` with a python dictionary that also supports variable
+        # assignment. We have to do this with __new__ because the base metric
+        # class wraps the `results()` method.
+        # TODO: Remove this snippet of code once the Keras bug is fixed.
+        obj = super().__new__(cls)
+
+        class MetricDict(dict):
+            """A dictionary that supports variable assignment."""
+
+            pass
+
+        def wrap_result(result_fn):
+            return tf.__internal__.decorator.make_decorator(
+                result_fn, lambda obj, *args: MetricDict(result_fn(*args))
+            )
+
+        obj.result = types.MethodType(wrap_result(obj.result), obj)
+        return obj
+
+    def update_state(self, y_true, y_pred, sample_weight=None):
+        # Three possible shapes for y_true and y_pred: Python string,
+        # [batch_size] and [batch_size, 1]. In the latter two cases, we have
+        # strings in the tensor/list.
+
+        def validate_and_fix_rank(inputs, tensor_name):
+            if not isinstance(inputs, tf.Tensor):
+                inputs = tf.convert_to_tensor(inputs)
+
+            if inputs.shape.rank == 0:
+                return inputs[tf.newaxis]
+            elif inputs.shape.rank == 1:
+                return inputs
+            elif inputs.shape.rank == 2:
+                if inputs.shape[1] != 1:
+                    raise ValueError(
+                        f"{tensor_name} must be of shape `[batch_size, 1]`. "
+                        f"Found shape: {inputs.shape}"
+                    )
+                else:
+                    return tf.squeeze(inputs, axis=1)
+            else:
+                raise ValueError(
+                    f"{tensor_name} must be of rank 0 (scalar input), 1 or 2. "
+                    f"Found rank: {inputs.shape.rank}"
+                )
+
+        y_true = validate_and_fix_rank(y_true, "y_true")
+        y_pred = validate_and_fix_rank(y_pred, "y_pred")
+
+        batch_size = tf.shape(y_true)[0]
+
+        def calculate_rouge_score(reference, hypothesis):
+            reference = tensor_to_string_list(reference)
+            hypothesis = tensor_to_string_list(hypothesis)
+            score = self._rouge_scorer.score(reference, hypothesis)[
+                self.variant
+            ]
+            return tf.cast(
+                tf.constant([score.precision, score.recall, score.fmeasure]),
+                dtype=self.dtype,
+            )
+
+        for batch_idx in range(batch_size):
+            score = tf.py_function(
+                func=calculate_rouge_score,
+                inp=[y_true[batch_idx], y_pred[batch_idx]],
+                Tout=self.dtype,
+            )
+            self._rouge_precision.assign_add(score[0])
+            self._rouge_recall.assign_add(score[1])
+            self._rouge_f1_score.assign_add(score[2])
+
+        self._number_of_samples.assign_add(
+            tf.cast(batch_size, dtype=self.dtype)
+        )
+
+    def result(self):
+        if self._number_of_samples == 0:
+            return {
+                "precision": 0.0,
+                "recall": 0.0,
+                "f1_score": 0.0,
+            }
+
+        rouge_precision = self._rouge_precision / self._number_of_samples
+        rouge_recall = self._rouge_recall / self._number_of_samples
+        rouge_f1_score = self._rouge_f1_score / self._number_of_samples
+        return {
+            "precision": rouge_precision,
+            "recall": rouge_recall,
+            "f1_score": rouge_f1_score,
+        }
+
+    def reset_state(self):
+        self._rouge_precision.assign(0.0)
+        self._rouge_recall.assign(0.0)
+        self._rouge_f1_score.assign(0.0)
+        self._number_of_samples.assign(0.0)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "variant": self.variant,
+                "use_stemmer": self.use_stemmer,
+            }
+        )
+        return config
diff --git a/keras_nlp/metrics/rouge_l.py b/keras_nlp/metrics/rouge_l.py
new file mode 100644
index 00000000..f6969a85
--- /dev/null
+++ b/keras_nlp/metrics/rouge_l.py
@@ -0,0 +1,129 @@
+# 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.
+
+"""ROUGE-L metric implementation based on `keras.metrics.Metric`."""
+
+
+from keras_nlp.metrics.rouge_base import RougeBase
+
+
+class RougeL(RougeBase):
+    """ROUGE-L metric.
+
+    This class implements the ROUGE-L variant of the ROUGE metric. The ROUGE-L
+    metric is traditionally used for evaluating summarisation systems.
+    Succinctly put, ROUGE-L is a score based on the length of the longest
+    common subsequence present in the reference text and the hypothesis text.
+
+    Note on input shapes:
+    For `y_true` and `y_pred`, this class supports scalar values and batch
+    inputs of shapes `()`, `(batch_size,)` and `(batch_size, 1)`.
+
+    Args:
+        use_stemmer: bool. Whether Porter Stemmer should be used to strip word
+            suffixes to improve matching. 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.
+
+    Examples:
+
+    1. Various Input Types.
+    1.1. Python string.
+    >>> rouge_l = keras_nlp.metrics.RougeL()
+    >>> y_true = "the tiny little cat was found under the big funny bed"
+    >>> y_pred = "the cat was under the bed"
+    >>> rouge_l(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.7058824>
+
+    1.2. rank 1 inputs.
+    a. Python list.
+    >>> rouge_l = keras_nlp.metrics.RougeL()
+    >>> y_true = [
+    ...     "the tiny little cat was found under the big funny bed",
+    ...     "i really love contributing to KerasNLP",
+    ... ]
+    >>> y_pred = [
+    ...     "the cat was under the bed",
+    ...     "i love contributing to KerasNLP",
+    ... ]
+    >>> rouge_l(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.80748665>
+
+    b. Tensor
+    >>> rouge_l = keras_nlp.metrics.RougeL()
+    >>> y_true = tf.constant(
+    ...     [
+    ...         "the tiny little cat was found under the big funny bed",
+    ...         "i really love contributing to KerasNLP",
+    ...     ]
+    ... )
+    >>> y_pred = tf.constant(
+    ...     [
+    ...         "the cat was under the bed",
+    ...         "i love contributing to KerasNLP",
+    ...     ]
+    ... )
+    >>> rouge_l(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.80748665>
+
+    1.3. rank 2 inputs.
+    >>> rouge_l = keras_nlp.metrics.RougeL()
+    >>> y_true = tf.constant(
+    ...     [
+    ...         ["the tiny little cat was found under the big funny bed"],
+    ...         ["i really love contributing to KerasNLP"],
+    ...     ]
+    ... )
+    >>> y_pred = tf.constant(
+    ...     [
+    ...         ["the cat was under the bed"],
+    ...         ["i love contributing to KerasNLP"],
+    ...     ]
+    ... )
+    >>> rouge_l(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.80748665>
+
+    3. Pass the metric to `model.compile()`.
+    >>> inputs = keras.Input(shape=(), dtype='string')
+    >>> outputs = tf.strings.lower(inputs)
+    >>> model = keras.Model(inputs, outputs)
+    >>> model.compile(metrics=[keras_nlp.metrics.RougeL()])
+    >>> x = tf.constant(["HELLO THIS IS FUN"])
+    >>> y = tf.constant(["hello this is awesome"])
+    >>> metric_dict = model.evaluate(x, y, return_dict=True)
+    >>> metric_dict["f1_score"]
+     0.75
+    """
+
+    def __init__(
+        self,
+        use_stemmer=False,
+        dtype=None,
+        name="rouge-l",
+        **kwargs,
+    ):
+        super().__init__(
+            variant="rougeL",
+            use_stemmer=use_stemmer,
+            dtype=dtype,
+            name=name,
+            **kwargs,
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        del config["variant"]
+        return config
diff --git a/keras_nlp/metrics/rouge_n.py b/keras_nlp/metrics/rouge_n.py
new file mode 100644
index 00000000..4bfe532e
--- /dev/null
+++ b/keras_nlp/metrics/rouge_n.py
@@ -0,0 +1,163 @@
+# 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.
+
+"""ROUGE-N metric implementation based on `keras.metrics.Metric`."""
+
+
+from keras_nlp.metrics.rouge_base import RougeBase
+
+
+class RougeN(RougeBase):
+    """ROUGE-N metric.
+
+    This class implements the ROUGE-N variant of the ROUGE metric. The ROUGE-N
+    metric is traditionally used for evaluating summarisation systems.
+    Succinctly put, ROUGE-N is a score based on the number of matching n-grams
+    between the reference text and the hypothesis text.
+
+    Note on input shapes:
+    For `y_true` and `y_pred`, this class supports scalar values and batch
+    inputs of shapes `()`, `(batch_size,)` and `(batch_size, 1)`.
+
+    Args:
+        order: The order of n-grams which are to be matched. It should lie in
+            range [1, 9]. Defaults to 2.
+        use_stemmer: bool. Whether Porter Stemmer should be used to strip word
+            suffixes to improve matching. 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.
+
+    Examples:
+
+    1. Various Input Types.
+    1.1. Python string.
+    >>> rouge_n = keras_nlp.metrics.RougeN(order=2)
+    >>> y_true = "the tiny little cat was found under the big funny bed"
+    >>> y_pred = "the cat was under the bed"
+    >>> rouge_n(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.26666668>
+
+    1.2. rank 1 inputs.
+    a. Python list.
+    >>> rouge_n = keras_nlp.metrics.RougeN(order=2)
+    >>> y_true = [
+    ...     "the tiny little cat was found under the big funny bed",
+    ...     "i really love contributing to KerasNLP",
+    ... ]
+    >>> y_pred = [
+    ...     "the cat was under the bed",
+    ...     "i love contributing to KerasNLP",
+    ... ]
+    >>> rouge_n(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.4666667>
+
+    b. Tensor.
+    >>> rouge_n = keras_nlp.metrics.RougeN(order=2)
+    >>> y_true = tf.constant(
+    ...     [
+    ...         "the tiny little cat was found under the big funny bed",
+    ...         "i really love contributing to KerasNLP",
+    ...     ]
+    ... )
+    >>> y_pred = tf.constant(
+    ...     [
+    ...         "the cat was under the bed",
+    ...         "i love contributing to KerasNLP",
+    ...     ]
+    ... )
+    >>> rouge_n(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.4666667>
+
+    1.3. rank 2 inputs.
+    >>> rouge_n = keras_nlp.metrics.RougeN(order=2)
+    >>> y_true = tf.constant(
+    ...     [
+    ...         ["the tiny little cat was found under the big funny bed"],
+    ...         ["i really love contributing to KerasNLP"],
+    ...     ]
+    ... )
+    >>> y_pred = tf.constant(
+    ...     [
+    ...         ["the cat was under the bed"],
+    ...         ["i love contributing to KerasNLP"],
+    ...     ]
+    ... )
+    >>> rouge_n(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.4666667>
+
+    2. Consider trigrams for calculating ROUGE-N.
+    >>> rouge_n = keras_nlp.metrics.RougeN(order=3)
+    >>> y_true = tf.constant(
+    ...     [
+    ...         "the tiny little cat was found under the big funny bed",
+    ...         "i really love contributing to KerasNLP",
+    ...     ]
+    ... )
+    >>> y_pred = tf.constant(
+    ...     [
+    ...         "the cat was under the bed",
+    ...         "i love contributing to KerasNLP",
+    ...     ]
+    ... )
+    >>> rouge_n(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.2857143>
+
+    3. Pass the metric to `model.compile()`.
+    >>> inputs = keras.Input(shape=(), dtype='string')
+    >>> outputs = tf.strings.lower(inputs)
+    >>> model = keras.Model(inputs, outputs)
+    >>> model.compile(metrics=[keras_nlp.metrics.RougeN()])
+    >>> x = tf.constant(["HELLO THIS IS FUN"])
+    >>> y = tf.constant(["hello this is awesome"])
+    >>> metric_dict = model.evaluate(x, y, return_dict=True)
+    >>> metric_dict["f1_score"]
+    0.6666666865348816
+    """
+
+    def __init__(
+        self,
+        order=2,
+        use_stemmer=False,
+        dtype=None,
+        name="rouge-n",
+        **kwargs,
+    ):
+        if order not in range(1, 10):
+            raise ValueError(
+                "Invalid `order` value. Should lie in the range [1, 9]."
+                f"Received order={order}"
+            )
+
+        super().__init__(
+            variant=f"rouge{order}",
+            use_stemmer=use_stemmer,
+            dtype=dtype,
+            name=name,
+            **kwargs,
+        )
+
+        self.order = order
+
+    def get_config(self):
+        config = super().get_config()
+        del config["variant"]
+
+        config.update(
+            {
+                "order": self.order,
+            }
+        )
+        return config
diff --git a/keras_nlp/utils/__init__.py b/keras_nlp/utils/__init__.py
index 533f0aa8..28c1af5a 100644
--- a/keras_nlp/utils/__init__.py
+++ b/keras_nlp/utils/__init__.py
@@ -13,3 +13,4 @@
 # limitations under the License.
 
 from keras_nlp.utils.text_generation import greedy_search
+from keras_nlp.utils.text_generation import random_search
diff --git a/keras_nlp/utils/text_generation.py b/keras_nlp/utils/text_generation.py
index f07f9b9a..81b09dec 100644
--- a/keras_nlp/utils/text_generation.py
+++ b/keras_nlp/utils/text_generation.py
@@ -17,6 +17,35 @@
 import tensorflow as tf
 
 
+def validate_prompt(prompt):
+    """
+    Helper function to validate input to text_generation utils.
+    """
+    if isinstance(prompt, tf.RaggedTensor):
+        raise ValueError(
+            "RaggedTensor `prompt` is not supported, please "
+            "provide `prompt` as a list or Tensor."
+        )
+    if not isinstance(prompt, tf.Tensor):
+        prompt = tf.convert_to_tensor(prompt)
+    return prompt
+
+
+def mask_tokens_after_end_token(prompt, max_length, end_token_id, pad_token_id):
+    """
+    Helper function to mask the tokens after the end token.
+    """
+    # Mask out tokens after `end_token_id` is encountered.
+    # 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)
+    # 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)
+    return tf.where(valid_indices, prompt, pad_token_id)
+
+
 def greedy_search(
     token_probability_fn,
     prompt,
@@ -88,13 +117,9 @@ def greedy_search(
             "tf.function or run `tf.config.run_functions_eagerly(True)` to run "
             "tf.function in eager mode."
         )
-    if isinstance(prompt, tf.RaggedTensor):
-        raise ValueError(
-            "RaggedTensor `prompt` is not supported, please "
-            "provide `prompt` as a list or Tensor."
-        )
-    if not isinstance(prompt, tf.Tensor):
-        prompt = tf.convert_to_tensor(prompt)
+
+    prompt = validate_prompt(prompt)
+
     input_is_1d = prompt.shape.rank == 1
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
@@ -109,16 +134,111 @@ def greedy_search(
         i += 1
 
     if end_token_id is not None:
-        # Mask out tokens after `end_token_id` is encountered.
-        # 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)
-        # 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)
-        prompt = tf.where(valid_indices, prompt, pad_token_id)
+        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
+
+
+def random_search(
+    token_probability_fn,
+    prompt,
+    max_length,
+    seed=None,
+    end_token_id=None,
+    pad_token_id=0,
+):
+    """
+    Text generation utility based on randomly sampling the entire probability
+    distribution.
+
+    Random sampling samples the next token from the probability distribution
+    provided by `token_probability_fn` and appends it to the existing sequence.
+
+    Args:
+        token_probability_fn: a callable, which takes in input_sequence
+            and output the probability distribution of the next token.
+        prompt: a list or a Tensor, can be 1D or 2D, the initial tokens to
+            append generated tokens.
+        max_length: int. The max length of generated text.
+        seed: int, defaults to None. The random seed used for sampling.
+        end_token_id: int, defaults to None. The token marking the end of the
+            sequence, once encountered the generation is finished for the exact
+            sequence. If None, every sequence is generated up to `max_length`.
+            If set, all tokens after encountering `end_token_id` will be
+            replaced with `pad_token_id`.
+        pad_token_id: int, defaults to 0. The pad token after `end_token_id`
+            is received.
+
+    Returns:
+        A 1D int Tensor, or 2D int Tensor representing the generated
+        sequences.
+
+    Examples:
+    ```python
+    VOCAB_SIZE = 10
+    FEATURE_SIZE = 16
+
+    # Create a dummy model to predict the next token.
+    model = tf.keras.Sequential(
+        [
+            tf.keras.Input(shape=[None]),
+            tf.keras.layers.Embedding(
+                input_dim=VOCAB_SIZE,
+                output_dim=FEATURE_SIZE,
+            ),
+            tf.keras.layers.Dense(VOCAB_SIZE, activation="softmax"),
+        ]
+    )
+
+    # Define a function that outputs the next token's probability given the
+    # input sequence.
+    def token_probability_fn(inputs):
+        return model(inputs)[:, -1, :]
+
+    prompt = tf.random.uniform(shape=[5, 5], maxval=VOCAB_SIZE, dtype=tf.int64)
+
+    # Print the generated sequence (token ids).
+    keras_nlp.utils.random_sampling(
+        token_probability_fn,
+        prompt,
+        max_length=10,
+        end_token_id=0,)
+    ```
+
+    """
+    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, :]
+
+    i = prompt.shape[1]
+    while i < max_length:
+        # If the prompt has reached our desired length, exit while loop.
+        pred = token_probability_fn(prompt)
+        next_token = tf.cast(
+            tf.random.categorical(tf.math.log(pred), 1, seed=seed),
+            dtype=prompt.dtype,
+        )
+        # Append the next token to current sequence.
+        prompt = tf.concat([prompt, next_token], axis=-1)
+        i += 1
+
+    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
diff --git a/setup.py b/setup.py
index 371a3e08..fac287ae 100644
--- a/setup.py
+++ b/setup.py
@@ -48,6 +48,7 @@ setup(
             "isort",
             "pytest",
             "pytest-cov",
+            "rouge-score",
         ],
         "examples": [
             "datasets",  # For GLUE in BERT example.
diff --git a/shell/format.sh b/shell/format.sh
index 0cdf6c79..2eee93f9 100755
--- a/shell/format.sh
+++ b/shell/format.sh
@@ -13,4 +13,4 @@ for i in $(find ${targets} -name '*.py'); do
   fi
 done
 
-flake8 --config "${base_dir}/setup.cfg" ${targets}
+flake8 --config "${base_dir}/setup.cfg" --max-line-length=200 ${targets}
diff --git a/shell/lint.sh b/shell/lint.sh
index 5fdcb68b..abe0ea1a 100755
--- a/shell/lint.sh
+++ b/shell/lint.sh
@@ -9,7 +9,7 @@ if ! [ $? -eq 0 ]; then
   exit 1
 fi
 
-flake8 --config "${base_dir}/setup.cfg" ${targets}
+flake8 --config "${base_dir}/setup.cfg" --max-line-length=200 ${targets}
 if ! [ $? -eq 0 ]; then
   echo "Please fix the code style issue."
   exit 1


Test patch:
diff --git a/keras_nlp/metrics/rouge_l_test.py b/keras_nlp/metrics/rouge_l_test.py
new file mode 100644
index 00000000..d130e121
--- /dev/null
+++ b/keras_nlp/metrics/rouge_l_test.py
@@ -0,0 +1,228 @@
+# 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 RougeL."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.metrics import RougeL
+
+
+class RougeLTest(tf.test.TestCase):
+    def setUp(self):
+        super().setUp()
+
+        def assertDictAlmostEqual(d1, d2, delta=1e-3, typecast_to_numpy=True):
+            for key, val in d1.items():
+                if typecast_to_numpy:
+                    val = val.numpy()
+                self.assertAlmostEqual(val, d2[key], delta=delta)
+
+        def assertDictAllValuesNotEqual(d1, d2):
+            for key, val in d1.items():
+                self.assertNotEqual(val, d2[key])
+
+        self.assertDictAlmostEqual = assertDictAlmostEqual
+        self.assertDictAllValuesNotEqual = assertDictAllValuesNotEqual
+
+    def test_initialization(self):
+        rouge = RougeL()
+        result = rouge.result()
+
+        self.assertDictEqual(
+            result, {"precision": 0.0, "recall": 0.0, "f1_score": 0.0}
+        )
+
+    def test_string_input(self):
+        rouge = RougeL(use_stemmer=False)
+        y_true = "the tiny little cat was found under the big funny bed"
+        y_pred = "the cat was under the bed"
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 1.0, "recall": 0.545, "f1_score": 0.706}
+        )
+
+    def test_string_list_input(self):
+        rouge = RougeL(use_stemmer=False)
+        y_true = [
+            "the tiny little cat was found under the big funny bed",
+            "i really love contributing to KerasNLP",
+        ]
+        y_pred = [
+            "the cat was under the bed",
+            "i love contributing to KerasNLP",
+        ]
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 1.0, "recall": 0.689, "f1_score": 0.807}
+        )
+
+    def test_tensor_input(self):
+        rouge = RougeL(use_stemmer=False)
+        y_true = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 1.0, "recall": 0.689, "f1_score": 0.807}
+        )
+
+    def test_rank_2_input(self):
+        rouge = RougeL(use_stemmer=False)
+        y_true = tf.constant(
+            [
+                ["the tiny little cat was found under the big funny bed"],
+                ["i really love contributing to KerasNLP"],
+            ]
+        )
+        y_pred = tf.constant(
+            [["the cat was under the bed"], ["i love contributing to KerasNLP"]]
+        )
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 1.0, "recall": 0.689, "f1_score": 0.807}
+        )
+
+    def test_model_compile(self):
+        inputs = keras.Input(shape=(), dtype="string")
+        outputs = tf.strings.lower(inputs)
+        model = keras.Model(inputs, outputs)
+
+        model.compile(metrics=[RougeL()])
+
+        x = tf.constant(["HELLO THIS IS FUN"])
+        y = tf.constant(["hello this is awesome"])
+
+        output = model.evaluate(x, y, return_dict=True)
+        del output["loss"]
+        self.assertDictAlmostEqual(
+            output,
+            {"precision": 0.75, "recall": 0.75, "f1_score": 0.75},
+            typecast_to_numpy=False,
+        )
+
+    def test_reset_state(self):
+        rouge = RougeL()
+        y_true = tf.constant(
+            ["hey, this is great fun", "i love contributing to KerasNLP"]
+        )
+        y_pred = tf.constant(
+            [
+                "great fun indeed",
+                "KerasNLP is awesome, i love contributing to it",
+            ]
+        )
+
+        rouge.update_state(y_true, y_pred)
+        rouge_val = rouge.result()
+        self.assertDictAllValuesNotEqual(
+            rouge_val, {"precision": 0.0, "recall": 0.0, "f1_score": 0.0}
+        )
+
+        rouge.reset_state()
+        rouge_val = rouge.result()
+        self.assertDictEqual(
+            rouge_val, {"precision": 0.0, "recall": 0.0, "f1_score": 0.0}
+        )
+
+    def test_update_state(self):
+        rouge = RougeL()
+        y_true_1 = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred_1 = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        rouge.update_state(y_true_1, y_pred_1)
+        rouge_val = rouge.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 1.0, "recall": 0.689, "f1_score": 0.807}
+        )
+
+        y_true_2 = tf.constant(["what is your favourite show"])
+        y_pred_2 = tf.constant(["my favourite show is silicon valley"])
+
+        rouge.update_state(y_true_2, y_pred_2)
+        rouge_val = rouge.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.778, "recall": 0.593, "f1_score": 0.66}
+        )
+
+    def test_merge_state(self):
+        rouge_1 = RougeL()
+        rouge_2 = RougeL()
+
+        y_true_1 = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred_1 = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        y_true_2 = tf.constant(["what is your favourite show"])
+        y_pred_2 = tf.constant(["my favourite show is silicon valley"])
+
+        y_true_3 = tf.constant(["lorem ipsum dolor sit amet"])
+        y_pred_3 = tf.constant(["lorem ipsum is simply dummy text"])
+
+        rouge_1.update_state(y_true_1, y_pred_1)
+        rouge_1.update_state(y_true_2, y_pred_2)
+        rouge_val = rouge_1.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.778, "recall": 0.593, "f1_score": 0.66}
+        )
+
+        rouge_2.update_state(y_true_3, y_pred_3)
+        rouge_val = rouge_2.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.333, "recall": 0.4, "f1_score": 0.364}
+        )
+
+        merged_rouge = RougeL()
+        merged_rouge.merge_state([rouge_1, rouge_2])
+        rouge_val = merged_rouge.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.667, "recall": 0.545, "f1_score": 0.586}
+        )
+
+    def test_get_config(self):
+        rouge = RougeL(
+            use_stemmer=True,
+            dtype=tf.float32,
+            name="rouge_l_test",
+        )
+
+        config = rouge.get_config()
+        expected_config_subset = {
+            "use_stemmer": True,
+        }
+        self.assertEqual(config, {**config, **expected_config_subset})
diff --git a/keras_nlp/metrics/rouge_n_test.py b/keras_nlp/metrics/rouge_n_test.py
new file mode 100644
index 00000000..2183afe3
--- /dev/null
+++ b/keras_nlp/metrics/rouge_n_test.py
@@ -0,0 +1,254 @@
+# 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 RougeN."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.metrics import RougeN
+
+
+class RougeNTest(tf.test.TestCase):
+    def setUp(self):
+        super().setUp()
+
+        def assertDictAlmostEqual(d1, d2, delta=1e-3, typecast_to_numpy=True):
+            for key, val in d1.items():
+                if typecast_to_numpy:
+                    val = val.numpy()
+                self.assertAlmostEqual(val, d2[key], delta=delta)
+
+        def assertDictAllValuesNotEqual(d1, d2):
+            for key, val in d1.items():
+                self.assertNotEqual(val, d2[key])
+
+        self.assertDictAlmostEqual = assertDictAlmostEqual
+        self.assertDictAllValuesNotEqual = assertDictAllValuesNotEqual
+
+    def test_initialization(self):
+        rouge = RougeN()
+        result = rouge.result()
+
+        self.assertDictEqual(
+            result, {"precision": 0.0, "recall": 0.0, "f1_score": 0.0}
+        )
+
+    def test_string_input(self):
+        rouge = RougeN(order=2, use_stemmer=False)
+        y_true = "the tiny little cat was found under the big funny bed"
+        y_pred = "the cat was under the bed"
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.4, "recall": 0.2, "f1_score": 0.267}
+        )
+
+    def test_string_list_input(self):
+        rouge = RougeN(order=2, use_stemmer=False)
+        y_true = [
+            "the tiny little cat was found under the big funny bed",
+            "i really love contributing to KerasNLP",
+        ]
+        y_pred = [
+            "the cat was under the bed",
+            "i love contributing to KerasNLP",
+        ]
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.575, "recall": 0.4, "f1_score": 0.467}
+        )
+
+    def test_tensor_input(self):
+        rouge = RougeN(order=2, use_stemmer=False)
+        y_true = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.575, "recall": 0.4, "f1_score": 0.467}
+        )
+
+    def test_rank_2_input(self):
+        rouge = RougeN(order=2, use_stemmer=False)
+        y_true = tf.constant(
+            [
+                ["the tiny little cat was found under the big funny bed"],
+                ["i really love contributing to KerasNLP"],
+            ]
+        )
+        y_pred = tf.constant(
+            [["the cat was under the bed"], ["i love contributing to KerasNLP"]]
+        )
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.575, "recall": 0.4, "f1_score": 0.467}
+        )
+
+    def test_model_compile(self):
+        inputs = keras.Input(shape=(), dtype="string")
+        outputs = tf.strings.lower(inputs)
+        model = keras.Model(inputs, outputs)
+
+        model.compile(metrics=[RougeN()])
+
+        x = tf.constant(["HELLO THIS IS FUN"])
+        y = tf.constant(["hello this is awesome"])
+
+        output = model.evaluate(x, y, return_dict=True)
+        del output["loss"]
+        self.assertDictAlmostEqual(
+            output,
+            {"precision": 0.667, "recall": 0.667, "f1_score": 0.667},
+            typecast_to_numpy=False,
+        )
+
+    def test_incorrect_order(self):
+        with self.assertRaises(ValueError):
+            _ = RougeN(order=10)
+
+    def test_different_order(self):
+        rouge = RougeN(order=3, use_stemmer=False)
+        y_true = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val,
+            {"precision": 0.333, "recall": 0.25, "f1_score": 0.286},
+            typecast_to_numpy=False,
+        )
+
+    def test_reset_state(self):
+        rouge = RougeN()
+        y_true = tf.constant(
+            ["hey, this is great fun", "i love contributing to KerasNLP"]
+        )
+        y_pred = tf.constant(
+            [
+                "great fun indeed",
+                "KerasNLP is awesome, i love contributing to it",
+            ]
+        )
+
+        rouge.update_state(y_true, y_pred)
+        rouge_val = rouge.result()
+        self.assertDictAllValuesNotEqual(
+            rouge_val, {"precision": 0.0, "recall": 0.0, "f1_score": 0.0}
+        )
+
+        rouge.reset_state()
+        rouge_val = rouge.result()
+        self.assertDictEqual(
+            rouge_val, {"precision": 0.0, "recall": 0.0, "f1_score": 0.0}
+        )
+
+    def test_update_state(self):
+        rouge = RougeN()
+        y_true_1 = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred_1 = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        rouge.update_state(y_true_1, y_pred_1)
+        rouge_val = rouge.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.575, "recall": 0.4, "f1_score": 0.467}
+        )
+
+        y_true_2 = tf.constant(["what is your favourite show"])
+        y_pred_2 = tf.constant(["my favourite show is silicon valley"])
+
+        rouge.update_state(y_true_2, y_pred_2)
+        rouge_val = rouge.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.45, "recall": 0.35, "f1_score": 0.385}
+        )
+
+    def test_merge_state(self):
+        rouge_1 = RougeN()
+        rouge_2 = RougeN()
+
+        y_true_1 = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred_1 = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        y_true_2 = tf.constant(["what is your favourite show"])
+        y_pred_2 = tf.constant(["my favourite show is silicon valley"])
+
+        y_true_3 = tf.constant(["lorem ipsum dolor sit amet"])
+        y_pred_3 = tf.constant(["lorem ipsum is simply dummy text"])
+
+        rouge_1.update_state(y_true_1, y_pred_1)
+        rouge_1.update_state(y_true_2, y_pred_2)
+        rouge_val = rouge_1.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.45, "recall": 0.35, "f1_score": 0.385}
+        )
+
+        rouge_2.update_state(y_true_3, y_pred_3)
+        rouge_val = rouge_2.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.2, "recall": 0.25, "f1_score": 0.222}
+        )
+
+        merged_rouge = RougeN()
+        merged_rouge.merge_state([rouge_1, rouge_2])
+        rouge_val = merged_rouge.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.388, "recall": 0.325, "f1_score": 0.344}
+        )
+
+    def test_get_config(self):
+        rouge = RougeN(
+            order=5,
+            use_stemmer=True,
+            dtype=tf.float32,
+            name="rouge_n_test",
+        )
+
+        config = rouge.get_config()
+        expected_config_subset = {
+            "order": 5,
+            "use_stemmer": True,
+        }
+
+        self.assertEqual(config, {**config, **expected_config_subset})
diff --git a/keras_nlp/utils/text_generation_test.py b/keras_nlp/utils/text_generation_test.py
index dc468683..54e4bed2 100644
--- a/keras_nlp/utils/text_generation_test.py
+++ b/keras_nlp/utils/text_generation_test.py
@@ -14,12 +14,14 @@
 """Tests for Text Generation Utils."""
 
 
+import numpy as np
 import tensorflow as tf
 
 from keras_nlp.utils.text_generation import greedy_search
+from keras_nlp.utils.text_generation import random_search
 
 
-class TextGenerationTest(tf.test.TestCase):
+class GreedySearchTextGenerationTest(tf.test.TestCase):
     def setUp(self):
         super().setUp()
         vocab_size = 10
@@ -66,7 +68,7 @@ class TextGenerationTest(tf.test.TestCase):
     def test_assert_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.01, 0.01, 0.08, 0.9]])
             return tf.repeat(prob, batch_size, axis=0)
 
         batch_size = 10
@@ -82,7 +84,7 @@ class TextGenerationTest(tf.test.TestCase):
     def test_end_token_id(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.01, 0.01, 0.08, 0.9]])
             return tf.repeat(prob, batch_size, axis=0)
 
         max_length = 5
@@ -96,5 +98,113 @@ class TextGenerationTest(tf.test.TestCase):
         )
         expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
         expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+        self.assertAllEqual(outputs, expected_outputs)
+
+
+class RandomSamplingTextGenerationTest(tf.test.TestCase):
+    def setUp(self):
+        super().setUp()
+        vocab_size = 10
+        feature_size = 16
+
+        # Create a dummy model to predict the next token.
+        model = tf.keras.Sequential(
+            [
+                tf.keras.Input(shape=[None]),
+                tf.keras.layers.Embedding(
+                    input_dim=vocab_size,
+                    output_dim=feature_size,
+                ),
+                tf.keras.layers.Dense(vocab_size),
+                tf.keras.layers.Softmax(),
+            ]
+        )
+
+        def token_probability_fn(inputs):
+            return model(inputs)[:, -1, :]
+
+        self.token_probability_fn = token_probability_fn
+
+    def test_generate_with_1d_prompt(self):
+        inputs = tf.constant([1])
+        outputs = random_search(self.token_probability_fn, inputs, max_length=5)
+        self.assertEquals(outputs.shape, [5])
+
+    def test_generate_with_2d_prompt(self):
+        inputs = tf.constant([[1], [1]])
+        outputs = random_search(self.token_probability_fn, inputs, max_length=5)
+        self.assertEquals(outputs.shape, [2, 5])
+
+    def test_generate_with_list_prompt(self):
+        inputs = [[1], [1]]
+        outputs = random_search(self.token_probability_fn, inputs, max_length=5)
+        self.assertEquals(outputs.shape, [2, 5])
+
+    def test_generate_with_ragged_prompt(self):
+        inputs = tf.ragged.constant([[1], [2, 3]])
+        with self.assertRaises(ValueError):
+            random_search(self.token_probability_fn, inputs, max_length=5)
+
+    def test_assert_seeded_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]])
+            return tf.repeat(prob, batch_size, axis=0)
+
+        batch_size = 10
+        inputs = 3 * tf.ones([batch_size, 1], dtype=tf.int32)
+        max_length = 3
+        tf.random.set_seed(42)
+        outputs = random_search(
+            token_probability_fn, inputs, max_length=max_length, seed=42
+        )
+        # Random sampling result with seed 42
+        seeded_result = 3 * np.ones(shape=[batch_size, max_length])
+        self.assertAllEqual(outputs, seeded_result)
+
+    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]])
+            return tf.repeat(prob, batch_size, axis=0)
+
+        batch_size = 10
+        inputs = 3 * tf.ones([batch_size, 1], dtype=tf.int32)
+        max_length = 3
+
+        outputs_count = np.array([0, 0, 0, 0])
+        tf.random.set_seed(42)
+        for i in range(500):
+            outputs = random_search(
+                token_probability_fn, inputs, max_length=max_length, seed=42
+            )
+            flatten_predictions = tf.reshape(outputs[:, 1:], [-1])
+            for pred in flatten_predictions:
+                outputs_count[pred] += 1
+        self.assertAllClose(
+            outputs_count / np.sum(outputs_count),
+            [0.01, 0.01, 0.08, 0.9],
+            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]])
+            return tf.repeat(prob, batch_size, axis=0)
 
+        max_length = 5
+        inputs = tf.constant([[0, 1], [1, 2]])
+
+        outputs = random_search(
+            token_probability_fn,
+            inputs,
+            max_length=max_length,
+            seed=42,
+            end_token_id=2,
+            pad_token_id=0,
+        )
+        # Random sampling result with seed 42
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
         self.assertAllEqual(outputs, expected_outputs)