TF generate refactor - Beam Search (#16374)

* refactor TF beam search

* refactored generate can now properly use attention masks

* add force bos/eos logit processors

docs/source/en/internal/generation_utils.mdx

@@ -178,6 +178,12 @@ generation.
 [[autodoc]] TFRepetitionPenaltyLogitsProcessor
     - __call__
 
+[[autodoc]] TFForcedBOSTokenLogitsProcessor
+    - __call__
+
+[[autodoc]] TFForcedEOSTokenLogitsProcessor
+    - __call__
+
 [[autodoc]] FlaxLogitsProcessor
     - __call__
 

src/transformers/__init__.py

@@ -1699,6 +1699,8 @@ if is_tf_available():
     _import_structure["benchmark.benchmark_args_tf"] = ["TensorFlowBenchmarkArguments"]
     _import_structure["benchmark.benchmark_tf"] = ["TensorFlowBenchmark"]
     _import_structure["generation_tf_logits_process"] = [
+        "TFForcedBOSTokenLogitsProcessor",
+        "TFForcedEOSTokenLogitsProcessor",
         "TFLogitsProcessor",
         "TFLogitsProcessorList",
         "TFLogitsWarper",
@@ -3827,6 +3829,8 @@ if TYPE_CHECKING:
         # Benchmarks
         from .benchmark.benchmark_tf import TensorFlowBenchmark
         from .generation_tf_logits_process import (
+            TFForcedBOSTokenLogitsProcessor,
+            TFForcedEOSTokenLogitsProcessor,
             TFLogitsProcessor,
             TFLogitsProcessorList,
             TFLogitsWarper,

src/transformers/generation_tf_logits_process.py

@@ -216,14 +216,10 @@ class TFMinLengthLogitsProcessor(TFLogitsProcessor):
         self.min_length = min_length
         self.eos_token_id = eos_token_id
 
-    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor) -> tf.Tensor:
-        # create boolean flag to decide if min length penalty should be applied
-        cur_len = input_ids.shape[-1]
-        apply_penalty = 1 - tf.clip_by_value(cur_len - self.min_length, 0, 1)
-
+    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
         # TODO(Matt) - this if statement has to be rewritten for XLA. Leaving it now though since
         # generate is not XLA - compileable anyways
-        if apply_penalty:
+        if cur_len < self.min_length:
             eos_token_id_mask = tf.broadcast_to(tf.range(scores.shape[-1]) == self.eos_token_id, scores.shape)
             scores = set_tensor_by_indices_to_value(scores, eos_token_id_mask, float("-inf"))
 
@@ -259,8 +255,8 @@ class TFRepetitionPenaltyLogitsProcessor(TFLogitsProcessor):
             np.put(token_penalties[i], prev_input_id, logit_penalties)
         return tf.convert_to_tensor(token_penalties, dtype=tf.float32)
 
-    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor) -> tf.Tensor:
-        score_penalties = self._create_score_penalties(input_ids, scores)
+    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+        score_penalties = self._create_score_penalties(input_ids[:, :cur_len], scores)
 
         scores = tf.math.multiply(scores, score_penalties)
 
@@ -330,12 +326,12 @@ class TFNoBadWordsLogitsProcessor(TFLogitsProcessor):
 
         return banned_tokens
 
-    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor) -> tf.Tensor:
+    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
 
         vocab_size = scores.shape[-1]
 
         # calculate a list of banned tokens according to bad words
-        banned_tokens = self.calc_banned_bad_words_ids(input_ids)
+        banned_tokens = self.calc_banned_bad_words_ids(input_ids[:, :cur_len])
 
         banned_tokens_indices_mask = []
         for banned_tokens_slice in banned_tokens:
@@ -365,12 +361,13 @@ class TFNoRepeatNGramLogitsProcessor(TFLogitsProcessor):
             raise ValueError(f"`ngram_size` has to be a strictly positive integer, but is {ngram_size}")
         self.ngram_size = ngram_size
 
