From 421b8fb3353e7b044c7fb6843b4134f183e8a008 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Mon, 3 Jun 2024 14:52:02 +0200 Subject: [PATCH] fix test --- tests/openvino/test_modeling.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/openvino/test_modeling.py b/tests/openvino/test_modeling.py index f2a873007..6a6c34c23 100644 --- a/tests/openvino/test_modeling.py +++ b/tests/openvino/test_modeling.py @@ -1189,13 +1189,6 @@ def test_pipeline(self, model_arch): model.to("cpu") model.compile() - # Text2Text generation - pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer) - inputs = "This is a test" - outputs = pipe(inputs) - self.assertEqual(pipe.device, model.device) - self.assertIsInstance(outputs[0]["generated_text"], str) - # Summarization pipe = pipeline("summarization", model=model, tokenizer=tokenizer) outputs = pipe(inputs) @@ -1208,9 +1201,16 @@ def test_pipeline(self, model_arch): self.assertEqual(pipe.device, model.device) self.assertIsInstance(outputs[0]["translation_text"], str) - ov_pipe = optimum_pipeline("translation_en_to_fr", model_id, accelerator="openvino") + # Text2Text generation + pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer) + inputs = "This is a test" + outputs = pipe(inputs) + self.assertEqual(pipe.device, model.device) + self.assertIsInstance(outputs[0]["generated_text"], str) + + ov_pipe = optimum_pipeline("text2text-generation", model_id, accelerator="openvino") ov_outputs = ov_pipe(inputs) - self.assertEqual(outputs[-1]["translation_text"], ov_outputs[-1]["translation_text"]) + self.assertEqual(outputs[-1]["generated_text"], ov_outputs[-1]["generated_text"]) del ov_pipe del pipe del model