Repository: keras-team/keras-nlp
Pull request #232 (2022-06-22 00:41:38)

added top k search util
Follow up from random search


Patch:
diff --git a/keras_nlp/utils/__init__.py b/keras_nlp/utils/__init__.py
index 28c1af5a..9bbfcbde 100644
--- a/keras_nlp/utils/__init__.py
+++ b/keras_nlp/utils/__init__.py
@@ -14,3 +14,4 @@
 
 from keras_nlp.utils.text_generation import greedy_search
 from keras_nlp.utils.text_generation import random_search
+from keras_nlp.utils.text_generation import top_k_search
diff --git a/keras_nlp/utils/text_generation.py b/keras_nlp/utils/text_generation.py
index 81b09dec..4c758363 100644
--- a/keras_nlp/utils/text_generation.py
+++ b/keras_nlp/utils/text_generation.py
@@ -18,9 +18,7 @@ import tensorflow as tf
 
 
 def validate_prompt(prompt):
-    """
-    Helper function to validate input to text_generation utils.
-    """
+    """Helper function to validate input to text_generation utils."""
     if isinstance(prompt, tf.RaggedTensor):
         raise ValueError(
             "RaggedTensor `prompt` is not supported, please "
@@ -31,10 +29,19 @@ def validate_prompt(prompt):
     return prompt
 
 
+def validate_token_probability_fn(token_probability_fn, prompt):
+    """Helper function to validate token probability fn output"""
+    test_pred = token_probability_fn(prompt)
+    if len(test_pred.shape) != 2:
+        raise ValueError(
+            "Output of `token_probability_fn` is not a 2D tensor, "
+            "please provide a function with the output shape "
+            "[batch_size, vocab_size]."
+        )
+
+
 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.
-    """
+    """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)
@@ -61,7 +68,8 @@ def greedy_search(
 
     Args:
         token_probability_fn: a callable, which takes in input_sequence
-            and output the probability distribution of the next token.
+            and output the probability distribution or the logits 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.
@@ -79,8 +87,11 @@ def greedy_search(
 
     Examples:
     ```python
+    BATCH_SIZE = 8
     VOCAB_SIZE = 10
     FEATURE_SIZE = 16
+    START_ID = 1
+    END_ID = 2
 
     # Create a dummy model to predict the next token.
     model = tf.keras.Sequential(
@@ -99,14 +110,15 @@ def greedy_search(
     def token_probability_fn(inputs):
         return model(inputs)[:, -1, :]
 
-    prompt = tf.random.uniform(shape=[5, 5], maxval=VOCAB_SIZE, dtype=tf.int64)
+    prompt = tf.fill((BATCH_SIZE, 1), START_ID)
 
     # Print the generated sequence (token ids).
     keras_nlp.utils.greedy_search(
         token_probability_fn,
         prompt,
         max_length=10,
-        end_token_id=0,)
+        end_token_id=END_ID,
+    )
     ```
 
     """
@@ -123,6 +135,7 @@ def greedy_search(
     input_is_1d = prompt.shape.rank == 1
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
+    validate_token_probability_fn(token_probability_fn, prompt)
 
     i = prompt.shape[1]
     while i < max_length:
@@ -148,6 +161,7 @@ def random_search(
     prompt,
     max_length,
     seed=None,
+    from_logits=False,
     end_token_id=None,
     pad_token_id=0,
 ):
@@ -160,11 +174,15 @@ def random_search(
 
     Args:
         token_probability_fn: a callable, which takes in input_sequence
-            and output the probability distribution of the next token.
+            and output the probability distribution of the next token. If
+            `from_logits` set to True, it should output the logits 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.
+        from_logits: bool. Indicates whether `token_probability_fn` outputs
+            logits or probabilities.
         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`.
@@ -179,8 +197,11 @@ def random_search(
 
     Examples:
     ```python
+    BATCH_SIZE = 8
     VOCAB_SIZE = 10
     FEATURE_SIZE = 16
+    START_ID = 1
+    END_ID = 2
 
     # Create a dummy model to predict the next token.
     model = tf.keras.Sequential(
@@ -199,14 +220,15 @@ def random_search(
     def token_probability_fn(inputs):
         return model(inputs)[:, -1, :]
 
-    prompt = tf.random.uniform(shape=[5, 5], maxval=VOCAB_SIZE, dtype=tf.int64)
+    prompt = tf.fill((BATCH_SIZE, 1), START_ID)
 
     # Print the generated sequence (token ids).
-    keras_nlp.utils.random_sampling(
+    keras_nlp.utils.random_search(
         token_probability_fn,
         prompt,
         max_length=10,
-        end_token_id=0,)
+        end_token_id=END_ID,
+    )
     ```
 
     """
@@ -222,11 +244,14 @@ def random_search(
     input_is_1d = prompt.shape.rank == 1
     if input_is_1d:
         prompt = prompt[tf.newaxis, :]
+    validate_token_probability_fn(token_probability_fn, prompt)
 
     i = prompt.shape[1]
     while i < max_length:
         # If the prompt has reached our desired length, exit while loop.
         pred = token_probability_fn(prompt)
+        if from_logits:
+            pred = tf.keras.activations.softmax(pred, axis=-1)
         next_token = tf.cast(
             tf.random.categorical(tf.math.log(pred), 1, seed=seed),
             dtype=prompt.dtype,
@@ -242,3 +267,129 @@ def random_search(
     if input_is_1d:
         return tf.squeeze(prompt)
     return prompt
+
+
+def top_k_search(
+    token_probability_fn,
+    prompt,
+    max_length,
+    k,
+    seed=None,
+    from_logits=False,
+    end_token_id=None,
+    pad_token_id=0,
+):
+    """
+    Text generation utility based on top-k sampling.
+
+    Top-k search samples the next token from the top-k tokens in 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. If
+            `from_logits` set to True, it should output the logits 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.
+        k: int. The number of top tokens to sample from. Should be non-negative
+            and less than the vocabulary size.
+        seed: int, defaults to None. The random seed used for sampling.
+        from_logits: bool. Indicates whether `token_probability_fn` outputs
+            logits or probabilities.
+        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
+    BATCH_SIZE = 8
+    VOCAB_SIZE = 10
+    FEATURE_SIZE = 16
+    START_ID = 1
+    END_ID = 2
+
+    # 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.fill((BATCH_SIZE, 1), START_ID)
+
+    # Print the generated sequence (token ids).
+    keras_nlp.utils.top_k_search(
+        token_probability_fn,
+        prompt,
+        max_length=10,
+        k=4,
+        end_token_id=END_ID,
+    )
+    ```
+
+    """
+    if not tf.executing_eagerly():
+        raise RuntimeError(
+            "`keras_nlp.utils.top_k_search` currently requires an eager "
+            "execution context. Please call `top_k_search` outside "
+            "tf.function or run `tf.config.run_functions_eagerly(True)` to run "
+            "tf.function in eager mode."
+        )
+    if k <= 0:
+        raise ValueError("k should be strictly positive (greater than 0).")
+
+    prompt = validate_prompt(prompt)
+    input_is_1d = prompt.shape.rank == 1
+    if input_is_1d:
+        prompt = prompt[tf.newaxis, :]
+    validate_token_probability_fn(token_probability_fn, prompt)
+
+    i = prompt.shape[1]
+    while i < max_length:
+        # If the prompt has reached our desired length, exit while loop.
+        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])
+        # Filter out top-k tokens.
+        top_k_pred, top_k_indices = tf.math.top_k(pred, k=k)
+        # Sample the next token from the probability distribution.
+        next_token = tf.random.categorical(
+            tf.math.log(top_k_pred), 1, seed=seed
+        )
+        # Rearrange to get the next token idx from the original order.
+        next_token = tf.gather_nd(top_k_indices, next_token, batch_dims=1)
+        next_token = tf.cast(next_token, dtype=prompt.dtype)
+        # Append the next token to current sequence.
+        prompt = tf.concat([prompt, next_token[:, tf.newaxis]], axis=-1)
+        i += 1
+
+    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


Test patch:
diff --git a/keras_nlp/utils/text_generation_test.py b/keras_nlp/utils/text_generation_test.py
index 54e4bed2..fab1e4e0 100644
--- a/keras_nlp/utils/text_generation_test.py
+++ b/keras_nlp/utils/text_generation_test.py
@@ -19,6 +19,7 @@ import tensorflow as tf
 
 from keras_nlp.utils.text_generation import greedy_search
 from keras_nlp.utils.text_generation import random_search
+from keras_nlp.utils.text_generation import top_k_search
 
 
 class GreedySearchTextGenerationTest(tf.test.TestCase):
@@ -158,7 +159,7 @@ class RandomSamplingTextGenerationTest(tf.test.TestCase):
         outputs = random_search(
             token_probability_fn, inputs, max_length=max_length, seed=42
         )
-        # Random sampling result with seed 42
+        # Random sampling result with seed 42.
         seeded_result = 3 * np.ones(shape=[batch_size, max_length])
         self.assertAllEqual(outputs, seeded_result)
 
@@ -195,7 +196,7 @@ class RandomSamplingTextGenerationTest(tf.test.TestCase):
 
         max_length = 5
         inputs = tf.constant([[0, 1], [1, 2]])
-
+        tf.random.set_seed(42)
         outputs = random_search(
             token_probability_fn,
             inputs,
@@ -204,7 +205,216 @@ class RandomSamplingTextGenerationTest(tf.test.TestCase):
             end_token_id=2,
             pad_token_id=0,
         )
-        # Random sampling result with seed 42
+        # 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)
+
+    def test_from_logits(self):
+        def token_logits_fn(inputs):
+            batch_size = inputs.shape[0]
+            prob = tf.constant([[1.0, 2.0, 3, 0, 4.0]])
+            return tf.repeat(prob, batch_size, axis=0)
+
+        def token_probability_fn(inputs):
+            batch_size = inputs.shape[0]
+            prob = tf.keras.activations.softmax(
+                tf.constant([[1.0, 2.0, 3, 0, 4.0]])
+            )
+            return tf.repeat(prob, batch_size, axis=0)
+
+        max_length = 5
+        inputs = tf.constant([[0, 1], [1, 2]])
+        tf.random.set_seed(42)
+        output_logit = random_search(
+            token_logits_fn,
+            inputs,
+            max_length=max_length,
+            from_logits=True,
+            seed=42,
+        )
+        tf.random.set_seed(42)
+        output_probs = random_search(
+            token_probability_fn,
+            inputs,
+            max_length=max_length,
+            from_logits=False,
+            seed=42,
+        )
+        self.assertAllEqual(output_logit, output_probs)
+
+
+class TopKSamplingTextGenerationTest(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 = top_k_search(
+            self.token_probability_fn, inputs, max_length=5, k=2
+        )
+        self.assertEquals(outputs.shape, [5])
+
+    def test_generate_with_2d_prompt(self):
+        inputs = tf.constant([[1], [1]])
+        outputs = top_k_search(
+            self.token_probability_fn, inputs, max_length=5, k=2
+        )
+        self.assertEquals(outputs.shape, [2, 5])
+
+    def test_generate_with_list_prompt(self):
+        inputs = [[1], [1]]
+        outputs = top_k_search(
+            self.token_probability_fn, inputs, max_length=5, k=2
+        )
+        self.assertEquals(outputs.shape, [2, 5])
+
+    def test_generate_with_ragged_prompt(self):
+        inputs = tf.ragged.constant([[1], [2, 3]])
+        with self.assertRaises(ValueError):
+            top_k_search(self.token_probability_fn, inputs, max_length=5, k=2)
+
+    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 = top_k_search(
+            token_probability_fn, inputs, max_length=max_length, k=2, seed=42
+        )
+        # Top-k sampling result with seed 42.
+        seeded_result = 3 * np.ones(shape=[batch_size, max_length])
+        seeded_result[3][1] = 2
+        seeded_result[7][1] = 2
+        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.1, 0.2, 0.3, 0.4]])
+            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 = top_k_search(
+                token_probability_fn,
+                inputs,
+                max_length=max_length,
+                k=2,
+                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.0, 0.0, 0.429, 0.571],
+            rtol=0.2,
+        )
+
+    def test_only_choose_from_top_k_tokens(self):
+        # Test that there are only the top-k tokens in the output.
+        def token_probability_fn(inputs):
+            batch_size = inputs.shape[0]
+            prob = tf.constant([[0.4, 0.3, 0.2, 0.1]])
+            return tf.repeat(prob, batch_size, axis=0)
+
+        # Test that it only samples from top-k tokens.
+        for k in [1, 2, 3]:
+            inputs = tf.constant([[0, 0], [0, 0]])
+            for _ in range(10):
+                outputs = top_k_search(
+                    token_probability_fn,
+                    inputs,
+                    max_length=5,
+                    k=k,
+                )
+                self.assertAllEqual(outputs < k, tf.ones_like(outputs))
+
+    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]])
+        tf.random.set_seed(42)
+        outputs = top_k_search(
+            token_probability_fn,
+            inputs,
+            max_length=max_length,
+            k=4,
+            seed=42,
+            end_token_id=2,
+            pad_token_id=0,
+        )
+        # Top-k 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)
+
+    def test_from_logits(self):
+        def token_logits_fn(inputs):
+            batch_size = inputs.shape[0]
+            prob = tf.constant([[1.0, 2.0, 3, 0, 4.0]])
+            return tf.repeat(prob, batch_size, axis=0)
+
+        def token_probability_fn(inputs):
+            batch_size = inputs.shape[0]
+            prob = tf.keras.activations.softmax(
+                tf.constant([[1.0, 2.0, 3, 0, 4.0]])
+            )
+            return tf.repeat(prob, batch_size, axis=0)
+
+        max_length = 5
+        inputs = tf.constant([[0, 1], [1, 2]])
+        tf.random.set_seed(42)
+        output_logit = top_k_search(
+            token_logits_fn,
+            inputs,
+            max_length=max_length,
+            k=3,
+            from_logits=True,
+            seed=42,
+        )
+        tf.random.set_seed(42)
+        output_probs = top_k_search(
+            token_probability_fn,
+            inputs,
+            max_length=max_length,
+            k=3,
+            from_logits=False,
+            seed=42,
+        )
+        self.assertAllEqual(output_logit, output_probs)