-    def calc_banned_ngram_tokens(self, prev_input_ids, num_hypos, cur_len):
+    def calc_banned_ngram_tokens(self, input_ids, num_hypos, cur_len):
         # Copied from fairseq for no_repeat_ngram in beam_search
         if cur_len + 1 < self.ngram_size:
             # return no banned tokens if we haven't generated ngram_size tokens yet
             return [[] for _ in range(num_hypos)]
         generated_ngrams = [{} for _ in range(num_hypos)]
+        prev_input_ids = input_ids[:, :cur_len]
         for idx in range(num_hypos):
             gen_tokens = prev_input_ids[idx].numpy().tolist()
             generated_ngram = generated_ngrams[idx]
@@ -388,10 +385,9 @@ class TFNoRepeatNGramLogitsProcessor(TFLogitsProcessor):
 
         return banned_tokens
 
-    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor) -> tf.Tensor:
+    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
 
         batch_size, vocab_size = scores.shape
-        cur_len = input_ids.shape[-1]
         banned_tokens = self.calc_banned_ngram_tokens(input_ids, batch_size, cur_len)
 
         # create banned_tokens boolean mask
@@ -406,3 +402,66 @@ class TFNoRepeatNGramLogitsProcessor(TFLogitsProcessor):
         )
 
         return scores
+
+
+class TFForcedBOSTokenLogitsProcessor(TFLogitsProcessor):
+    r"""
+    [`TFLogitsProcessor`] that enforces the specified token as the first generated token.
+
+    Args:
+        bos_token_id (`int`):
+            The id of the token to force as the first generated token.
+    """
+
+    def __init__(self, bos_token_id: int):
+        if bos_token_id < 0:
+            raise ValueError(f"The forced bos token id  must be a non-negative integer, got {bos_token_id}")
+        self.bos_token_id = bos_token_id
+
+    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+        if cur_len == 1:
+            batch_size, num_tokens = scores.shape
+            # sets the score to 0 in the bos_token_id column
+            scores = tf.zeros((batch_size, 1))
+            # sets the score to -inf everywhere else
+            if self.bos_token_id > 0:
+                scores = tf.concat((tf.broadcast_to(-float("inf"), (batch_size, self.bos_token_id)), scores), axis=-1)
+            if self.bos_token_id < (num_tokens - 1):
+                scores = tf.concat(
+                    (scores, tf.broadcast_to(-float("inf"), (batch_size, (num_tokens - 1) - self.bos_token_id))),
+                    axis=-1,
+                )
+        return scores
+
+
+class TFForcedEOSTokenLogitsProcessor(TFLogitsProcessor):
+    r"""
+    [`TFLogitsProcessor`] that enforces the specified token as the last generated token when `max_length` is reached.
+
+    Args:
+        max_length (`int`):
+            The maximum length of the sequence to be generated.
+        eos_token_id (`int`):
+            The id of the token to force as the last generated token when `max_length` is reached.
+    """
+
+    def __init__(self, max_length: int, eos_token_id: int):
+        self.max_length = max_length
+        if eos_token_id < 0:
+            raise ValueError(f"The forced eos token id must be a non-negative integer, got {eos_token_id}")
+        self.eos_token_id = eos_token_id
+
+    def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+        if cur_len == self.max_length - 1:
+            batch_size, num_tokens = scores.shape
+            # sets the score to 0 in the eos_token_id column
+            scores = tf.zeros((batch_size, 1))
+            # sets the score to -inf everywhere else
+            if self.eos_token_id > 0:
+                scores = tf.concat((tf.broadcast_to(-float("inf"), (batch_size, self.eos_token_id)), scores), axis=-1)
+            if self.eos_token_id < (num_tokens - 1):
+                scores = tf.concat(
+                    (scores, tf.broadcast_to(-float("inf"), (batch_size, (num_tokens - 1) - self.eos_token_id))),
+                    axis=-1,
+                )
+        return scores

