fix beam search with static cache

huggingface · Aug 30, 2024 · 4dd1494 · 4dd1494
1 parent 2b47ea9
commit 4dd1494
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py
@@ -1548,7 +1548,7 @@ def _prepare_cache_for_generation(
                     )
                 model_kwargs[cache_name] = self._get_cache(
                     cache_implementation=generation_config.cache_implementation,
-                    batch_size=generation_config.num_beams * generation_config.num_return_sequences * batch_size,
+                    batch_size=max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size,
                     max_cache_len=max_cache_length,
                     device=device,
                     model_kwargs=model_kwargs,

diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py
@@ -1420,7 +1420,7 @@ def _prepare_model_kwargs(input_ids, attention_mask, signature):
             signature = inspect.signature(model.forward).parameters.keys()
 
             # no cache as some models require special cache classes to be init outside forward
-            model.genertation_config.use_cache = False
+            model.generation_config.use_cache = False
 
             # Without padding
             model_kwargs = _prepare_model_kwargs(input_ids, attention_mask, signature)