fix ModelStorage save/load

Author: NotBioWaste905

dialogue2graph/pipelines/model_storage.py

@@ -98,8 +98,8 @@ def load(self, path: Path):
         try:
             with open(path, "r") as f:
                 loaded_storage = yaml.safe_load(f)
-                for key, config in loaded_storage.items():
-                    self.add(key=key, config=config, model_type=config.pop("model_type"))
+                for key in loaded_storage:
+                    self.add(key=key, config=loaded_storage[key].get("config"), model_type=loaded_storage[key].get("model_type"))
                     logger.debug(f"Loaded model configuration for '{key}'")
             logger.info(f"Successfully loaded {len(loaded_storage)} models from {path}")
         except Exception as e:
@@ -148,7 +148,11 @@ def save(self, path: str):
         logger.debug(f"Attempting to save model storage to {path}")
         try:
             with open(path, "w") as f:
-                storage_dump = {k: v.config for k, v in self.storage.items()}
+                storage_dump = {}
+                for model_key in self.storage:
+                    storage_dump[model_key] = {}
+                    storage_dump[model_key]["config"] = self.storage[model_key].config
+                    storage_dump[model_key]["model_type"] = self.storage[model_key].model_type
                 yaml.dump(storage_dump, f)
             logger.info(f"Saved {len(self.storage)} models to {path}")
         except Exception as e:

examples/pipeline_usage/models_config.yml

@@ -1,12 +1,20 @@
 d2g_llm_filling_llm:v1:
-  name: o3-mini
-  temperature: 1
+  config:
+    name: o3-mini
+    temperature: 1
+  model_type: llm
 d2g_llm_grouping_llm:v1:
-  name: gpt-4o-latest
-  temperature: 0
+  config:
+    name: gpt-4o-latest
+    temperature: 0
+  model_type: llm
 d2g_llm_sim_model:v1:
-  model_kwargs:
-    device: cpu
-  model_name: cointegrated/LaBSE-en-ru
+  config:
+    model_kwargs:
+      device: cpu
+    model_name: cointegrated/LaBSE-en-ru
+  model_type: emb
 my_model:
-  name: gpt-3.5-turbo
+  config:
+    name: gpt-3.5-turbo
+  model_type: llm