Removed saving to temporary directory; added core property handling for OVModelForCausalLM

Author: nikita-savelyevv

optimum/commands/export/openvino.py

@@ -15,7 +15,6 @@
 
 import logging
 import sys
-import tempfile
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
@@ -357,17 +356,12 @@ def run(self):
             if quantize_after_export:
                 from optimum.intel import OVModelForCausalLM, OVQuantizer
 
+                # TODO: remove disabling mmap once OV is updated to 2024.3
                 model = OVModelForCausalLM.from_pretrained(
-                    self.args.output, trust_remote_code=self.args.trust_remote_code
+                    self.args.output, trust_remote_code=self.args.trust_remote_code, ov_config={"ENABLE_MMAP": "NO"}
                 )
                 quantizer = OVQuantizer(model)
                 quantization_config.tokenizer = quantization_config.tokenizer or str(self.args.output)
-                # TODO: set save_directory=self.args.output once OV is updated to 2024.3
-                quantizer.quantize(ov_config=OVConfig(quantization_config=quantization_config))
-                with tempfile.TemporaryDirectory() as temp_dir:
-                    import shutil
-
-                    model.save_pretrained(temp_dir)
-                    ov_config.save_pretrained(self.args.output)
-                    shutil.copy(f"{temp_dir}/openvino_model.xml", f"{self.args.output}/openvino_model.xml")
-                    shutil.copy(f"{temp_dir}/openvino_model.bin", f"{self.args.output}/openvino_model.bin")
+                quantizer.quantize(
+                    ov_config=OVConfig(quantization_config=quantization_config), save_directory=self.args.output
+                )

optimum/intel/openvino/modeling_base.py

@@ -103,6 +103,7 @@ def __init__(
     def load_model(
         file_name: Union[str, Path],
         quantization_config: Union[OVWeightQuantizationConfig, Dict] = None,
+        ov_core_properties: Optional[Dict] = None,
     ):
         """
         Loads the model.
@@ -112,6 +113,8 @@ def load_model(
                 The path of the model ONNX or XML file.
             quantization_config (`OVWeightQuantizationConfig` or `Dict`, *optional*):
                 Quantization config to apply after model is loaded.
+            ov_core_properties (`Dict`, *optional*):
+                OpenVINO core properties to set before model loading.
         """
 
         def fix_op_names_duplicates(model: openvino.runtime.Model):
@@ -128,6 +131,8 @@ def fix_op_names_duplicates(model: openvino.runtime.Model):
 
         if isinstance(file_name, str):
             file_name = Path(file_name)
+        if ov_core_properties:
+            core.set_property(ov_core_properties)
         model = core.read_model(file_name) if not file_name.suffix == ".onnx" else convert_model(file_name)
         if file_name.suffix == ".onnx":
             model = fix_op_names_duplicates(model)  # should be called during model conversion to IR

optimum/intel/openvino/modeling_decoder.py

@@ -748,9 +748,16 @@ def _from_pretrained(
 
         load_in_4bit = quantization_config.bits == 4 if quantization_config else False
 
+        ov_config = kwargs.get("ov_config", None)
+        ov_core_properties = {}
+        if ov_config and "ENABLE_MMAP" in ov_config:
+            ov_core_properties["ENABLE_MMAP"] = ov_config["ENABLE_MMAP"]
+            del ov_config["ENABLE_MMAP"]
+
         model = cls.load_model(
             model_cache_path,
             quantization_config=None if load_in_4bit else quantization_config,
+            ov_core_properties=ov_core_properties,
         )
 
         model_type = config.model_type.replace("_", "-")