Fix the truncation strategy

Author: chenmoneygithub

keras_nlp/layers/masked_lm_mask_generator.py

@@ -147,11 +147,7 @@ def call(self, inputs):
             # convert dense to ragged.
             inputs = tf.RaggedTensor.from_tensor(inputs)
 
-        (
-            token_ids,
-            mask_positions,
-            mask_ids,
-        ) = tf_text.mask_language_model(
+        (token_ids, mask_positions, mask_ids,) = tf_text.mask_language_model(
             inputs,
             item_selector=self._random_selector,
             mask_values_chooser=self._mask_values_chooser,

keras_nlp/samplers/beam_sampler_test.py

@@ -125,13 +125,25 @@ def token_probability_fn(inputs, mask):
             prob = tf.constant([[[0.0, 0.0, 0.0, 1.0]]])
             return tf.tile(prob, [batch_size, seq_length, 1])
 
-        max_length = 5
+        max_length = 4
         inputs = tf.constant([[0, 1], [1, 2]])
         outputs = self.sampler(
             inputs,
             token_probability_fn,
             max_length=max_length,
             end_token_id=2,
         )
-        expected_outputs = tf.ragged.constant([[0, 1, 3, 3, 3], [1]])
+        # end_token in prompt does not trigger truncation.
+        expected_outputs = tf.ragged.constant([[0, 1, 3, 3], [1, 2, 3, 3]])
+        self.assertAllEqual(outputs, expected_outputs)
+
+        max_length = 4
+        inputs = tf.constant([[0, 1], [1, 3]])
+        outputs = self.sampler(
+            inputs,
+            token_probability_fn,
+            max_length=max_length,
+            end_token_id=3,
+        )
+        expected_outputs = tf.ragged.constant([[0, 1], [1, 3]])
         self.assertAllEqual(outputs, expected_outputs)

keras_nlp/samplers/greedy_sampler_test.py

@@ -93,15 +93,14 @@ def token_probability_fn(inputs, mask):
         )
 
     def test_end_token_id(self):
-        max_length = 5
-
         def token_probability_fn(inputs, mask):
             batch_size = inputs.shape[0]
             prob = tf.constant([[[0.0, 0.0, 0.0, 1.0]]])
             return tf.repeat(
                 tf.repeat(prob, batch_size, axis=0), max_length, axis=1
             )
 
+        max_length = 4
         sampler = GreedySampler()
         inputs = tf.constant([[0, 1], [1, 2]])
         outputs = sampler(
@@ -110,7 +109,18 @@ def token_probability_fn(inputs, mask):
             max_length=max_length,
             end_token_id=2,
         )
-        expected_outputs = tf.ragged.constant([[0, 1, 3, 3, 3], [1]])
+        # end_token in prompt does not trigger truncation.
+        expected_outputs = tf.ragged.constant([[0, 1, 3, 3], [1, 2, 3, 3]])
+        self.assertAllEqual(outputs, expected_outputs)
+
+        outputs = sampler(
+            inputs,
+            token_probability_fn,
+            max_length=max_length,
+            end_token_id=3,
+        )
+        # Generated end_token will be truncated.
+        expected_outputs = tf.ragged.constant([[0, 1], [1, 2]])
         self.assertAllEqual(outputs, expected_outputs)
 
     def test_compare_xla_noxla_results(self):

keras_nlp/samplers/sampler.py

@@ -189,14 +189,27 @@ def _pad_prompt(self, prompt, max_length):
 
     def _mask_tokens_after_end_token(
         self,
-        prompt,
+        generated_result,
+        original_prompt,
         max_length,
         end_token_id,
     ):
         """Helper function to truncate the tokens after the end token."""
+        # Exclude original prompts from being truncated.
+        # Add a big int to tokens in original prompt to move token ids over
+        # the vocab size. Vocab size should not be large as 1e9.
+        increment_value = int(1e9)
+        increment = tf.cast(
+            tf.zeros_like(original_prompt) + increment_value,
+            dtype=tf.int32,
+        )
+        if not isinstance(increment, tf.RaggedTensor):
+            increment = tf.RaggedTensor.from_tensor(increment)
+        increment = increment.to_tensor(shape=tf.shape(generated_result))
+        generated_result += increment
         # 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)
+        end_indices = tf.math.argmax(generated_result == end_token_id, -1)
         # Use max_length if no `end_token_id` is found.
         end_indices = tf.where(
             end_indices == 0,
@@ -205,7 +218,9 @@ def _mask_tokens_after_end_token(
         )
         # Truncate out tokens after (including) the end token.
         mask_indices = tf.sequence_mask(end_indices, maxlen=max_length)
-        return tf.ragged.boolean_mask(prompt, mask_indices)
+        # Revert the increment added earlier.
+        generated_result -= increment
+        return tf.ragged.boolean_mask(generated_result, mask_indices)
 
     def __call__(
         self,
@@ -217,10 +232,10 @@ def __call__(
         from_logits=True,
     ):
         prompt, mask = self._validate_prompt_and_mask(prompt, mask)
-
         input_is_1d = prompt.shape.rank == 1
         if input_is_1d:
             prompt = tf.RaggedTensor.from_tensor(prompt[tf.newaxis, :])
+        original_prompt = tf.identity(prompt)
 
         shortest_prompt_len = tf.reduce_min(prompt.row_lengths())
         # Pad prompt to be a dense Tensor of shape [batch_size, max_length].
@@ -247,6 +262,7 @@ def __call__(
         if end_token_id is not None:
             prompt = self._mask_tokens_after_end_token(
                 prompt,
+                original_prompt,
                 max_length,
                 end_token_id,
             )

keras_nlp/samplers/top_k_sampler_test.py

@@ -144,17 +144,27 @@ def token_probability_fn(inputs, mask):
             prob = tf.constant([[[0.0, 0.0, 0.0, 1.0]]])
             return tf.tile(prob, [batch_size, seq_length, 1])
 
-        max_length = 5
-        inputs = tf.constant([[0, 1], [1, 2]])
         tf.random.set_seed(42)
         sampler = TopKSampler(k=4, seed=42)
+        max_length = 4
+        inputs = tf.constant([[0, 1], [1, 2]])
         outputs = sampler(
             inputs,
             token_probability_fn,
             max_length=max_length,
             end_token_id=2,
             from_logits=False,
         )
-        # Top-k sampling result with seed 42.
-        expected_outputs = tf.ragged.constant([[0, 1, 3, 3, 3], [1]])
+        # end_token in prompt does not trigger truncation.
+        expected_outputs = tf.ragged.constant([[0, 1, 3, 3], [1, 2, 3, 3]])
+        self.assertAllEqual(outputs, expected_outputs)
+
+        outputs = sampler(
+            inputs,
+            token_probability_fn,
+            max_length=max_length,
+            end_token_id=3,
+        )
+        # Generated end_token will be truncated.
+        expected_outputs = tf.ragged.constant([[0, 1], [1, 2]])
         self.assertAllEqual(outputs, expected_outputs)

keras_nlp/samplers/top_p_sampler_test.py

@@ -144,17 +144,28 @@ def token_probability_fn(inputs, mask):
             prob = tf.constant([[[0.0, 0.0, 0.0, 1.0]]])
             return tf.tile(prob, [batch_size, seq_length, 1])
 
-        max_length = 5
-        inputs = tf.constant([[0, 1], [1, 2]])
         tf.random.set_seed(42)
         sampler = TopPSampler(p=0.1, seed=42)
+        max_length = 4
+        inputs = tf.constant([[0, 1], [1, 2]])
         outputs = sampler(
             inputs,
             token_probability_fn,
             max_length=max_length,
             end_token_id=2,
             from_logits=False,
         )
-        # Top-p sampling result with seed 42.
-        expected_outputs = tf.ragged.constant([[0, 1, 3, 3, 3], [1]])
+        # end_token in prompt does not trigger truncation.
+        expected_outputs = tf.ragged.constant([[0, 1, 3, 3], [1, 2, 3, 3]])
+        self.assertAllEqual(outputs, expected_outputs)
+
+        outputs = sampler(
+            inputs,
+            token_probability_fn,
+            max_length=max_length,
+            end_token_id=3,
+            from_logits=False,
+        )
+        # Generated end_token will be truncated.
+        expected_outputs = tf.ragged.constant([[0, 1], [1, 2]])
         self.assertAllEqual(outputs, expected_outputs)