add auto model loading to test

huggingface · Sep 28, 2023 · cb04a62 · cb04a62
1 parent 23b9627
commit cb04a62
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 9 deletions.
diff --git a/optimum/intel/openvino/utils.py b/optimum/intel/openvino/utils.py
@@ -78,6 +78,7 @@
 
 
 _HEAD_TO_AUTOMODELS = {
+    "feature-extraction" : "OVModelForFeatureExtraction",
     "fill-mask": "OVModelForMaskedLM",
     "text-generation": "OVModelForCausalLM",
     "text2text-generation": "OVModelForSeq2SeqLM",
@@ -87,6 +88,7 @@
     "image-classification": "OVModelForImageClassification",
     "audio-classification": "OVModelForAudioClassification",
     "stable-diffusion": "OVStableDiffusionPipeline",
+    "stable-diffusion-xl": "OVStableDiffusionXLPipeline",
 }
 
 

diff --git a/tests/openvino/test_exporters_cli.py b/tests/openvino/test_exporters_cli.py
@@ -19,23 +19,38 @@
 from utils_tests import MODEL_NAMES
 
 from optimum.exporters.openvino.__main__ import main_export
-
+from optimum.intel.openvino.utils import _HEAD_TO_AUTOMODELS
+from optimum.intel import (
+    OVModelForAudioClassification,
+    OVModelForCausalLM,
+    OVModelForFeatureExtraction,
+    OVModelForImageClassification,
+    OVModelForMaskedLM,
+    OVModelForQuestionAnswering,
+    OVModelForSeq2SeqLM,
+    OVModelForSequenceClassification,
+    OVModelForTokenClassification,
+    OVStableDiffusionPipeline,
+    OVStableDiffusionXLPipeline,
+)
 
 class OVCLIExportTestCase(unittest.TestCase):
     """
     Integration tests ensuring supported models are correctly exported.
     """
 
     SUPPORTED_ARCHITECTURES = (
-        ["causal-lm", "gpt2"],
-        ["causal-lm-with-past", "gpt2"],
-        ["seq2seq-lm", "t5"],
-        ["seq2seq-lm-with-past", "t5"],
-        ["sequence-classification", "bert"],
+        ["text-generation", "gpt2"],
+        ["text-generation-with-past", "gpt2"],
+        ["text2text-generation", "t5"],
+        ["text2text-generation-with-past", "t5"],
+        ["text-classification", "bert"],
         ["question-answering", "distilbert"],
-        ["masked-lm", "bert"],
-        ["default", "blenderbot"],
-        ["default-with-past", "blenderbot"],
+        ["token-classification", "roberta"],
+        ["image-classification", "vit"],
+        ["audio-classification", "wav2vec2"],
+        ["fill-mask", "bert"],
+        ["feature-extraction", "blenderbot"],
         ["stable-diffusion", "stable-diffusion"],
         ["stable-diffusion-xl", "stable-diffusion-xl"],
         ["stable-diffusion-xl", "stable-diffusion-xl-refiner"],
@@ -57,3 +72,5 @@ def test_exporters_cli(self, task: str, model_type: str):
                 shell=True,
                 check=True,
             )
+            model_kwargs = {"use_cache" : task.endswith("with-past")} if "generation" in task else {}
+            ov_model = eval(_HEAD_TO_AUTOMODELS[task.replace("-with-past", "")]).from_pretrained(tmpdirname, **model_kwargs)