From 4869104f4454ab610b5655f0f20d4a28f9eca38d Mon Sep 17 00:00:00 2001
From: Helena Kloosterman <helena.kloosterman@intel.com>
Date: Mon, 29 Apr 2024 17:48:44 +0200
Subject: [PATCH] Add version info to OpenVINO models (#690)

* Add version info to OpenVINO models

* Add more dependencies
---
 optimum/exporters/openvino/convert.py | 44 ++++++++++++++++++++++++++-
 optimum/intel/utils/import_utils.py   |  1 +
 tests/openvino/test_export.py         |  4 +++
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py
index 55e3318017..6c86c2c2df 100644
--- a/optimum/exporters/openvino/convert.py
+++ b/optimum/exporters/openvino/convert.py
@@ -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
@@ -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
 
@@ -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)
 
 
@@ -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
diff --git a/optimum/intel/utils/import_utils.py b/optimum/intel/utils/import_utils.py
index 06848dbaf2..ac6306923d 100644
--- a/optimum/intel/utils/import_utils.py
+++ b/optimum/intel/utils/import_utils.py
@@ -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"
diff --git a/tests/openvino/test_export.py b/tests/openvino/test_export.py
index 9d1daaab63..851f8355f5 100644
--- a/tests/openvino/test_export.py
+++ b/tests/openvino/test_export.py
@@ -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
 
 
@@ -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):