From 7d5d9d671d5ea99c567c2de94b5bf43bfdd6995b Mon Sep 17 00:00:00 2001 From: eaidova Date: Wed, 4 Dec 2024 13:52:18 +0400 Subject: [PATCH] avoid import onnx in modeling_base and stateful --- optimum/exporters/openvino/stateful.py | 2 +- optimum/intel/openvino/modeling_base.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/optimum/exporters/openvino/stateful.py b/optimum/exporters/openvino/stateful.py index 4b4374ab51..90630503a6 100644 --- 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,7 @@ 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 index 3799bf8567..99422f1a54 100644 --- 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,