update

Author: ydshieh

tests/generation/test_utils.py

@@ -1263,6 +1263,9 @@ def test_dola_decoding_sample(self):
 
             if model.get_output_embeddings() is None:
                 self.skipTest("DoLa is not supported for models that don't have output embeddings")
+
+            logits_processor_kwargs = self._get_logits_processor_kwargs(do_sample=True, config=model.config)
+
             # Sets dola generation arguments such that:
             # a) no EOS is generated, to ensure generation doesn't break early
             # b) there are at least two forward passes in the main model, to ensure the input preparation of
@@ -1279,9 +1282,8 @@ def test_dola_decoding_sample(self):
                 "return_dict_in_generate": True,
                 "use_cache": getattr(config, "use_cache", False),  # Some models don't support the cache
                 "dola_layers": "low",
-                "bad_words_ids": [[model.config.image_token_index]] if hasattr(model.config, "image_token_index") else None,
             }
-            output_dola = model.generate(**generation_kwargs, **inputs_dict)
+            output_dola = model.generate(**generation_kwargs, **logits_processor_kwargs, **inputs_dict)
             self._check_outputs(output_dola, model.config, use_cache=getattr(config, "use_cache", False))
 
     @pytest.mark.generate