avoid import onnx in modeling_base and stateful

huggingface · Dec 4, 2024 · 97e9141 · 97e9141
1 parent e26f516
commit 97e9141
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/optimum/exporters/openvino/stateful.py b/optimum/exporters/openvino/stateful.py
@@ -20,7 +20,6 @@
 
 import openvino as ov
 from openvino.runtime import opset13
-from optimum.exporters import TasksManager
 from optimum.intel.utils.import_utils import _openvino_version, is_openvino_version, is_transformers_version
 
 from .utils import MULTI_MODAL_TEXT_GENERATION_MODELS
@@ -192,6 +191,8 @@ def ensure_stateful_is_available(warn=True):
 
 
 def ensure_export_task_support_stateful(task: str):
+    from optimum.exporters import TasksManager
+
     task = TasksManager.map_from_synonym(task)
     return task in ["text-generation-with-past"]
 

diff --git a/optimum/intel/openvino/modeling_base.py b/optimum/intel/openvino/modeling_base.py
@@ -30,7 +30,7 @@
 from transformers.generation import GenerationMixin
 from transformers.utils import is_offline_mode
 
-from optimum.exporters.onnx import OnnxConfig
+from optimum.exporters.base import ExportConfig
 from optimum.modeling_base import FROM_PRETRAINED_START_DOCSTRING, OptimizedModel
 
 from ...exporters.openvino import export, main_export
@@ -623,7 +623,7 @@ def _to_load(
         cls,
         model,
         config: PretrainedConfig,
-        onnx_config: OnnxConfig,
+        onnx_config: ExportConfig,
         token: Optional[Union[bool, str]] = None,
         revision: Optional[str] = None,
         force_download: bool = False,