fix: implement generate method

Author: ChenZiHong-Gavin

graphgen/generate.py

@@ -53,20 +53,20 @@ def main():
     mode = config["generate"]["mode"]
     unique_id = int(time.time())
 
-    output_path = os.path.join(working_dir, "data", "graphgen", f"{unique_id}_{mode}")
+    output_path = os.path.join(working_dir, "data", "graphgen", f"{unique_id}")
     set_working_dir(output_path)
 
     set_logger(
-        os.path.join(output_path, f"{unique_id}.log"),
+        os.path.join(output_path, f"{unique_id}_{mode}.log"),
         if_stream=True,
     )
     logger.info(
         "GraphGen with unique ID %s logging to %s",
         unique_id,
-        os.path.join(working_dir, f"{unique_id}.log"),
+        os.path.join(working_dir, f"{unique_id}_{mode}.log"),
     )
 
-    graph_gen = GraphGen(working_dir=working_dir, output_path=output_path)
+    graph_gen = GraphGen(unique_id=unique_id, working_dir=working_dir)
 
     graph_gen.insert(read_config=config["read"], split_config=config["split"])
 
@@ -81,8 +81,11 @@ def main():
             logger.warning(
                 "Quiz and Judge strategy is disabled. Edge sampling falls back to random."
             )
-            # TODO: make edge sampling random
-        #     graph_gen.traverse_strategy.edge_sampling = "random"
+            assert (
+                config["partition"]["method"] == "ece"
+                and "ece_params" in config["partition"]
+            ), "Only ECE partition with edge sampling is supported."
+            config["partition"]["ece_params"]["edge_sampling"] = "random"
     elif mode == "cot":
         logger.info("Generation mode set to 'cot'. Start generation.")
     else:

graphgen/graphgen.py

@@ -39,10 +39,8 @@
 
 @dataclass
 class GraphGen:
+    unique_id: int = int(time.time())
     working_dir: str = os.path.join(sys_path, "cache")
-    output_path: str = os.path.join(
-        working_dir, "data", "graphgen", str(int(time.time()))
-    )
 
     # llm
     tokenizer_instance: Tokenizer = None
@@ -86,7 +84,7 @@ def __post_init__(self):
             self.working_dir, namespace="rephrase"
         )
         self.qa_storage: JsonListStorage = JsonListStorage(
-            self.working_dir,
+            os.path.join(self.working_dir, "data", "graphgen", f"{self.unique_id}"),
             namespace="qa",
         )
 
@@ -238,59 +236,49 @@ async def quiz_and_judge(self, quiz_and_judge_config: Dict):
     async def generate(self, partition_config: Dict, generate_config: Dict):
         # Step 1: partition the graph
         # TODO: implement graph partitioning, e.g. Partitioner().partition(self.graph_storage)
-        pass
-
-    @async_to_sync_method
-    async def traverse(self):
-        output_data_type = self.config["output_data_type"]
-
-        if output_data_type == "atomic":
+        mode = generate_config["mode"]
+        if mode == "atomic":
             results = await traverse_graph_for_atomic(
                 self.synthesizer_llm_client,
                 self.tokenizer_instance,
                 self.graph_storage,
-                self.traverse_strategy,
+                partition_config["ece_params"],
                 self.text_chunks_storage,
                 self.progress_bar,
             )
-        elif output_data_type == "multi_hop":
+        elif mode == "multi_hop":
             results = await traverse_graph_for_multi_hop(
                 self.synthesizer_llm_client,
                 self.tokenizer_instance,
                 self.graph_storage,
-                self.traverse_strategy,
+                partition_config["ece_params"],
                 self.text_chunks_storage,
                 self.progress_bar,
             )
-        elif output_data_type == "aggregated":
+        elif mode == "aggregated":
             results = await traverse_graph_for_aggregated(
                 self.synthesizer_llm_client,
                 self.tokenizer_instance,
                 self.graph_storage,
-                self.traverse_strategy,
+                partition_config["ece_params"],
                 self.text_chunks_storage,
                 self.progress_bar,
             )
+        elif mode == "cot":
+            method_params = generate_config.get("method_params", {})
+            results = await generate_cot(
+                self.graph_storage,
+                self.synthesizer_llm_client,
+                method_params=method_params,
+            )
         else:
