Avoid failure TFBlipModelTest::test_pipeline_image_to_text (#31827)

* fix * fix * fix --------- Co-authored-by: ydshieh <[email protected]>
huggingface · Jul 8, 2024 · 4879ac2 · 4879ac2
1 parent ba74370
commit 4879ac2
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/tests/models/efficientnet/test_modeling_efficientnet.py b/tests/models/efficientnet/test_modeling_efficientnet.py
@@ -215,6 +215,12 @@ def test_model_from_pretrained(self):
     def test_pipeline_image_feature_extraction(self):
         super().test_pipeline_image_feature_extraction()
 
+    @is_pipeline_test
+    @require_vision
+    @slow
+    def test_pipeline_image_feature_extraction_fp16(self):
+        super().test_pipeline_image_feature_extraction_fp16()
+
     @is_pipeline_test
     @require_vision
     @slow

diff --git a/tests/models/vits/test_modeling_vits.py b/tests/models/vits/test_modeling_vits.py
@@ -181,6 +181,10 @@ def test_config(self):
     def test_pipeline_feature_extraction(self):
         super().test_pipeline_feature_extraction()
 
+    @is_flaky(description="torch 2.2.0 gives `Timeout >120.0s`")
+    def test_pipeline_feature_extraction_fp16(self):
+        super().test_pipeline_feature_extraction_fp16()
+
     @unittest.skip(reason="Need to fix this after #26538")
     def test_model_forward(self):
         set_seed(12345)

diff --git a/tests/pipelines/test_pipelines_image_to_text.py b/tests/pipelines/test_pipelines_image_to_text.py
@@ -17,7 +17,7 @@
 import requests
 
 from transformers import MODEL_FOR_VISION_2_SEQ_MAPPING, TF_MODEL_FOR_VISION_2_SEQ_MAPPING, is_vision_available
-from transformers.pipelines import pipeline
+from transformers.pipelines import ImageToTextPipeline, pipeline
 from transformers.testing_utils import (
     is_pipeline_test,
     require_tf,
@@ -46,8 +46,8 @@ class ImageToTextPipelineTests(unittest.TestCase):
     tf_model_mapping = TF_MODEL_FOR_VISION_2_SEQ_MAPPING
 
     def get_test_pipeline(self, model, tokenizer, processor, torch_dtype="float32"):
-        pipe = pipeline(
-            "image-to-text", model=model, tokenizer=tokenizer, image_processor=processor, torch_dtype=torch_dtype
+        pipe = ImageToTextPipeline(
+            model=model, tokenizer=tokenizer, image_processor=processor, torch_dtype=torch_dtype
         )
         examples = [
             Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"),