[BugFix] Fix test_phi3v.py (#5725)

2024-06-20 21:45:34 -07:00
parent 1f5674218f
commit c35e4a3dd7
2 changed files with 9 additions and 5 deletions
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -233,11 +233,13 @@ class HfRunner:
        prompts: List[str],
        max_tokens: int,
        images: Optional[List[Image.Image]] = None,
+        **kwargs,
    ) -> List[Tuple[List[int], str]]:
        outputs = self.generate(prompts,
                                do_sample=False,
                                max_new_tokens=max_tokens,
-                                images=images)
+                                images=images,
+                                **kwargs)

        return [(output_ids[0], output_str[0])
                for output_ids, output_str in outputs]