Skip to content

Commit

Permalink
Add version info to OpenVINO models (#690)
Browse files Browse the repository at this point in the history
* Add version info to OpenVINO models

* Add more dependencies
  • Loading branch information
helena-intel authored Apr 29, 2024
1 parent 805335c commit 4869104
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 1 deletion.
44 changes: 43 additions & 1 deletion optimum/exporters/openvino/convert.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,10 @@
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union

import onnx
from transformers.utils import is_tf_available, is_torch_available

from openvino.runtime import PartialShape, save_model
from openvino.runtime import Model, PartialShape, save_model
from openvino.runtime.exceptions import OVTypeError
from openvino.runtime.utils.types import get_element_type
from openvino.tools.ovc import convert_model
Expand All @@ -32,6 +33,14 @@
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.intel.utils.import_utils import (
_nncf_version,
_optimum_intel_version,
_optimum_version,
_timm_version,
_torch_version,
_transformers_version,
)
from optimum.utils import DEFAULT_DUMMY_SHAPES, is_diffusers_available
from optimum.utils.save_utils import maybe_save_preprocessors

Expand Down Expand Up @@ -81,6 +90,8 @@ def _save_model(model, path: str, ov_config: Optional["OVConfig"] = None):

compress_to_fp16 = ov_config.dtype == "fp16"

library_name = TasksManager.infer_library_from_model(Path(path).parent)
model = _add_version_info_to_model(model, library_name)
save_model(model, path, compress_to_fp16)


Expand Down Expand Up @@ -689,3 +700,34 @@ def export_tokenizer(

for model, file_name in zip(converted, (OV_TOKENIZER_NAME, OV_DETOKENIZER_NAME)):
save_model(model, output / file_name.format(suffix))


def _add_version_info_to_model(model: Model, library_name: Optional[str] = None):
"""
Add dependency versions to OpenVINO model
"""
try:
model.set_rt_info(_transformers_version, ["optimum", "transformers_version"])
model.set_rt_info(_torch_version, ["optimum", "pytorch_version"])
model.set_rt_info(_optimum_intel_version, ["optimum", "optimum_intel_version"])
model.set_rt_info(_optimum_version, ["optimum", "optimum_version"])

if any("token_embeddings" in output.get_names() for output in model.outputs):
import sentence_transformers

model.set_rt_info(sentence_transformers.__version__, ["optimum", "sentence_transformers_version"])
if library_name == "diffusers":
model.set_rt_info(_optimum_version, ["optimum", "diffusers_version"])
elif library_name == "timm":
model.set_rt_info(_timm_version, ["optimum", "timm_version"])
rt_info = model.get_rt_info()
if "nncf" in rt_info:
model.set_rt_info(_nncf_version, ["optimum", "nncf_version"])
input_model = rt_info["conversion_parameters"].get("input_model", None)
if input_model is not None and "onnx" in input_model.value:
model.set_rt_info(onnx.__version__, ["optimum", "onnx_version"])

except Exception:
pass

return model
1 change: 1 addition & 0 deletions optimum/intel/utils/import_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt}

_optimum_version = importlib_metadata.version("optimum")
_optimum_intel_version = importlib_metadata.version("optimum-intel")

_transformers_available = importlib.util.find_spec("transformers") is not None
_transformers_version = "N/A"
Expand Down
4 changes: 4 additions & 0 deletions tests/openvino/test_export.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@
OVStableDiffusionXLPipeline,
)
from optimum.intel.openvino.modeling_base import OVBaseModel
from optimum.intel.utils.import_utils import _transformers_version
from optimum.utils.save_utils import maybe_load_preprocessors


Expand Down Expand Up @@ -113,6 +114,9 @@ def _openvino_export(

if task == "text-generation":
self.assertEqual(ov_model.stateful, stateful and use_cache)
self.assertEqual(
ov_model.model.get_rt_info()["optimum"]["transformers_version"], _transformers_version
)

@parameterized.expand(SUPPORTED_ARCHITECTURES)
def test_export(self, model_type: str):
Expand Down

0 comments on commit 4869104

Please sign in to comment.