feat: save config as yaml file when generating

Author: ChenZiHong-Gavin

generate.py

@@ -2,6 +2,7 @@
 import json
 import time
 import argparse
+import yaml
 from dotenv import load_dotenv
 
 from graphgen.graphgen import GraphGen
@@ -75,3 +76,5 @@
     graph_gen.judge(re_judge=False)
 
     graph_gen.traverse()
+    with open(os.path.join(sys_path, "cache", "configs", f"graphgen_{unique_id}.yaml"), "w", encoding='utf-8') as f:
+        yaml.dump(traverse_strategy.to_yaml(), f)

graphgen/graphgen.py

@@ -67,7 +67,7 @@ async def async_split_chunks(self, data: Union[List[list], List[dict]], data_typ
             }
             _add_doc_keys = await self.full_docs_storage.filter_keys(list(new_docs.keys()))
             new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys}
-            if not len(new_docs):
+            if len(new_docs) == 0:
                 logger.warning("All docs are already in the storage")
                 return {}
             logger.info(f"[New Docs] inserting {len(new_docs)} docs")
@@ -86,11 +86,12 @@ async def async_split_chunks(self, data: Union[List[list], List[dict]], data_typ
         elif data_type == "chunked":
             assert isinstance(data, list) and isinstance(data[0], list)
             new_docs = {
-                compute_content_hash("".join(chunk['content']), prefix="doc-"): {'content': "".join(chunk['content'])} for doc in data for chunk in doc
+                compute_content_hash("".join(chunk['content']), prefix="doc-"): {'content': "".join(chunk['content'])}
+                for doc in data for chunk in doc
             }
             _add_doc_keys = await self.full_docs_storage.filter_keys(list(new_docs.keys()))
             new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys}
-            if not len(new_docs):
+            if len(new_docs) == 0:
                 logger.warning("All docs are already in the storage")
                 return {}
             logger.info(f"[New Docs] inserting {len(new_docs)} docs")

models/strategy/travserse_strategy.py

@@ -19,3 +19,16 @@ class TraverseStrategy(BaseStrategy):
     edge_sampling: str = "max_loss" # "max_loss" or "min_loss" or "random"
     # 孤立节点的处理策略
     isolated_node_strategy: str = "add" # "add" or "ignore"
+
+    def to_yaml(self):
+        return {
+            "traverse_strategy": {
+                "expand_method": self.expand_method,
+                "bidirectional": self.bidirectional,
+                "max_extra_edges": self.max_extra_edges,
+                "max_tokens": self.max_tokens,
+                "max_depth": self.max_depth,
+                "edge_sampling": self.edge_sampling,
+                "isolated_node_strategy": self.isolated_node_strategy
+            }
+        }