resolve dictionary as input

Author: eaidova

optimum/exporters/openvino/__main__.py

@@ -148,7 +148,6 @@ def main_export(
         else:
             possible_synonyms = ""
         logger.info(f"Automatic task detection to {task}{possible_synonyms}.")
-
     onnx_config, models_and_onnx_configs = optimum_main._get_submodels_and_onnx_configs(
         model=model,
         task=task,

optimum/exporters/openvino/convert.py

@@ -21,7 +21,7 @@
 
 from transformers.utils import is_tf_available, is_torch_available
 
-from openvino.runtime import PartialShape, serialize
+from openvino.runtime import PartialShape, save_model
 from openvino.runtime.utils.types import get_element_type
 from openvino.tools.ovc import convert_model
 from optimum.exporters.onnx.base import OnnxConfig
@@ -151,8 +151,6 @@ def export_pytorch(
 
     with torch.no_grad():
         model.config.return_dict = True
-        custom_patcher = type(config).patch_model_for_export != OnnxConfig.patch_model_for_export
-        model.config.torchscript = not custom_patcher
         model.eval()
 
         # Check if we need to override certain configuration item
@@ -182,24 +180,30 @@ def export_pytorch(
         else:
             sig = inspect.signature(model.call)
 
-        dummy_inputs = remove_none_from_dummy_inputs(dummy_inputs)
+        dummy_inputs, dict_inputs = remove_none_from_dummy_inputs(dummy_inputs)
         input_info = get_input_shapes(dummy_inputs, inputs)
         try:
-            if custom_patcher:
-                patcher = config.patch_model_for_export(model, model_kwargs=model_kwargs)
-                patched_forward = patcher.patched_forward
-
-                @functools.wraps(patched_forward)
-                def ts_patched_forward(*args, **kwargs):
-                    outputs = patched_forward(*args, **kwargs)
-                    return tuple(outputs.values())
-
-                patcher.patched_forward = ts_patched_forward
-                with patcher:
-                    ov_model = convert_model(model, example_input=dummy_inputs, input=input_info)
-            else:
+            patcher = config.patch_model_for_export(model, model_kwargs=model_kwargs)
+            patched_forward = patcher.patched_forward
+
+            @functools.wraps(patched_forward)
+            def ts_patched_forward(*args, **kwargs):
+                for i in range(len(dict_inputs)):
+                    input_name = dict_inputs[i][0]
+                    keys = dict_inputs[i][1]
+                    tuple_input = kwargs[input_name]
+                    input_dict = dict(zip(keys, tuple_input))
+                    kwargs[input_name] = input_dict
+                outputs = patched_forward(*args, **kwargs)
+                return tuple(outputs.values())
+
+            patcher.patched_forward = ts_patched_forward
+            with patcher:
                 ov_model = convert_model(model, example_input=dummy_inputs, input=input_info)
         except Exception:
+            orig_torch_onnx_export = torch.onnx.export
+
+            torch.onnx.export = functools.partial(orig_torch_onnx_export, do_constant_folding=True)
             model.config.torchscript = False
             model.config.return_dict = True
             onnx_output = (
@@ -210,13 +214,19 @@ def ts_patched_forward(*args, **kwargs):
             input_names, output_names = export_pytorch_to_onnx(
                 model, config, opset, onnx_output, device, input_shapes, model_kwargs
             )
+            torch.onnx.export = orig_torch_onnx_export
             ov_model = convert_model(str(onnx_output))
-            serialize(ov_model, output.parent / OV_XML_FILE_NAME if output.suffix != ".xml" else output)
+            save_model(
+                ov_model,
+                output.parent / OV_XML_FILE_NAME if output.suffix != ".xml" else output,
+                compress_to_fp16=False,
+            )
             return input_names, output_names, True
         clear_class_registry()
         ordered_dummy_inputs = {param: dummy_inputs[param] for param in sig.parameters if param in dummy_inputs}
         ordered_input_names = list(inputs)
         flatten_inputs = flattenize_inputs(ordered_dummy_inputs.values())
+        ov_model.validate_nodes_and_infer_types()
         for idx, out_tensor in enumerate(ov_model.outputs):
             if idx < len(output_names):
                 out_tensor.get_tensor().set_names({output_names[idx]})
@@ -233,7 +243,9 @@ def ts_patched_forward(*args, **kwargs):
             inp_tensor.get_node().set_partial_shape(static_shape)
             inp_tensor.get_node().set_element_type(get_element_type(inp_data.cpu().numpy().dtype))
         ov_model.validate_nodes_and_infer_types()
-        serialize(ov_model, output.parent / OV_XML_FILE_NAME if output.suffix != ".xml" else output)
+        save_model(
+            ov_model, output.parent / OV_XML_FILE_NAME if output.suffix != ".xml" else output, compress_to_fp16=False
+        )
         del model
         gc.collect()
     return input_names, output_names, False

optimum/exporters/openvino/utils.py

@@ -46,18 +46,19 @@ def remove_none_from_list_tuple(item):
         return type(item)(new_item)
 
     upd_dummy = {}
+    dict_dummy = []
     for k, v in dummy_inputs.items():
         if v is None:
             continue
         if isinstance(v, dict):
-            for kk, vv in v.items():
-                upd_dummy[kk] = vv
+            dict_dummy.append((k, list(v.keys())))
+            upd_dummy[k] = remove_none_from_list_tuple(tuple(v.values()))
             continue
         if isinstance(v, (tuple, list)):
             upd_dummy[k] = remove_none_from_list_tuple(v)
             continue
         upd_dummy[k] = v
-    return upd_dummy
+    return upd_dummy, dict_dummy
 
 
 def get_input_shapes(dummy_inputs, inputs):

optimum/intel/openvino/modeling_base.py

@@ -150,7 +150,7 @@ def _save_pretrained(self, save_directory: Union[str, Path]):
                 The directory where to save the model files.
         """
         dst_path = os.path.join(save_directory, OV_XML_FILE_NAME)
-        openvino.runtime.serialize(self.model, dst_path)
+        openvino.save_model(self.model, dst_path, compress_to_fp16=False)
 
     @classmethod
     def _from_pretrained(
@@ -201,7 +201,7 @@ def _from_pretrained(
             model_save_dir = model_id
         # Download the model from the hub
         else:
-            model_file_names = [file_name]
+            model_file_names = [file_name] if from_onnx else []
             # If not ONNX then OpenVINO IR
             if not from_onnx:
                 model_file_names.append(file_name.replace(".xml", ".bin"))

optimum/intel/openvino/modeling_base_seq2seq.py

@@ -105,7 +105,7 @@ def _save_pretrained(self, save_directory: Union[str, Path]):
 
         for src_file, dst_file_name in zip(src_files, dst_file_names):
             dst_path = os.path.join(save_directory, dst_file_name)
-            openvino.runtime.serialize(src_file, dst_path)
+            openvino.save_model(src_file, dst_path, compress_to_fp16=False)
 
     @classmethod
     def _from_pretrained(

optimum/intel/openvino/modeling_decoder.py

@@ -191,7 +191,7 @@ def _save_pretrained(self, save_directory: Union[str, Path]):
         """
         model_to_save = self.model if self._pkv_precision == Type.f32 else self._original_model
         dst_path = os.path.join(save_directory, OV_XML_FILE_NAME)
-        openvino.runtime.serialize(model_to_save, dst_path)
+        openvino.save_model(model_to_save, dst_path, compress_to_fp16=False)
 
     @classmethod
     def _from_transformers(

optimum/intel/openvino/modeling_diffusion.py

@@ -155,7 +155,7 @@ def _save_pretrained(self, save_directory: Union[str, Path]):
             if ov_model is not None:
                 dst_path = save_directory / dst_path / OV_XML_FILE_NAME
                 dst_path.parent.mkdir(parents=True, exist_ok=True)
-                openvino.runtime.serialize(ov_model.model, dst_path)
+                openvino.save_model(ov_model.model, dst_path, compress_to_fp16=False)
                 model_dir = ov_model.config.get("_name_or_path", None) or ov_model._model_dir / ov_model._model_name
                 config_path = Path(model_dir) / ov_model.CONFIG_NAME
                 if config_path.is_file():

optimum/intel/openvino/quantization.py

@@ -414,7 +414,7 @@ def _quantize_torchmodel(
     @staticmethod
     def _save_pretrained(model: openvino.runtime.Model, output_path: str):
         compress_quantize_weights_transformation(model)
-        openvino.runtime.serialize(model, output_path)
+        openvino.save_model(model, output_path, compress_to_fp16=False)
 
     def _set_task(self):
         if self.task is None:

optimum/intel/openvino/trainer.py

@@ -772,7 +772,7 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None):
                     compress_quantize_weights_transformation(ov_model)
 
             # Serialize IR xml and bin
-            save_model(ov_model, output_path)
+            save_model(ov_model, output_path, compress_to_fp16=False)
 
     def _get_compression_controller_by_cls(
         self, controller_cls: Type[PTCompressionAlgorithmController]