add weight_map

Change-Id: I82fe09fb33ef5fc48139874ec9fb9aebd178a459
Signed-off-by: Yi Liu <yiliu4@habana.ai>

Author: Yi4Liu

examples/ds/post_process.py

@@ -53,15 +53,13 @@ def update_config(model_path, qmodel_path):
     import json
     import os
 
-    # open config
     with open(os.path.join(model_path, "config.json"), "r") as f:
         config = json.load(f)
         config["quantization_config"] = quantization_config
-        # save new config to qmodel_path
         logger.info(f"Updated config: {config}")
-        logger.debug(f"Saving config to {qmodel_path}")
-
-        with open(os.path.join(qmodel_path, "config.json"), "w") as f:
+        config_filepath = os.path.join(qmodel_path, "config.json")
+        logger.debug(f"Saving config to {config_filepath}")
+        with open(config_filepath, "w") as f:
             json.dump(config, f, indent=4)
 
 

examples/ds/quant.py

@@ -112,8 +112,12 @@ def quant_model_weight_with_low_cpu_usage(model_path, qmodel_path):
     # Dump tensor mapping into json file
     model_state_dict_mapping_file_path = os.path.join(qmodel_path, MODEL_STATE_DICT_MAPPING_FILENAME)
     logger.info(f"Saving tensor mapping to {model_state_dict_mapping_file_path}")
+    state_dict_mapping = {
+        "metadata":{},
+        "weight_map": qtensor_mappping,
+    }
     with open(model_state_dict_mapping_file_path, "w") as f:
-        json.dump(qtensor_mappping, f, indent=4)
+        json.dump(state_dict_mapping, f, indent=4)
 
 
 def _import_oh():