Incorporated changes suggested in comments

Signed-off-by: Dhiraj Kumar Sah <[email protected]>

Author: quic-dhirajku

QEfficient/base/modeling_qeff.py

@@ -5,7 +5,7 @@
 #
 # ----------------------------------------------------------------------------
 
-# import hashlib
+import copy
 import inspect
 import json
 import logging
@@ -52,10 +52,9 @@ def __init__(self, model: torch.nn.Module, **kwargs) -> None:
 
         # Store Model parameters to Calculate Hash for caching
         self.model_params = {}
-        self.model_params.update(kwargs)
+        self.model_params = copy.deepcopy(kwargs)
         self.model_params["config"] = self.model.config.to_diff_dict()
         self.model_params["_transform_names"] = self._transform_names()
-        self.compile_params = {}
 
         if hasattr(self.model.config, "architectures"):
             self.model_architecture = self.model.config.architectures[0]
@@ -142,13 +141,15 @@ def _export(
             :onnx_transform_kwargs (dict): Additional arguments to be passed to `Transform.apply` for this class.
             :export_dir (str): Specify the export directory. The export_dir will be suffixed with a hash corresponding to current model.
         """
-        self.model_params["output_names"] = output_names
-        self.model_params["dynamic_axes"] = dynamic_axes
+        export_params = {}
+        export_params["output_names"] = output_names
+        export_params["dynamic_axes"] = dynamic_axes
+
+        self.model_params["export_params"] = export_params
+
+        self.model_params.update(export_kwargs) if export_kwargs is not None else None
+        self.model_params.update(onnx_transform_kwargs) if export_kwargs is not None else None
 
-        if export_kwargs is not None:
-            self.model_params.update(export_kwargs)
-        if onnx_transform_kwargs is not None:
-            self.model_params.update(onnx_transform_kwargs)
         export_dir = Path(export_dir or (QEFF_HOME / self.model_architecture / self.model_name))
 
         export_hash = hash_dict_params(self.model_params)
@@ -163,17 +164,6 @@ def _export(
         tmp_onnx_path = tmp_onnx_dir / f"{self.model_name}.onnx"
         tmp_onnx_dir.mkdir(parents=True, exist_ok=True)
 
-        model_params_json = export_dir / "model_params.json"
-        with open(model_params_json, "w") as fp:
-            json.dump(
-                {
-                    "model_params": [
-                        {k: make_serializable(self.model_params[k]) for k in sorted(self.model_params.keys())}
-                    ]
-                },
-                fp,
-                indent=4,
-            )
         # Create input_names from example_inputs
 
         input_names = []
@@ -231,6 +221,20 @@ def _export(
             onnx.save(model, onnx_path)
             logger.info("Transformed onnx saved")
 
+            # Dumping model paramters in a JSON file after successful ONNX export
+            model_params_json = export_dir / "model_params.json"
+            with open(model_params_json, "w") as fp:
+                json.dump(
+                    {
+                        "model_params": {
+                            k: make_serializable(self.model_params[k]) for k in sorted(self.model_params.keys())
+                        }
+                    },
+                    fp,
+                    indent=4,
+                )
+            logger.info("Parameters used for export hash dumped in a JSON file successfully")
+
         except Exception as e:
             logger.error(f"ONNX export (or) ONNXTransforms failed: {e}")
 
@@ -277,6 +281,8 @@ def _compile(
         if onnx_path is None and self.onnx_path is None:
             self.export()
 
+        self.compile_params = {}
+
         onnx_path = Path(onnx_path or self.onnx_path)
         compile_dir = Path(compile_dir or onnx_path.parent)
         qpc_path = compile_dir / "qpc"
@@ -339,18 +345,6 @@ def _compile(
             # Probably compilation failure last time, delete directory to start over
             shutil.rmtree(qpc_path)
 
-        compile_params_json = compile_dir / "compile_params.json"
-        with open(compile_params_json, "w") as fp:
-            json.dump(
-                {
-                    "compile_params": [
-                        {k: make_serializable(self.compile_params[k]) for k in sorted(self.compile_params.keys())}
-                    ]
-                },
-                fp,
-                indent=4,
-            )
-
         # Write specializations.json file
         if specializations is not None:
             specializations_json = compile_dir / "specializations.json"
@@ -394,6 +388,19 @@ def _compile(
         logger.info(f"Running compiler: {' '.join(command)}")
         try:
             subprocess.run(command, capture_output=True, check=True)
+
+            # Dumping compile paramters in a JSON file after successful ONNX export
+            compile_params_json = compile_dir / "compile_params.json"
+            with open(compile_params_json, "w") as fp:
+                json.dump(
+                    {
+                        "compile_params": {
+                            k: make_serializable(self.compile_params[k]) for k in sorted(self.compile_params.keys())
+                        }
+                    },
+                    fp,
+                    indent=4,
+                )
         except subprocess.CalledProcessError as e:
             raise RuntimeError(
                 "\n".join(

QEfficient/transformers/models/modeling_auto.py

@@ -171,8 +171,6 @@ def __init__(self, model: nn.Module, pooling=None, **kwargs):
         self.model.base_model.config.use_cache = True
         self.model_params["qeff_class"] = self.__class__.__name__
 
-        # self.pretrained_model_name_or_path = kwargs.get("pretrained_model_name_or_path", None)
-
     @classmethod
     @with_replaced_quantizers
     def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, *args, **kwargs):
@@ -913,7 +911,6 @@ def __init__(
             self.model.config.vision_config.use_flash_attn = "false"
         else:
             self.model.config.text_config.use_cache = True
-        self.pretrained_model_name_or_path = kwargs.get("pretrained_model_name_or_path", None)
 
     @classmethod
     def from_pretrained(