fix macos

Author: IlyasMoutawwakil

tests/pipelines/test_pipelines.py

@@ -22,10 +22,11 @@
 from transformers import AutoTokenizer
 from transformers.pipelines import Pipeline
 
-from optimum.pipelines import pipeline
+from optimum.pipelines import pipeline as optimum_pipeline
 from optimum.utils.testing_utils import remove_directory
 
 
+DEVICE = "cpu"
 GENERATE_KWARGS = {"max_new_tokens": 10, "min_new_tokens": 5, "do_sample": True}
 
 
@@ -49,7 +50,7 @@ def _create_dummy_audio(self) -> Dict[str, Any]:
 
     def test_text_classification_pipeline(self):
         """Test text classification ORT pipeline"""
-        pipe = pipeline(task="text-classification", accelerator="ort")
+        pipe = optimum_pipeline(task="text-classification", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         text = self._create_dummy_text()
         result = pipe(text)
@@ -61,7 +62,7 @@ def test_text_classification_pipeline(self):
 
     def test_token_classification_pipeline(self):
         """Test token classification ORT pipeline"""
-        pipe = pipeline(task="token-classification", accelerator="ort")
+        pipe = optimum_pipeline(task="token-classification", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         text = self._create_dummy_text()
         result = pipe(text)
@@ -74,7 +75,7 @@ def test_token_classification_pipeline(self):
 
     def test_question_answering_pipeline(self):
         """Test question answering ORT pipeline"""
-        pipe = pipeline(task="question-answering", accelerator="ort")
+        pipe = optimum_pipeline(task="question-answering", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         question = "What animal jumps?"
         context = "The quick brown fox jumps over the lazy dog."
@@ -88,7 +89,7 @@ def test_question_answering_pipeline(self):
 
     def test_fill_mask_pipeline(self):
         """Test fill mask ORT pipeline"""
-        pipe = pipeline(task="fill-mask", accelerator="ort")
+        pipe = optimum_pipeline(task="fill-mask", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         text = "The weather is <mask> today."
         result = pipe(text)
@@ -100,7 +101,7 @@ def test_fill_mask_pipeline(self):
 
     def test_feature_extraction_pipeline(self):
         """Test feature extraction ORT pipeline"""
-        pipe = pipeline(task="feature-extraction", accelerator="ort")
+        pipe = optimum_pipeline(task="feature-extraction", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         text = self._create_dummy_text()
         result = pipe(text)
@@ -111,7 +112,7 @@ def test_feature_extraction_pipeline(self):
 
     def test_text_generation_pipeline(self):
         """Test text generation ORT pipeline"""
-        pipe = pipeline(task="text-generation", accelerator="ort")
+        pipe = optimum_pipeline(task="text-generation", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         text = "The future of AI is"
         result = pipe(text, **GENERATE_KWARGS)
@@ -123,7 +124,7 @@ def test_text_generation_pipeline(self):
 
     def test_summarization_pipeline(self):
         """Test summarization ORT pipeline"""
-        pipe = pipeline(task="summarization", accelerator="ort")
+        pipe = optimum_pipeline(task="summarization", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         text = "The quick brown fox jumps over the lazy dog."
         result = pipe(text, **GENERATE_KWARGS)
@@ -134,7 +135,7 @@ def test_summarization_pipeline(self):
 
     def test_translation_pipeline(self):
         """Test translation ORT pipeline"""
-        pipe = pipeline(task="translation_en_to_de", accelerator="ort")
+        pipe = optimum_pipeline(task="translation_en_to_de", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         text = "Hello, how are you?"
         result = pipe(text, **GENERATE_KWARGS)
@@ -145,7 +146,7 @@ def test_translation_pipeline(self):
 
     def test_text2text_generation_pipeline(self):
         """Test text2text generation ORT pipeline"""
-        pipe = pipeline(task="text2text-generation", accelerator="ort")
+        pipe = optimum_pipeline(task="text2text-generation", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         text = "translate English to German: Hello, how are you?"
         result = pipe(text, **GENERATE_KWARGS)
@@ -156,7 +157,7 @@ def test_text2text_generation_pipeline(self):
 
     def test_zero_shot_classification_pipeline(self):
         """Test zero shot classification ORT pipeline"""
-        pipe = pipeline(task="zero-shot-classification", accelerator="ort")
+        pipe = optimum_pipeline(task="zero-shot-classification", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         text = "This is a great movie with excellent acting."
         candidate_labels = ["positive", "negative", "neutral"]
@@ -169,7 +170,7 @@ def test_zero_shot_classification_pipeline(self):
 
     def test_image_classification_pipeline(self):
         """Test image classification ORT pipeline"""
-        pipe = pipeline(task="image-classification", accelerator="ort")
+        pipe = optimum_pipeline(task="image-classification", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         image = self._create_dummy_image()
         result = pipe(image)
@@ -181,7 +182,7 @@ def test_image_classification_pipeline(self):
 
     def test_image_segmentation_pipeline(self):
         """Test image segmentation ORT pipeline"""
-        pipe = pipeline(task="image-segmentation", accelerator="ort")
+        pipe = optimum_pipeline(task="image-segmentation", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         image = self._create_dummy_image()
         result = pipe(image)
@@ -194,7 +195,7 @@ def test_image_segmentation_pipeline(self):
 
     def test_image_to_text_pipeline(self):
         """Test image to text ORT pipeline"""
-        pipe = pipeline(task="image-to-text", accelerator="ort")
+        pipe = optimum_pipeline(task="image-to-text", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         image = self._create_dummy_image()
         result = pipe(image, generate_kwargs=GENERATE_KWARGS)
@@ -205,7 +206,7 @@ def test_image_to_text_pipeline(self):
 
     def test_image_to_image_pipeline(self):
         """Test image to image ORT pipeline"""
-        pipe = pipeline(task="image-to-image", accelerator="ort")
+        pipe = optimum_pipeline(task="image-to-image", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         image = self._create_dummy_image()
         result = pipe(image)
@@ -215,7 +216,7 @@ def test_image_to_image_pipeline(self):
     # TODO: Enable when fixed in optimum-onnx
     # def test_automatic_speech_recognition_pipeline(self):
     #     """Test automatic speech recognition ORT pipeline"""
-    #     pipe = pipeline(task="automatic-speech-recognition", accelerator="ort")
+    #     pipe = optimum_pipeline(task="automatic-speech-recognition", accelerator="ort", device=DEVICE)
     #     audio = self._create_dummy_audio()
     #     result = pipe(audio, generate_kwargs=GENERATE_KWARGS)
 
@@ -224,7 +225,7 @@ def test_image_to_image_pipeline(self):
 
     def test_audio_classification_pipeline(self):
         """Test audio classification ORT pipeline"""
-        pipe = pipeline(task="audio-classification", accelerator="ort")
+        pipe = optimum_pipeline(task="audio-classification", accelerator="ort", device=DEVICE)
         self.assertIsInstance(pipe, Pipeline)
         audio = self._create_dummy_audio()
         result = pipe(audio)
@@ -240,7 +241,7 @@ def test_pipeline_with_ort_model(self):
 
         tokenizer = AutoTokenizer.from_pretrained("distilbert-base-cased")
         model = ORTModelForFeatureExtraction.from_pretrained("distilbert-base-cased", export=True)
-        pipe = pipeline(task="feature-extraction", model=model, tokenizer=tokenizer, accelerator="ort")
+        pipe = optimum_pipeline(task="feature-extraction", model=model, tokenizer=tokenizer, accelerator="ort")
         self.assertIsInstance(pipe, Pipeline)
         text = self._create_dummy_text()
         result = pipe(text)
@@ -251,7 +252,9 @@ def test_pipeline_with_ort_model(self):
 
     def test_pipeline_with_model_id(self):
         """Test ORT pipeline with a custom model id"""
-        pipe = pipeline(task="feature-extraction", model="distilbert-base-cased", accelerator="ort")
+        pipe = optimum_pipeline(
+            task="feature-extraction", model="distilbert-base-cased", accelerator="ort", device=DEVICE
+        )
         self.assertIsInstance(pipe, Pipeline)
         text = self._create_dummy_text()
         result = pipe(text)
@@ -262,13 +265,13 @@ def test_pipeline_with_model_id(self):
     def test_pipeline_with_invalid_task(self):
         """Test ORT pipeline with an unsupported task"""
         with self.assertRaises(KeyError) as context:
-            _ = pipeline(task="invalid-task", accelerator="ort")
+            _ = optimum_pipeline(task="invalid-task", accelerator="ort", device=DEVICE)
         self.assertIn("Unknown task invalid-task", str(context.exception))
 
     def test_pipeline_with_invalid_accelerator(self):
         """Test ORT pipeline with an unsupported accelerator"""
         with self.assertRaises(ValueError) as context:
-            _ = pipeline(task="feature-extraction", accelerator="invalid-accelerator")
+            _ = optimum_pipeline(task="feature-extraction", accelerator="invalid-accelerator", device=DEVICE)
         self.assertIn("Accelerator invalid-accelerator not recognized", str(context.exception))
 
     def tearDown(self):