Image Feature Extraction pipeline (#28216)
* Draft pipeline * Fixup * Fix docstrings * Update doctest * Update pipeline_model_mapping * Update docstring * Update tests * Update src/transformers/pipelines/image_feature_extraction.py Co-authored-by: Omar Sanseviero <osanseviero@gmail.com> * Fix docstrings - review comments * Remove pipeline mapping for composite vision models * Add to pipeline tests * Remove for flava (multimodal) * safe pil import * Add requirements for pipeline run * Account for super slow efficientnet * Review comments * Fix tests * Swap order of kwargs * Use build_pipeline_init_args * Add back FE pipeline for Vilt * Include image_processor_kwargs in docstring * Mark test as flaky * Update TODO * Update tests/pipelines/test_pipelines_image_feature_extraction.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * Add license header --------- Co-authored-by: Omar Sanseviero <osanseviero@gmail.com> Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com>
This commit is contained in:
@@ -130,7 +130,7 @@ class EfficientNetModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.Test
|
||||
|
||||
all_model_classes = (EfficientNetModel, EfficientNetForImageClassification) if is_torch_available() else ()
|
||||
pipeline_model_mapping = (
|
||||
{"feature-extraction": EfficientNetModel, "image-classification": EfficientNetForImageClassification}
|
||||
{"image-feature-extraction": EfficientNetModel, "image-classification": EfficientNetForImageClassification}
|
||||
if is_torch_available()
|
||||
else {}
|
||||
)
|
||||
@@ -216,6 +216,12 @@ class EfficientNetModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.Test
|
||||
model = EfficientNetModel.from_pretrained(model_name)
|
||||
self.assertIsNotNone(model)
|
||||
|
||||
@is_pipeline_test
|
||||
@require_vision
|
||||
@slow
|
||||
def test_pipeline_image_feature_extraction(self):
|
||||
super().test_pipeline_image_feature_extraction()
|
||||
|
||||
@is_pipeline_test
|
||||
@require_vision
|
||||
@slow
|
||||
|
||||
Reference in New Issue
Block a user