Fixed import issues for onnx>=1.17 (openvinotoolkit#1301)

Wovchena · Dec 5, 2024 · e6a435f · e6a435f
1 parent 0b81108
commit e6a435f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 22 deletions.
diff --git a/tools/who_what_benchmark/whowhatbench/visualtext_evaluator.py b/tools/who_what_benchmark/whowhatbench/visualtext_evaluator.py
@@ -4,8 +4,6 @@
 import datasets
 import pandas as pd
 from diffusers.utils.loading_utils import load_image
-from optimum.intel.openvino.modeling_visual_language import \
-    MODEL_TYPE_TO_CLS_MAPPING
 from tqdm import tqdm
 from transformers import set_seed
 
@@ -114,6 +112,9 @@ def _generate_data(self, model, gen_answer_fn=None, generation_config=None):
         def default_gen_answer(
             model, prompt, image, processor, tokenizer, max_new_tokens, crop_question
         ):
+
+            from optimum.intel.openvino.modeling_visual_language import \
+                MODEL_TYPE_TO_CLS_MAPPING
             preprocess_inputs = MODEL_TYPE_TO_CLS_MAPPING[
                 model.config.model_type
             ].preprocess_inputs

diff --git a/tools/who_what_benchmark/whowhatbench/wwb.py b/tools/who_what_benchmark/whowhatbench/wwb.py
@@ -9,31 +9,20 @@
 import logging
 import os
 
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, AutoProcessor, AutoModel, AutoModelForVision2Seq
+import openvino as ov
+
 import pandas as pd
 from datasets import load_dataset
 from diffusers import DiffusionPipeline
-from optimum.exporters.tasks import TasksManager
-from optimum.intel import OVPipelineForText2Image
-from optimum.intel.openvino import OVModelForCausalLM, OVModelForVisualCausalLM
-from optimum.utils import NormalizedConfigManager, NormalizedTextConfig
 from PIL import Image
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, AutoProcessor, AutoModel, AutoModelForVision2Seq
-import openvino as ov
 
 from whowhatbench import EVALUATOR_REGISTRY
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-TasksManager._SUPPORTED_MODEL_TYPE["stablelm-epoch"] = (
-    TasksManager._SUPPORTED_MODEL_TYPE["llama"]
-)
-NormalizedConfigManager._conf["stablelm-epoch"] = NormalizedTextConfig.with_args(
-    num_layers="num_hidden_layers",
-    num_attention_heads="num_attention_heads",
-)
-
 
 class GenAIModelWrapper:
     """
@@ -81,6 +70,7 @@ def load_text_model(
         model = load_text_genai_pipeline(model_id, device, ov_config)
     else:
         logger.info("Using Optimum API")
+        from optimum.intel.openvino import OVModelForCausalLM
         try:
             model = OVModelForCausalLM.from_pretrained(
                 model_id, trust_remote_code=True, device=device, ov_config=ov_config
@@ -100,11 +90,6 @@ def load_text_model(
     return model
 
 
-TEXT2IMAGE_TASK2CLASS = {
-    "text-to-image": OVPipelineForText2Image,
-}
-
-
 def load_text2image_genai_pipeline(model_dir, device="CPU", ov_config=None):
     try:
         import openvino_genai
@@ -132,7 +117,8 @@ def load_text2image_model(
             model_id, trust_remote_code=True)
     else:
         logger.info("Using Optimum API")
-        TEXT2IMAGEPipeline = TEXT2IMAGE_TASK2CLASS[model_type]
+        from optimum.intel import OVPipelineForText2Image
+        TEXT2IMAGEPipeline = OVPipelineForText2Image
 
         try:
             model = TEXT2IMAGEPipeline.from_pretrained(
@@ -192,6 +178,7 @@ def load_visual_text_model(
         model = load_visual_text_genai_pipeline(model_id, device, ov_config)
     else:
         logger.info("Using Optimum API")
+        from optimum.intel.openvino import OVModelForVisualCausalLM
         try:
             model = OVModelForVisualCausalLM.from_pretrained(
                 model_id, trust_remote_code=True, device=device, ov_config=ov_config