src/transformers/models/rag/modeling_tf_rag.py

@@ -1245,7 +1245,10 @@ class TFRagTokenForGeneration(TFRagPreTrainedModel, TFCausalLanguageModelingLoss
                 no_repeat_ngram_size=no_repeat_ngram_size,
                 bad_words_ids=bad_words_ids,
                 min_length=min_length,
+                max_length=max_length,
                 eos_token_id=eos_token_id,
+                forced_bos_token_id=None,
+                forced_eos_token_id=None,
             )
             model_kwargs["attention_mask"] = context_attention_mask
 

src/transformers/utils/dummy_tf_objects.py

@@ -17,6 +17,20 @@ class TensorFlowBenchmark(metaclass=DummyObject):
         requires_backends(self, ["tf"])
 
 
+class TFForcedBOSTokenLogitsProcessor(metaclass=DummyObject):
+    _backends = ["tf"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["tf"])
+
+
+class TFForcedEOSTokenLogitsProcessor(metaclass=DummyObject):
+    _backends = ["tf"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["tf"])
+
+
 class TFLogitsProcessor(metaclass=DummyObject):
     _backends = ["tf"]
 

tests/generation/test_generation_logits_process.py

@@ -472,14 +472,14 @@ class LogitsProcessorTest(unittest.TestCase):
 
         logits_processor = ForcedEOSTokenLogitsProcessor(max_length=max_length, eos_token_id=eos_token_id)
 
-        # check that all scores are -inf except the eos_token_id when max_length is reached
+        # check that all scores are -inf except the eos_token_id when max_length-1 is reached
         input_ids = ids_tensor((batch_size, 4), vocab_size=20)
         scores = self._get_uniform_logits(batch_size, vocab_size)
         scores = logits_processor(input_ids, scores)
         self.assertTrue(torch.isneginf(scores[:, eos_token_id + 1 :]).all())
         self.assertListEqual(scores[:, eos_token_id].tolist(), 4 * [0])  # score for eos_token_id should be zero
 
-        # check that eos_token_id is not forced if max_length is not reached
+        # check that eos_token_id is not forced if max_length-1 is not reached
         input_ids = ids_tensor((batch_size, 3), vocab_size=20)
         scores = self._get_uniform_logits(batch_size, vocab_size)
         scores = logits_processor(input_ids, scores)

tests/generation/test_generation_tf_logits_process.py

