Generation: fix test (#34369)

* fix test

* fix copies

tests/generation/test_utils.py

@@ -671,29 +671,6 @@ class GenerationTesterMixin:
             else:
                 self.assertTrue(output_generate.shape[-1] == self.max_new_tokens + inputs_dict["input_ids"].shape[-1])
 
-            # for VLMs inputs embeds won't match input ids unless images are encoded and merged with ids properly
-            # no quick fix available, since obtaining image embeddings step is very model-specific
-            if any(name in model.__class__.__name__.lower() for name in ("blip", "llava", "paligemma")):
-                prepare_inputs_for_generation_args = set(
-                    inspect.signature(model.prepare_inputs_for_generation).parameters
-                )
-                # `inputs_embeds` input is well supported when `cache_positions` is used, because it means the modeling
-                # code is up to date with our most recent standards
-                if (
-                    "inputs_embeds" in prepare_inputs_for_generation_args
-                    and "cache_positions" in prepare_inputs_for_generation_args
-                ):
-                    input_embeds = model.get_input_embeddings()(inputs_dict["input_ids"])
-                    beam_kwargs.update({"inputs_embeds": input_embeds})
-                    output_generate2 = self._beam_sample_generate(
-                        model=model,
-                        input_ids=None,
-                        inputs_dict={},
-                        beam_kwargs=beam_kwargs,
-                    )
-
-                    torch.testing.assert_close(output_generate[:, input_embeds.shape[1] :], output_generate2)
-
     @pytest.mark.generate
     def test_beam_sample_generate_dict_output(self):
         for model_class in self.all_generative_model_classes:
@@ -1570,7 +1547,8 @@ class GenerationTesterMixin:
                     )
 
     @pytest.mark.generate
-    def test_generate_from_inputs_embeds_decoder_only(self):
+    @parameterized.expand([(1,), (2,)])
+    def test_generate_from_inputs_embeds_decoder_only(self, num_beams):
         # When supported, tests that the decoder model can generate from `inputs_embeds` instead of `input_ids`
         # if fails, you should probably update the `prepare_inputs_for_generation` function
         for model_class in self.all_generative_model_classes:
@@ -1597,11 +1575,15 @@ class GenerationTesterMixin:
                 continue
 
             input_ids = inputs_dict.pop("input_ids")
+            generation_kwargs = {
+                "return_dict_in_generate": True,
+                "output_scores": True,
+                "num_beams": num_beams,
+                "do_sample": False,
+            }
 
             # Traditional way of generating text
-            outputs_from_ids = model.generate(
+            outputs_from_ids = model.generate(input_ids, max_new_tokens=5, **generation_kwargs)
-                input_ids, max_new_tokens=5, return_dict_in_generate=True, output_scores=True
-            )
             self.assertEqual(outputs_from_ids.sequences.shape, (input_ids.shape[0], input_ids.shape[1] + 5))
 
             # Same thing, but from input embeddings (`input_ids` is passed so the prompt is present in the output)
@@ -1610,8 +1592,7 @@ class GenerationTesterMixin:
                 input_ids,
                 inputs_embeds=inputs_embeds,
                 max_new_tokens=5,
-                return_dict_in_generate=True,
+                **generation_kwargs,
-                output_scores=True,
             )
             self.assertListEqual(outputs_from_ids.sequences.tolist(), outputs_from_embeds.sequences.tolist())
 
@@ -1622,15 +1603,14 @@ class GenerationTesterMixin:
                 input_ids,
                 inputs_embeds=random_embeds,
                 max_new_tokens=5,
