Add openvino export configs (#568)

* add openvino export configs

* more libs

* more libs

* mixtral and model patcher

* chatglm export

* rework chatglm config

* more testing models

* rework config registration

* add chatglm in tests

* Update tests/openvino/test_modeling.py

* fix style

* gemma

* add test models

* qwen

* fix failed tests

* add comment for gemma

Author: eaidova

optimum/exporters/openvino/__init__.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import optimum.exporters.openvino.model_configs
+
 from .__main__ import main_export
 from .convert import export, export_from_model, export_models, export_pytorch_via_onnx
 from .stateful import ensure_stateful_is_available, patch_stateful

optimum/exporters/openvino/__main__.py

@@ -58,7 +58,7 @@ def main_export(
     local_files_only: bool = False,
     use_auth_token: Optional[Union[bool, str]] = None,
     model_kwargs: Optional[Dict[str, Any]] = None,
-    custom_onnx_configs: Optional[Dict[str, "OnnxConfig"]] = None,
+    custom_export_configs: Optional[Dict[str, "OnnxConfig"]] = None,
     fn_get_submodels: Optional[Callable] = None,
     compression_option: Optional[str] = None,
     compression_ratio: Optional[float] = None,
@@ -112,11 +112,11 @@ def main_export(
             when running `transformers-cli login` (stored in `~/.huggingface`).
         model_kwargs (`Optional[Dict[str, Any]]`, defaults to `None`):
             Experimental usage: keyword arguments to pass to the model during
-            the export. This argument should be used along the `custom_onnx_configs` argument
+            the export. This argument should be used along the `custom_export_configs` argument
             in case, for example, the model inputs/outputs are changed (for example, if
             `model_kwargs={"output_attentions": True}` is passed).
-        custom_onnx_configs (`Optional[Dict[str, OnnxConfig]]`, defaults to `None`):
-            Experimental usage: override the default ONNX config used for the given model. This argument may be useful for advanced users that desire a finer-grained control on the export. An example is available [here](https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model).
+        custom_export_configs (`Optional[Dict[str, OnnxConfig]]`, defaults to `None`):
+            Experimental usage: override the default export config used for the given model. This argument may be useful for advanced users that desire a finer-grained control on the export. An example is available [here](https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model).
         fn_get_submodels (`Optional[Callable]`, defaults to `None`):
             Experimental usage: Override the default submodels that are used at the export. This is
             especially useful when exporting a custom architecture that needs to split the ONNX (e.g. encoder-decoder). If unspecified with custom models, optimum will try to use the default submodels used for the given task, with no guarantee of success.
@@ -134,7 +134,7 @@ def main_export(
     ```python
     >>> from optimum.exporters.openvino import main_export
 
-    >>> main_export("gpt2", output="gpt2_onnx/")
+    >>> main_export("gpt2", output="gpt2_ov/")
     ```
     """
 
@@ -206,14 +206,14 @@ def main_export(
         if model_type not in TasksManager._SUPPORTED_MODEL_TYPE:
             custom_architecture = True
         elif task not in TasksManager.get_supported_tasks_for_model_type(
-            model_type, exporter="onnx", library_name=library_name
+            model_type, exporter="openvino", library_name=library_name
         ):
             if original_task == "auto":
                 autodetected_message = " (auto-detected)"
             else:
                 autodetected_message = ""
             model_tasks = TasksManager.get_supported_tasks_for_model_type(
-                model_type, exporter="onnx", library_name=library_name
+                model_type, exporter="openvino", library_name=library_name
             )
             raise ValueError(
                 f"Asked to export a {model_type} model for the task {task}{autodetected_message}, but the Optimum OpenVINO exporter only supports the tasks {', '.join(model_tasks.keys())} for {model_type}. Please use a supported task. Please open an issue at https://github.com/huggingface/optimum/issues if you would like the task {task} to be supported in the ONNX export for {model_type}."
@@ -288,7 +288,7 @@ class StoreAttr(object):
         not custom_architecture
         and library_name != "diffusers"
         and task + "-with-past"
-        in TasksManager.get_supported_tasks_for_model_type(model_type, exporter="onnx", library_name=library_name)
+        in TasksManager.get_supported_tasks_for_model_type(model_type, exporter="openvino", library_name=library_name)
     ):
         # Make -with-past the default if --task was not explicitely specified
         if original_task == "auto":
@@ -319,7 +319,7 @@ class StoreAttr(object):
         ov_config=ov_config,
         stateful=stateful,
         model_kwargs=model_kwargs,
-        custom_onnx_configs=custom_onnx_configs,
+        custom_export_configs=custom_export_configs,
         fn_get_submodels=fn_get_submodels,
         preprocessors=preprocessors,
         device=device,

optimum/exporters/openvino/convert.py

@@ -32,10 +32,11 @@
 from optimum.exporters.onnx.convert import check_dummy_inputs_are_allowed
 from optimum.exporters.onnx.convert import export_pytorch as export_pytorch_to_onnx
 from optimum.exporters.onnx.convert import export_tensorflow as export_tensorflow_onnx
+from optimum.exporters.utils import _get_submodels_and_export_configs
 from optimum.utils import DEFAULT_DUMMY_SHAPES, is_diffusers_available
 from optimum.utils.save_utils import maybe_save_preprocessors
 
-from ...intel.utils.import_utils import is_nncf_available, is_optimum_version
+from ...intel.utils.import_utils import is_nncf_available
 from .model_patcher import patch_model_with_bettertransformer
 from .stateful import ensure_export_task_support_stateful, ensure_stateful_is_available, patch_stateful
 from .utils import (
@@ -48,13 +49,6 @@
 )
 
 
-if is_optimum_version(">=", "1.16.99"):
-    from optimum.exporters.onnx.utils import _get_submodels_and_onnx_configs
-
-else:
-    from optimum.exporters.onnx.__main__ import _get_submodels_and_onnx_configs
-
-
 UNSUPPORTED_TOKENIZER_CLASSES = (T5Tokenizer, T5TokenizerFast)
 
 
@@ -418,7 +412,7 @@ def ts_patched_forward(*args, **kwargs):
 
 
 def export_models(
-    models_and_onnx_configs: Dict[
+    models_and_export_configs: Dict[
         str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], "OnnxConfig"]
     ],
     output_dir: Path,
@@ -434,7 +428,7 @@ def export_models(
     Export the models to OpenVINO IR format
 
     Args:
-        models_and_onnx_configs (Dict[ str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], "OnnxConfig"]):
+        models_and_export_configs (Dict[ str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], "OnnxConfig"]):
         output_dir (Path): output directory for saving models
         opset (Optional[int], optional, Default to None): ONNX export opset
         output_names (Optional[List[str]], optional, Defaults to None): model output names
@@ -459,20 +453,20 @@ def export_models(
 
     outputs = []
 
-    if output_names is not None and len(output_names) != len(models_and_onnx_configs):
+    if output_names is not None and len(output_names) != len(models_and_export_configs):
         raise ValueError(
-            f"Provided custom names {output_names} for the export of {len(models_and_onnx_configs)} models. Please provide the same number of names as models to export."
+            f"Provided custom names {output_names} for the export of {len(models_and_export_configs)} models. Please provide the same number of names as models to export."
         )
 
-    for i, model_name in enumerate(models_and_onnx_configs.keys()):
-        submodel, sub_onnx_config = models_and_onnx_configs[model_name]
+    for i, model_name in enumerate(models_and_export_configs.keys()):
+        submodel, sub_export_config = models_and_export_configs[model_name]
         output_name = output_names[i] if output_names is not None else Path(model_name + ".xml")
         output_path = output_dir / output_name
         output_path.parent.mkdir(parents=True, exist_ok=True)
         outputs.append(
             export(
                 model=submodel,
-                config=sub_onnx_config,
+                config=sub_export_config,
                 output=output_path,
                 opset=opset,
                 device=device,
@@ -495,7 +489,7 @@ def export_from_model(
     stateful: bool = True,
     opset: Optional[int] = None,
     model_kwargs: Optional[Dict[str, Any]] = None,
-    custom_onnx_configs: Optional[Dict[str, "OnnxConfig"]] = None,
+    custom_export_configs: Optional[Dict[str, "OnnxConfig"]] = None,
     fn_get_submodels: Optional[Callable] = None,
     preprocessors: List = None,
     device: str = "cpu",
@@ -524,14 +518,14 @@ def export_from_model(
             task = TasksManager._infer_task_from_model_or_model_class(model=model)
         except (ValueError, KeyError) as e:
             raise RuntimeError(
-                f"The model task could not be automatically inferred in `onnx_export_from_model`. Please provide the argument `task` with the relevant task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
+                f"The model task could not be automatically inferred in `export_from_model`. Please provide the argument `task` with the relevant task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
             )
 
         if (
             not custom_architecture
             and library_name != "diffusers"
             and task + "-with-past"
-            in TasksManager.get_supported_tasks_for_model_type(model_type, "onnx", library_name=library_name)
+            in TasksManager.get_supported_tasks_for_model_type(model_type, "openvino", library_name=library_name)
         ):
             # -with-past is the default.
             task = task + "-with-past"
@@ -541,9 +535,9 @@ def export_from_model(
     stateful = stateful and ensure_export_task_support_stateful(task)
 
     # TODO: support onnx_config.py in the model repo
-    if custom_architecture and custom_onnx_configs is None:
+    if custom_architecture and custom_export_configs is None:
         raise ValueError(
-            f"Trying to export a {model_type} model, that is a custom or unsupported architecture, but no custom onnx configuration was passed as `custom_onnx_configs`. Please refer to https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model#custom-export-of-transformers-models for an example on how to export custom models. Please open an issue at https://github.com/huggingface/optimum/issues if you would like the model type {model_type} to be supported natively in the ONNX export."
+            f"Trying to export a {model_type} model, that is a custom or unsupported architecture, but no custom export configuration was passed as `custom_export_configs`. Please refer to https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model#custom-export-of-transformers-models for an example on how to export custom models. Please open an issue at https://github.com/huggingface/optimum/issues if you would like the model type {model_type} to be supported natively in the ONNX export."
         )
 
     if task.startswith("text-generation") and model.config.is_encoder_decoder:
@@ -569,18 +563,19 @@ def export_from_model(
             kwargs_shapes[input_name] if input_name in kwargs_shapes else DEFAULT_DUMMY_SHAPES[input_name]
         )
 
-    onnx_config, models_and_onnx_configs = _get_submodels_and_onnx_configs(
+    export_config, models_and_export_configs = _get_submodels_and_export_configs(
         model=model,
         task=task,
         monolith=False,
-        custom_onnx_configs=custom_onnx_configs if custom_onnx_configs is not None else {},
+        custom_export_configs=custom_export_configs if custom_export_configs is not None else {},
         custom_architecture=custom_architecture,
         fn_get_submodels=fn_get_submodels,
         preprocessors=preprocessors,
         library_name=library_name,
         model_kwargs=model_kwargs,
         _variant="default",
         legacy=False,
+        exporter="openvino",
     )
 
     if ov_config is None:
@@ -612,18 +607,18 @@ def export_from_model(
         model_name_or_path = model.config._name_or_path
         maybe_save_preprocessors(model_name_or_path, output)
 
-        files_subpaths = ["openvino_" + model_name + ".xml" for model_name in models_and_onnx_configs.keys()]
+        files_subpaths = ["openvino_" + model_name + ".xml" for model_name in models_and_export_configs.keys()]
 
     else:
         # save the subcomponent configuration
-        for model_name in models_and_onnx_configs:
-            subcomponent = models_and_onnx_configs[model_name][0]
+        for model_name in models_and_export_configs:
+            subcomponent = models_and_export_configs[model_name][0]
             if hasattr(subcomponent, "save_config"):
                 subcomponent.save_config(output / model_name)
             elif hasattr(subcomponent, "config") and hasattr(subcomponent.config, "save_pretrained"):
                 subcomponent.config.save_pretrained(output / model_name)
 
-        files_subpaths = [os.path.join(name_dir, OV_XML_FILE_NAME) for name_dir in models_and_onnx_configs]
+        files_subpaths = [os.path.join(name_dir, OV_XML_FILE_NAME) for name_dir in models_and_export_configs]
 
         # Saving the additional components needed to perform inference.
         model.scheduler.save_pretrained(output.joinpath("scheduler"))
@@ -643,7 +638,7 @@ def export_from_model(
         model.save_config(output)
 
     export_models(
-        models_and_onnx_configs=models_and_onnx_configs,
+        models_and_export_configs=models_and_export_configs,
         output_dir=output,
         output_names=files_subpaths,
         input_shapes=input_shapes,