@@ -26,6 +26,8 @@ if is_tf_available():
     import tensorflow as tf
 
     from transformers.generation_tf_logits_process import (
+        TFForcedBOSTokenLogitsProcessor,
+        TFForcedEOSTokenLogitsProcessor,
         TFLogitsProcessorList,
         TFMinLengthLogitsProcessor,
         TFNoBadWordsLogitsProcessor,
@@ -43,7 +45,7 @@ if is_tf_available():
 @require_tf
 class TFLogitsProcessorTest(unittest.TestCase):
     def _get_uniform_logits(self, batch_size: int, length: int):
-        scores = np.ones((batch_size, length), dtype=np.float32) / length
+        scores = tf.ones((batch_size, length), dtype=tf.float32) / length
         return scores
 
     def test_min_length_dist_processor(self):
@@ -54,15 +56,17 @@ class TFLogitsProcessorTest(unittest.TestCase):
         min_dist_processor = TFMinLengthLogitsProcessor(min_length=10, eos_token_id=eos_token_id)
 
         # check that min length is applied at length 5
-        input_ids = ids_tensor((batch_size, 5), vocab_size=20)
+        cur_len = 5
+        input_ids = ids_tensor((batch_size, cur_len), vocab_size=20)
         scores = self._get_uniform_logits(batch_size, vocab_size)
-        scores_before_min_length = min_dist_processor(input_ids, scores)
+        scores_before_min_length = min_dist_processor(input_ids, scores, cur_len)
         self.assertListEqual(scores_before_min_length[:, eos_token_id].numpy().tolist(), 4 * [-float("inf")])
 
         # check that min length is not applied anymore at length 15
-        input_ids = ids_tensor((batch_size, 15), vocab_size=20)
+        cur_len = 15
+        input_ids = ids_tensor((batch_size, cur_len), vocab_size=20)
         scores = self._get_uniform_logits(batch_size, vocab_size)
-        scores_before_min_length = min_dist_processor(input_ids, scores)
+        scores_before_min_length = min_dist_processor(input_ids, scores, cur_len)
         self.assertFalse(tf.math.reduce_any(tf.math.is_inf(scores_before_min_length)).numpy())
 
     def test_temperature_dist_warper(self):
@@ -72,8 +76,10 @@ class TFLogitsProcessorTest(unittest.TestCase):
         scores = self._get_uniform_logits(batch_size=2, length=length)
 
         # tweak scores to not be uniform anymore
+        scores = scores.numpy()
         scores[1, 5] = (1 / length) + 0.1  # peak, 1st batch
         scores[1, 10] = (1 / length) - 0.4  # valley, 1st batch
+        scores = tf.convert_to_tensor(scores)
 
         # compute softmax
         probs = tf.nn.softmax(scores, axis=-1)
@@ -97,8 +103,11 @@ class TFLogitsProcessorTest(unittest.TestCase):
         self.assertLess(tf.math.reduce_min(probs[1, :]), tf.math.reduce_min(warped_prob_smooth[1, :]))
 
     def test_repetition_penalty_dist_process(self):
-        input_ids = tf.constant([[0, 1], [5, 0]], dtype=tf.int32)
         vocab_size = 10
+        cur_len = 2
+
+        input_ids = tf.constant([[0, 1], [5, 0]], dtype=tf.int32)
+        self.assertEqual(cur_len, input_ids.shape[1])
 
         scores = self._get_uniform_logits(batch_size=2, length=vocab_size)
 
@@ -109,7 +118,7 @@ class TFLogitsProcessorTest(unittest.TestCase):
 
         rep_penalty_proc = TFRepetitionPenaltyLogitsProcessor(penalty=2.0)
 
-        scores = rep_penalty_proc(input_ids, tf.identity(scores))
+        scores = rep_penalty_proc(input_ids, tf.identity(scores), cur_len)
 
         # check that values were correctly changed
         self.assertAlmostEqual(scores[0, 0].numpy(), -(1 / vocab_size) * 2)
@@ -188,15 +197,18 @@ class TFLogitsProcessorTest(unittest.TestCase):
     def test_no_repeat_ngram_dist_processor(self):
         vocab_size = 3
         batch_size = 2
+        cur_len = 4
 
         input_ids = tf.constant([[1, 1, 2, 1], [0, 1, 0, 1]], dtype=tf.int32)
+        self.assertEqual(cur_len, input_ids.shape[1])
+
         scores = self._get_uniform_logits(batch_size, vocab_size)
 
         no_repeat_proc_2_gram = TFNoRepeatNGramLogitsProcessor(2)
         no_repeat_proc_3_gram = TFNoRepeatNGramLogitsProcessor(3)
 
-        filtered_scores_2_gram = no_repeat_proc_2_gram(input_ids, tf.identity(scores))
-        filtered_scores_3_gram = no_repeat_proc_3_gram(input_ids, tf.identity(scores))
+        filtered_scores_2_gram = no_repeat_proc_2_gram(input_ids, tf.identity(scores), cur_len)
+        filtered_scores_3_gram = no_repeat_proc_3_gram(input_ids, tf.identity(scores), cur_len)
 
         # 2-gram would forbid 2nd and 3rd token (1,2) at 1st batch and 1st token (0) at 2nd batch
         self.assertListEqual(
@@ -212,14 +224,17 @@ class TFLogitsProcessorTest(unittest.TestCase):
         vocab_size = 5
         batch_size = 2
         eos_token_id = 4
+        cur_len = 4
 
         input_ids = tf.constant([[0, 1, 3, 1], [0, 1, 0, 1]], dtype=tf.int32)
+        self.assertEqual(cur_len, input_ids.shape[1])
+
         bad_word_tokens = [[1], [4], [1, 0], [0, 1, 2], [1, 3, 1, 3]]
         scores = self._get_uniform_logits(batch_size, vocab_size)
 
         no_bad_words_dist_proc = TFNoBadWordsLogitsProcessor(bad_words_ids=bad_word_tokens, eos_token_id=eos_token_id)
 
-        filtered_scores = no_bad_words_dist_proc(input_ids, tf.identity(scores))
+        filtered_scores = no_bad_words_dist_proc(input_ids, tf.identity(scores), cur_len)
 
         # batch 1: 1st, 2nd, and 4th (0, 1, 3) token are forbidden
         # batch 2: 1st, 2nd, and 3rd (0, 1, 2) token are forbidden
@@ -228,14 +243,65 @@ class TFLogitsProcessorTest(unittest.TestCase):
             [[True, True, False, True, True], [True, True, True, False, True]],
         )
 
+    def test_forced_bos_token_logits_processor(self):
+        vocab_size = 20
+        batch_size = 4
+        bos_token_id = 0
+
+        logits_processor = TFForcedBOSTokenLogitsProcessor(bos_token_id=bos_token_id)
+
+        # check that all scores are -inf except the bos_token_id score
+        cur_len = 1
+        input_ids = ids_tensor((batch_size, cur_len), vocab_size=20)
+        scores = self._get_uniform_logits(batch_size, vocab_size)
+        scores = logits_processor(input_ids, scores, cur_len)
+        self.assertTrue(
+            tf.math.reduce_all(tf.math.is_inf(scores[:, bos_token_id + 1 :]) & (scores[:, bos_token_id + 1 :] < 0))
+        )
+        self.assertListEqual(scores[:, bos_token_id].numpy().tolist(), 4 * [0])  # score for bos_token_id shold be zero
+
+        # check that bos_token_id is not forced if current length is greater than 1
+        cur_len = 4
+        input_ids = ids_tensor((batch_size, cur_len), vocab_size=20)
+        scores = self._get_uniform_logits(batch_size, vocab_size)
+        scores = logits_processor(input_ids, scores, cur_len)
+        self.assertFalse(tf.math.reduce_any(tf.math.is_inf((scores))))
+
+    def test_forced_eos_token_logits_processor(self):
+        vocab_size = 20
+        batch_size = 4
+        eos_token_id = 0
+        max_length = 5
+
+        logits_processor = TFForcedEOSTokenLogitsProcessor(max_length=max_length, eos_token_id=eos_token_id)
+
+        # check that all scores are -inf except the eos_token_id when max_length-1 is reached
+        cur_len = 4
+        input_ids = ids_tensor((batch_size, cur_len), vocab_size=20)
+        scores = self._get_uniform_logits(batch_size, vocab_size)
+        scores = logits_processor(input_ids, scores, cur_len)
+        self.assertTrue(
+            tf.math.reduce_all(tf.math.is_inf(scores[:, eos_token_id + 1 :]) & (scores[:, eos_token_id + 1 :] < 0))
+        )
+        self.assertListEqual(
+            scores[:, eos_token_id].numpy().tolist(), 4 * [0]
+        )  # score for eos_token_id should be zero
+
+        # check that eos_token_id is not forced if max_length-1 is not reached
+        cur_len = 3
+        input_ids = ids_tensor((batch_size, cur_len), vocab_size=20)
+        scores = self._get_uniform_logits(batch_size, vocab_size)
+        scores = logits_processor(input_ids, scores, cur_len)
+        self.assertFalse(tf.math.reduce_any(tf.math.is_inf((scores))))
+
     def test_processor_list(self):
         batch_size = 4
-        sequence_length = 10
+        cur_len = 10
         vocab_size = 15
         eos_token_id = 0
 
         # dummy input_ids and scores
-        input_ids = ids_tensor((batch_size, sequence_length), vocab_size)
+        input_ids = ids_tensor((batch_size, cur_len), vocab_size)
         input_ids_comp = tf.identity(input_ids)
 
         scores = self._get_uniform_logits(batch_size, vocab_size)
@@ -251,13 +317,13 @@ class TFLogitsProcessorTest(unittest.TestCase):
         no_bad_words_dist_proc = TFNoBadWordsLogitsProcessor(bad_words_ids=[[1]], eos_token_id=eos_token_id)
 
         # no processor list
-        scores = min_dist_proc(input_ids, scores)
+        scores = min_dist_proc(input_ids, scores, cur_len)
         scores = temp_dist_warp(input_ids, scores)
-        scores = rep_penalty_proc(input_ids, scores)
+        scores = rep_penalty_proc(input_ids, scores, cur_len)
         scores = top_k_warp(input_ids, scores)
         scores = top_p_warp(input_ids, scores)
-        scores = no_repeat_proc(input_ids, scores)
-        scores = no_bad_words_dist_proc(input_ids, scores)
+        scores = no_repeat_proc(input_ids, scores, cur_len)
+        scores = no_bad_words_dist_proc(input_ids, scores, cur_len)
 
         # with processor list
         processor = TFLogitsProcessorList(
@@ -271,7 +337,7 @@ class TFLogitsProcessorTest(unittest.TestCase):
                 no_bad_words_dist_proc,
             ]
         )
