fix(graphgen): add unique_id for log

Author: ChenZiHong-Gavin

baselines/Genie/genie.py

@@ -1,11 +1,11 @@
 # https://arxiv.org/pdf/2401.14367
 import os
 import json
-from dotenv import load_dotenv
 import argparse
 import asyncio
 
 from dataclasses import dataclass
+from dotenv import load_dotenv
 from models import OpenAIModel
 from typing import List
 from utils import create_event_loop, compute_content_hash
@@ -58,7 +58,7 @@ def generate(self, docs: List[List[dict]]) -> List[dict]:
         loop = create_event_loop()
         return loop.run_until_complete(self.async_generate(docs))
 
-    async def async_generate(self, docs: List[List[dict]]) -> List[dict]:
+    async def async_generate(self, docs: List[List[dict]]) -> dict:
         final_results = {}
         semaphore = asyncio.Semaphore(self.max_concurrent)
 

generate.py

@@ -1,5 +1,6 @@
 import os
 import json
+import time
 import argparse
 from dotenv import load_dotenv
 
@@ -8,7 +9,8 @@
 from utils import set_logger
 
 sys_path = os.path.abspath(os.path.dirname(__file__))
-set_logger(os.path.join(sys_path, "cache", "logs", "graphgen.log"), if_stream=False)
+unique_id = int(time.time())
+set_logger(os.path.join(sys_path, "cache", "logs", f"graphgen_{unique_id}.log"), if_stream=False)
 
 load_dotenv()
 
@@ -58,6 +60,7 @@
     traverse_strategy = TraverseStrategy()
 
     graph_gen = GraphGen(
+        unique_id=unique_id,
         teacher_llm_client=teacher_llm_client,
         student_llm_client=student_llm_client,
         if_web_search=args.web_search,

graphgen/graphgen.py

@@ -1,15 +1,14 @@
 # https://github.com/HKUDS/LightRAG
 
 import os
-import time
 import asyncio
+import time
+from typing import List, cast, Union
+from dataclasses import dataclass
 from tqdm.asyncio import tqdm as tqdm_async
 
 from .operators import *
 from models import Chunk, JsonKVStorage, OpenAIModel, NetworkXStorage, WikiSearch, Tokenizer, TraverseStrategy
-from typing import List, cast, Union
-
-from dataclasses import dataclass
 from utils import create_event_loop, logger, compute_content_hash
 from models.storage.base_storage import StorageNameSpace
 
@@ -19,6 +18,7 @@
 
 @dataclass
 class GraphGen:
+    unique_id: int = int(time.time())
     working_dir: str = os.path.join(sys_path, "cache")
     full_docs_storage: JsonKVStorage = JsonKVStorage(
         working_dir, namespace="full_docs"
@@ -33,7 +33,7 @@ class GraphGen:
         working_dir, namespace="graph"
     )
     qa_storage: JsonKVStorage = JsonKVStorage(
-        os.path.join(working_dir, "data", "graphgen"), namespace=f"qa-{int(time.time())}"
+        os.path.join(working_dir, "data", "graphgen"), namespace=f"qa-{unique_id}"
     )
 
     # text chunking

models/storage/networkx_storage.py

@@ -1,17 +1,17 @@
 import os
 import html
 import networkx as nx
-import numpy as np
 
-from typing import Any, Union, cast
+from typing import Any, Union, cast, Optional
 from dataclasses import dataclass
+
 from .base_storage import BaseGraphStorage
 from utils import logger
 
 @dataclass
 class NetworkXStorage(BaseGraphStorage):
     @staticmethod
-    def load_nx_graph(file_name) -> nx.Graph:
+    def load_nx_graph(file_name) -> Optional[nx.Graph]:
         if os.path.exists(file_name):
             return nx.read_graphml(file_name)
         return None
@@ -85,9 +85,6 @@ def __post_init__(self):
                 f"Loaded graph from {self._graphml_xml_file} with {preloaded_graph.number_of_nodes()} nodes, {preloaded_graph.number_of_edges()} edges"
             )
         self._graph = preloaded_graph or nx.Graph()
-        # self._node_embed_algorithms = {
-        #     "node2vec": self._node2vec_embed,
-        # }
 
     async def index_done_callback(self):
         NetworkXStorage.write_nx_graph(self._graph, self._graphml_xml_file)
@@ -154,8 +151,3 @@ async def delete_node(self, node_id: str):
             logger.info(f"Node {node_id} deleted from the graph.")
         else:
             logger.warning(f"Node {node_id} not found in the graph for deletion.")
-
-    async def embed_nodes(self, algorithm: str) -> tuple[np.ndarray, list[str]]:
-        if algorithm not in self._node_embed_algorithms:
-            raise ValueError(f"Node embedding algorithm {algorithm} not supported")
-        return await self._node_embed_algorithms[algorithm]()

requirements.txt

@@ -14,3 +14,5 @@ jieba
 torch
 plotly
 pandas
+gradio
+kaleido