-                return_dict_in_generate=True,
+                **generation_kwargs,
-                output_scores=True,
             )
             for i in range(len(outputs_from_rand_embeds.scores)):
                 self.assertFalse(torch.allclose(outputs_from_embeds.scores[i], outputs_from_rand_embeds.scores[i]))
 
             # input_ids is not a required input -- if we don't pass it, the newly generated tokens will be the same
             outputs_from_embeds_wo_ids = model.generate(
-                inputs_embeds=inputs_embeds, max_new_tokens=5, return_dict_in_generate=True, output_scores=True
+                inputs_embeds=inputs_embeds, max_new_tokens=5, **generation_kwargs
             )
             self.assertListEqual(
                 outputs_from_embeds.sequences[:, inputs_embeds.shape[1] :].tolist(),

tests/models/idefics/test_modeling_idefics.py

@@ -773,7 +773,8 @@ class IdeficsForVisionText2TextTest(IdeficsModelTest, GenerationTesterMixin, uni
     @unittest.skip(
         reason="IDEFICS has specific requirements for working with inputs embeds like passing also the ids and pixels"
     )
-    def test_generate_from_inputs_embeds_decoder_only(self):
+    @parameterized.expand([(1,), (2,)])
+    def test_generate_from_inputs_embeds_decoder_only(self, num_beams):
         pass
 
     @unittest.skip(reason="IDEFICS cannot compile due to dynamic control flow when checking inputs")

tests/models/mamba2/test_modeling_mamba2.py

@@ -204,7 +204,8 @@ class Mamba2ModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMix
         pass
 
     @unittest.skip(reason="To fix, Mamba 2 cache slicing test case is an edge case")
-    def test_generate_from_inputs_embeds_decoder_only(self):
+    @parameterized.expand([(1,), (2,)])
+    def test_generate_from_inputs_embeds_decoder_only(self, num_beams):
         pass
 
     @unittest.skip(reason="To fix, Mamba 2 cache slicing test case is an edge case")

tests/models/moshi/test_modeling_moshi.py

@@ -656,16 +656,21 @@ class MoshiTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
                         )
 
     @pytest.mark.generate
-    def test_generate_from_inputs_embeds_decoder_only(self):
+    @parameterized.expand([(1,), (2,)])
+    def test_generate_from_inputs_embeds_decoder_only(self, num_beams):
         for model_class in self.all_generative_model_classes:
             config, input_ids, _, inputs_dict = self._get_input_ids_and_config()
 
             model = model_class(config).to(torch_device).eval()
+            generation_kwargs = {
+                "return_dict_in_generate": True,
+                "output_scores": True,
+                "num_beams": num_beams,
+                "do_sample": False,
+            }
 
             # Traditional way of generating text
-            outputs_from_ids = model.generate(
+            outputs_from_ids = model.generate(input_ids, max_new_tokens=5, **generation_kwargs, **inputs_dict)
-                input_ids, max_new_tokens=5, return_dict_in_generate=True, output_scores=True, **inputs_dict
-            )
             self.assertEqual(outputs_from_ids.sequences.shape, (input_ids.shape[0], input_ids.shape[1] + 5))
 
             # Same thing, but from input embeddings (`input_ids` is passed so the prompt is present in the output)
@@ -674,8 +679,7 @@ class MoshiTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
                 input_ids,
                 inputs_embeds=inputs_embeds,
                 max_new_tokens=5,
-                return_dict_in_generate=True,
+                **generation_kwargs,
-                output_scores=True,
                 **inputs_dict,
             )
 
@@ -686,8 +690,7 @@ class MoshiTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
                 input_ids,
                 inputs_embeds=random_embeds,
                 max_new_tokens=5,
-                return_dict_in_generate=True,
+                **generation_kwargs,
-                output_scores=True,
                 **inputs_dict,
             )
             for i in range(len(outputs_from_rand_embeds.scores)):
@@ -697,8 +700,7 @@ class MoshiTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
             outputs_from_embeds_wo_ids = model.generate(
                 inputs_embeds=inputs_embeds,
                 max_new_tokens=5,
-                return_dict_in_generate=True,
+                **generation_kwargs,
-                output_scores=True,
                 **inputs_dict,
             )
             self.assertListEqual(