-        scores_comp = processor(input_ids, scores_comp)
+        scores_comp = processor(input_ids, scores_comp, cur_len=cur_len)
 
         # remove inf
         scores = set_tensor_by_indices_to_value(scores, tf.math.is_inf(scores), -1e9)

tests/gpt2/test_modeling_tf_gpt2.py

@@ -536,7 +536,6 @@ class TFGPT2ModelLanguageGenerationTest(unittest.TestCase):
             "bad_words_ids": [tokenizer("is").input_ids, tokenizer("angry about").input_ids],
             "no_repeat_ngram_size": 2,
             "do_sample": False,
-            "repetition_penalty": 1.3,
             "num_beams": 2,
         }
 
@@ -544,8 +543,8 @@ class TFGPT2ModelLanguageGenerationTest(unittest.TestCase):
 
         output_strings = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
         expected_output_string = [
-            "Today is a beautiful day and I hope you enjoy it.\nI am very happy to announce that",
-            "Yesterday was the first time I've ever seen a game where you can play with",
+            "Today is a beautiful day and a great day for all of us.\n\nI’m",
+            "Yesterday was the first day of the year for the second time in a row,",
         ]
         self.assertListEqual(output_strings, expected_output_string)
 

tests/speech_to_text/test_modeling_tf_speech_to_text.py

@@ -508,7 +508,7 @@ class TFSpeech2TextModelTest(TFModelTesterMixin, unittest.TestCase):
                 # if bos token id is not defined model needs input_ids, num_return_sequences = 1
                 self._check_generated_ids(model.generate(input_features, do_sample=True, num_beams=2))
 
-            with self.assertRaises(AssertionError):
+            with self.assertRaises(ValueError):
                 # generating more sequences than having beams leads is not possible
                 model.generate(input_features, do_sample=False, num_return_sequences=3, num_beams=2)
 

tests/test_modeling_tf_common.py

@@ -1179,7 +1179,7 @@ class TFModelTesterMixin:
                 # num_return_sequences = 1
                 self._check_generated_ids(model.generate(do_sample=True, max_length=5, num_beams=2))
 
-            with self.assertRaises(AssertionError):
+            with self.assertRaises(ValueError):
                 # generating more sequences than having beams leads is not possible
                 model.generate(input_ids, do_sample=False, num_return_sequences=3, num_beams=2)