fix generation test

huggingface · Oct 11, 2023 · d59bcc6 · d59bcc6
1 parent 6aa60fe
commit d59bcc6
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py
@@ -2953,7 +2953,8 @@ def forward(self, input_ids, foo=False, **kwargs):
 
                 return outs
 
-            def prepare_inputs_for_generation(self, *args, foo=False, **kwargs):
+            def prepare_inputs_for_generation(self, *args, foo=False, encoder_outputs=None, **kwargs):
+                kwargs["encoder_outputs"] = encoder_outputs
                 inputs = super().prepare_inputs_for_generation(*args, **kwargs)
 
                 inputs["foo"] = foo
@@ -2992,3 +2993,14 @@ def prepare_inputs_for_generation(self, *args, foo=False, **kwargs):
             assistant_model=assistant,
         )
         self.assertListEqual(outputs_assisted.tolist(), outputs_foo.tolist())
+
+        # Check that passing encoder_outputs directly also works as expected
+        encoder_outputs = assistant.get_encoder()(input_ids)
+
+        outputs_assisted = model.generate(
+            foo=True,
+            assistant_model=assistant,
+            encoder_outputs=encoder_outputs,
+            assistant_encoder_outputs=encoder_outputs,
+        )
+        self.assertListEqual(outputs_assisted.tolist(), outputs_foo.tolist())