-            raise ValueError(f"Unknown qa_form: {output_data_type}")
-
-        results = format_generation_results(
-            results, output_data_format=self.config["output_data_format"]
-        )
-
-        await self.qa_storage.upsert(results)
-        await self.qa_storage.index_done_callback()
-
-    @async_to_sync_method
-    async def generate_reasoning(self, method_params):
-        results = await generate_cot(
-            self.graph_storage,
-            self.synthesizer_llm_client,
-            method_params=method_params,
-        )
+            raise ValueError(f"Unknown generation mode: {mode}")
+        # Step 2： generate QA pairs
+        # TODO
 
+        # Step 3: format
         results = format_generation_results(
-            results, output_data_format=self.config["output_data_format"]
+            results, output_data_format=generate_config["data_format"]
         )
 
         await self.qa_storage.upsert(results)

graphgen/models/__init__.py

@@ -13,5 +13,4 @@
 from .splitter import ChineseRecursiveTextSplitter, RecursiveCharacterSplitter
 from .storage.json_storage import JsonKVStorage, JsonListStorage
 from .storage.networkx_storage import NetworkXStorage
-from .strategy.travserse_strategy import TraverseStrategy
 from .tokenizer import Tokenizer

graphgen/models/strategy/__init__.py

@@ -1,9 +1,10 @@
 import random
 from collections import defaultdict
+from typing import Dict
 
 from tqdm.asyncio import tqdm as tqdm_async
 
-from graphgen.models import NetworkXStorage, TraverseStrategy
+from graphgen.models import NetworkXStorage
 from graphgen.utils import logger
 
 
@@ -247,18 +248,18 @@ async def get_batches_with_strategy(  # pylint: disable=too-many-branches
     nodes: list,
     edges: list,
     graph_storage: NetworkXStorage,
-    traverse_strategy: TraverseStrategy,
+    traverse_strategy: Dict,
 ):
-    expand_method = traverse_strategy.expand_method
+    expand_method = traverse_strategy["expand_method"]
     if expand_method == "max_width":
         logger.info("Using max width strategy")
     elif expand_method == "max_tokens":
         logger.info("Using max tokens strategy")
     else:
         raise ValueError(f"Invalid expand method: {expand_method}")
 
-    max_depth = traverse_strategy.max_depth
-    edge_sampling = traverse_strategy.edge_sampling
+    max_depth = traverse_strategy["max_depth"]
+    edge_sampling = traverse_strategy["edge_sampling"]
 
     # 构建临接矩阵
     edge_adj_list = defaultdict(list)
@@ -275,16 +276,16 @@ async def get_cached_node_info(node_id: str) -> dict:
     for i, (node_name, _) in enumerate(nodes):
         node_dict[node_name] = i
 
-    if traverse_strategy.loss_strategy == "both":
+    if traverse_strategy["loss_strategy"] == "both":
         er_tuples = [
             ([nodes[node_dict[edge[0]]], nodes[node_dict[edge[1]]]], edge)
             for edge in edges
         ]
         edges = _sort_tuples(er_tuples, edge_sampling)
-    elif traverse_strategy.loss_strategy == "only_edge":
+    elif traverse_strategy["loss_strategy"] == "only_edge":
         edges = _sort_edges(edges, edge_sampling)
     else:
-        raise ValueError(f"Invalid loss strategy: {traverse_strategy.loss_strategy}")
+        raise ValueError(f"Invalid loss strategy: {traverse_strategy['loss_strategy']}")
 
     for i, (src, tgt, _) in enumerate(edges):
         edge_adj_list[src].append(i)
@@ -315,10 +316,10 @@ async def get_cached_node_info(node_id: str) -> dict:
                 nodes,
                 edge,
                 max_depth,
-                traverse_strategy.bidirectional,
-                traverse_strategy.max_extra_edges,
+                traverse_strategy["bidirectional"],
+                traverse_strategy["max_extra_edges"],
                 edge_sampling,
