feat: add community2batch method

Author: ChenZiHong-Gavin

graphgen/bases/base_partitioner.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, List
+from typing import Any, List, Tuple
 
 from graphgen.bases.base_storage import BaseGraphStorage
 from graphgen.bases.datatypes import Community
@@ -29,6 +29,41 @@ def split_communities(self, communities: List[Community]) -> List[Community]:
         :return:
         """
 
+    @staticmethod
+    async def community2batch(
+        communities: List[Community], g: BaseGraphStorage
+    ) -> list[
+        tuple[
+            list[tuple[str, dict]], list[tuple[Any, Any, dict] | tuple[Any, Any, Any]]
+        ]
+    ]:
+        """
+        Convert communities to batches of nodes and edges.
+        :param communities
+        :param g: Graph storage instance
+        :return: List of batches, each batch is a tuple of (nodes, edges)
+        """
+        batches = []
+        for comm in communities:
+            nodes = comm.nodes
+            edges = comm.edges
+            nodes_data = []
+            for node in nodes:
+                node_data = await g.get_node(node)
+                if node_data:
+                    nodes_data.append((node, node_data))
+            edges_data = []
+            for u, v in edges:
+                edge_data = await g.get_edge(u, v)
+                if edge_data:
+                    edges_data.append((u, v, edge_data))
+                else:
+                    edge_data = await g.get_edge(v, u)
+                    if edge_data:
+                        edges_data.append((v, u, edge_data))
+            batches.append((nodes_data, edges_data))
+        return batches
+
     @staticmethod
     def _build_adjacency_list(
         nodes: List[tuple[str, dict]], edges: List[tuple[str, str, dict]]

graphgen/configs/atomic_config.yaml

@@ -11,16 +11,9 @@ quiz_and_judge: # quiz and test whether the LLM masters the knowledge points
   quiz_samples: 2 # number of quiz samples to generate
   re_judge: false # whether to re-judge the existing quiz samples
 partition: # graph partition configuration
-  method: ece # ece is a custom partition method based on comprehension loss
+  method: dfs # partition method, support: dfs, bfs, ece, leiden
   method_params:
-    bidirectional: true # whether to traverse the graph in both directions
-    edge_sampling: max_loss # edge sampling strategy, support: random, max_loss, min_loss
-    expand_method: max_width # expand method, support: max_width, max_depth
-    isolated_node_strategy: ignore # strategy for isolated nodes, support: ignore, add
-    max_depth: 1 # maximum depth for graph traversal
-    max_extra_edges: 0 # max edges per direction (if expand_method="max_width")
-    max_tokens: 256 # restricts input length (if expand_method="max_tokens")
-    loss_strategy: only_edge # defines loss computation focus, support: only_edge, both
+    max_units_per_community: 1 # atomic partition, one node or edge per community
 generate:
   mode: atomic # atomic, aggregated, multi_hop, cot
   data_format: Alpaca # Alpaca, Sharegpt, ChatML

graphgen/graphgen.py

@@ -5,6 +5,7 @@
 from typing import Dict, cast
 
 import gradio as gr
+from jieba.lac_small.predict import results
 
 from graphgen.bases.base_storage import StorageNameSpace
 from graphgen.bases.datatypes import Chunk
@@ -234,10 +235,13 @@ async def quiz_and_judge(self, quiz_and_judge_config: Dict):
 
     @async_to_sync_method
     async def generate(self, partition_config: Dict, generate_config: Dict):
-        pass
         # Step 1: partition the graph
-        # mode = generate_config["mode"]
-        # batches = partition_kg(self.graph_storage, partition_config)
+        batches = await partition_kg(self.graph_storage, partition_config)
+
+        # Step 2： generate QA pairs
+        mode = generate_config["mode"]
+        logger.info("[Generation] mode: %s, batches: %d", mode, len(batches))
+        # results = generate_qa_pairs(generate_config)
         # if mode == "atomic":
         #     results = await traverse_graph_for_atomic(
         #         self.synthesizer_llm_client,
@@ -273,8 +277,6 @@ async def generate(self, partition_config: Dict, generate_config: Dict):
         #     )
         # else:
         #     raise ValueError(f"Unknown generation mode: {mode}")
-        # Step 2： generate QA pairs
-        # TODO
 
         # Step 3: format
         # results = format_generation_results(

graphgen/operators/partition/partition_kg.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import Any, List, Tuple
 
 from graphgen.bases import BaseGraphStorage
 from graphgen.bases.datatypes import Community
@@ -11,10 +11,12 @@
 from graphgen.utils import logger
 
 
-def partition_kg(
+async def partition_kg(
     kg_instance: BaseGraphStorage,
     partition_config: dict = None,
-) -> List[Community]:
+) -> list[
+    tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict] | tuple[Any, Any, Any]]]
+]:
     method = partition_config["method"]
     method_params = partition_config["method_params"]
     if method == "bfs":
@@ -32,6 +34,7 @@ def partition_kg(
     else:
         raise ValueError(f"Unsupported partition method: {method}")
 
-    communities = partitioner.partition(g=kg_instance, **method_params)
-    logger.info(f"Partitioned the graph into {len(communities)} communities.")
-    return communities
+    communities = await partitioner.partition(g=kg_instance, **method_params)
+    logger.info("Partitioned the graph into %d communities.", len(communities))
+    batches = await partitioner.community2batch(communities, g=kg_instance)
+    return batches