-                traverse_strategy.loss_strategy,
+                traverse_strategy["loss_strategy"],
             )
         else:
             level_n_edges = _get_level_n_edges_by_max_tokens(
@@ -328,10 +329,10 @@ async def get_cached_node_info(node_id: str) -> dict:
                 nodes,
                 edge,
                 max_depth,
-                traverse_strategy.bidirectional,
-                traverse_strategy.max_tokens,
+                traverse_strategy["bidirectional"],
+                traverse_strategy["max_tokens"],
                 edge_sampling,
-                traverse_strategy.loss_strategy,
+                traverse_strategy["loss_strategy"],
             )
 
         for _edge in level_n_edges:
@@ -352,7 +353,7 @@ async def get_cached_node_info(node_id: str) -> dict:
     logger.info("Processing batches: %d", len(processing_batches))
 
     # isolate nodes
-    isolated_node_strategy = traverse_strategy.isolated_node_strategy
+    isolated_node_strategy = traverse_strategy["isolated_node_strategy"]
     if isolated_node_strategy == "add":
         processing_batches = await _add_isolated_nodes(
             nodes, processing_batches, graph_storage

graphgen/models/strategy/travserse_strategy.py

@@ -1,15 +1,10 @@
 import asyncio
+from typing import Dict
 
 import gradio as gr
 from tqdm.asyncio import tqdm as tqdm_async
 
-from graphgen.models import (
-    JsonKVStorage,
-    NetworkXStorage,
-    OpenAIClient,
-    Tokenizer,
-    TraverseStrategy,
-)
+from graphgen.models import JsonKVStorage, NetworkXStorage, OpenAIClient, Tokenizer
 from graphgen.operators.build_kg.split_kg import get_batches_with_strategy
 from graphgen.templates import (
     ANSWER_REPHRASING_PROMPT,
@@ -164,7 +159,7 @@ async def traverse_graph_for_aggregated(
     llm_client: OpenAIClient,
     tokenizer: Tokenizer,
     graph_storage: NetworkXStorage,
-    traverse_strategy: TraverseStrategy,
+    traverse_strategy: Dict,
     text_chunks_storage: JsonKVStorage,
     progress_bar: gr.Progress = None,
     max_concurrent: int = 1000,
@@ -240,7 +235,7 @@ async def _process_single_batch(
                         "question": question,
                         "answer": context,
                         "loss": get_average_loss(
-                            _process_batch, traverse_strategy.loss_strategy
+                            _process_batch, traverse_strategy["loss_strategy"]
                         ),
                     }
                 }
@@ -272,7 +267,7 @@ async def _process_single_batch(
                     "question": qa["question"],
                     "answer": qa["answer"],
                     "loss": get_average_loss(
-                        _process_batch, traverse_strategy.loss_strategy
+                        _process_batch, traverse_strategy["loss_strategy"]
                     ),
                 }
             return final_results
@@ -313,7 +308,7 @@ async def traverse_graph_for_atomic(
     llm_client: OpenAIClient,
     tokenizer: Tokenizer,
     graph_storage: NetworkXStorage,
-    traverse_strategy: TraverseStrategy,
+    traverse_strategy: Dict,
     text_chunks_storage: JsonKVStorage,
     progress_bar: gr.Progress = None,
     max_concurrent: int = 1000,
@@ -331,7 +326,6 @@ async def traverse_graph_for_atomic(
     :return: question and answer
     """
 
-    assert traverse_strategy.qa_form == "atomic"
     semaphore = asyncio.Semaphore(max_concurrent)
 
     def _parse_qa(qa: str) -> tuple:
@@ -429,7 +423,7 @@ async def traverse_graph_for_multi_hop(
     llm_client: OpenAIClient,
     tokenizer: Tokenizer,
     graph_storage: NetworkXStorage,
-    traverse_strategy: TraverseStrategy,
+    traverse_strategy: Dict,
     text_chunks_storage: JsonKVStorage,
     progress_bar: gr.Progress = None,
     max_concurrent: int = 1000,
@@ -517,7 +511,7 @@ async def _process_single_batch(_process_batch: tuple) -> dict:
                         "question": question,
                         "answer": answer,
                         "loss": get_average_loss(
-                            _process_batch, traverse_strategy.loss_strategy
+                            _process_batch, traverse_strategy["loss_strategy"]
                         ),
                     }
                 }