feat: add DFSPartitioner & BFSPartitioner

Author: ChenZiHong-Gavin

graphgen/bases/__init__.py

@@ -1,5 +1,6 @@
 from .base_kg_builder import BaseKGBuilder
 from .base_llm_client import BaseLLMClient
+from .base_partitioner import BasePartitioner
 from .base_reader import BaseReader
 from .base_splitter import BaseSplitter
 from .base_storage import (

graphgen/bases/base_partitioner.py

@@ -0,0 +1,49 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, List
+
+from graphgen.bases.base_storage import BaseGraphStorage
+from graphgen.bases.datatypes import Community
+
+
+@dataclass
+class BasePartitioner(ABC):
+    @abstractmethod
+    async def partition(
+        self,
+        g: BaseGraphStorage,
+        **kwargs: Any,
+    ) -> List[Community]:
+        """
+        Graph -> Communities
+        :param g: Graph storage instance
+        :param kwargs: Additional parameters for partitioning
+        :return: List of communities
+        """
+
+    @abstractmethod
+    def split_communities(self, communities: List[Community]) -> List[Community]:
+        """
+        Split large communities into smaller ones based on max_size.
+        :param communities
+        :return:
+        """
+
+    @staticmethod
+    def _build_adjacency_list(
+        nodes: List[tuple[str, dict]], edges: List[tuple[str, str, dict]]
+    ) -> tuple[dict[str, List[str]], set[tuple[str, str]]]:
+        """
+        Build adjacency list and edge set from nodes and edges.
+        :param nodes
+        :param edges
+        :return: adjacency list, edge set
+        """
+        adj: dict[str, List[str]] = {n[0]: [] for n in nodes}
+        edge_set: set[tuple[str, str]] = set()
+        for e in edges:
+            adj[e[0]].append(e[1])
+            adj[e[1]].append(e[0])
+            edge_set.add((e[0], e[1]))
+            edge_set.add((e[1], e[0]))
+        return adj, edge_set

graphgen/bases/base_storage.py

@@ -78,7 +78,7 @@ async def get_node(self, node_id: str) -> Union[dict, None]:
     async def update_node(self, node_id: str, node_data: dict[str, str]):
         raise NotImplementedError
 
-    async def get_all_nodes(self) -> Union[list[dict], None]:
+    async def get_all_nodes(self) -> Union[list[tuple[str, dict]], None]:
         raise NotImplementedError
 
     async def get_edge(
@@ -91,7 +91,7 @@ async def update_edge(
     ):
         raise NotImplementedError
 
-    async def get_all_edges(self) -> Union[list[dict], None]:
+    async def get_all_edges(self) -> Union[list[tuple[str, str, dict]], None]:
         raise NotImplementedError
 
     async def get_node_edges(

graphgen/bases/datatypes.py

@@ -30,3 +30,11 @@ class Token:
     @property
     def logprob(self) -> float:
         return math.log(self.prob)
+
+
+@dataclass
+class Community:
+    id: Union[int, str]
+    nodes: List[str] = field(default_factory=list)
+    edges: List[tuple] = field(default_factory=list)
+    metadata: dict = field(default_factory=dict)

graphgen/configs/atomic_config.yaml

@@ -17,8 +17,8 @@ partition: # graph partition configuration
     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: 3 # maximum depth for graph traversal
-    max_extra_edges: 5 # max edges per direction (if expand_method="max_width")
+    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
 generate:

graphgen/configs/cot_config.yaml

@@ -9,7 +9,7 @@ search: # web search configuration
 quiz_and_judge: # quiz and test whether the LLM masters the knowledge points
   enabled: false
 partition: # graph partition configuration
-  method: leiden # leiden is a community detection algorithm
+  method: leiden # leiden is a partitioner detection algorithm
   method_params:
     max_size: 20 # Maximum size of communities
     use_lcc: false

graphgen/evaluate.py

@@ -13,7 +13,7 @@
 from .utils import logger, set_logger
 
 sys_path = os.path.abspath(os.path.dirname(__file__))
-set_logger(os.path.join(sys_path, "cache", "logs", "evaluate.log"))
+set_logger(os.path.join(sys_path, "cache", "logs", "evaluator.log"))
 
 load_dotenv()
 

graphgen/graphgen.py

@@ -18,14 +18,11 @@
 from graphgen.operators import (
     build_kg,
     chunk_documents,
-    generate_cot,
     judge_statement,
+    partition_kg,
     quiz,
     read_files,
     search_all,
-    traverse_graph_for_aggregated,
-    traverse_graph_for_atomic,
-    traverse_graph_for_multi_hop,
 )
 from graphgen.utils import (
     async_to_sync_method,
@@ -237,54 +234,55 @@ 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
-        # TODO: implement graph partitioning, e.g. Partitioner().partition(self.graph_storage)
-        mode = generate_config["mode"]
-        if mode == "atomic":
-            results = await traverse_graph_for_atomic(
-                self.synthesizer_llm_client,
-                self.tokenizer_instance,
-                self.graph_storage,
-                partition_config["method_params"],
-                self.text_chunks_storage,
-                self.progress_bar,
-            )
-        elif mode == "multi_hop":
-            results = await traverse_graph_for_multi_hop(
-                self.synthesizer_llm_client,
-                self.tokenizer_instance,
-                self.graph_storage,
-                partition_config["method_params"],
-                self.text_chunks_storage,
-                self.progress_bar,
-            )
-        elif mode == "aggregated":
-            results = await traverse_graph_for_aggregated(
-                self.synthesizer_llm_client,
-                self.tokenizer_instance,
-                self.graph_storage,
-                partition_config["method_params"],
-                self.text_chunks_storage,
-                self.progress_bar,
-            )
-        elif mode == "cot":
-            results = await generate_cot(
-                self.graph_storage,
-                self.synthesizer_llm_client,
-                method_params=partition_config["method_params"],
-            )
-        else:
-            raise ValueError(f"Unknown generation mode: {mode}")
+        # mode = generate_config["mode"]
+        # batches = partition_kg(self.graph_storage, partition_config)
+        # if mode == "atomic":
+        #     results = await traverse_graph_for_atomic(
+        #         self.synthesizer_llm_client,
+        #         self.tokenizer_instance,
+        #         self.graph_storage,
+        #         partition_config["method_params"],
+        #         self.text_chunks_storage,
+        #         self.progress_bar,
+        #     )
+        # elif mode == "multi_hop":
+        #     results = await traverse_graph_for_multi_hop(
+        #         self.synthesizer_llm_client,
+        #         self.tokenizer_instance,
+        #         self.graph_storage,
+        #         partition_config["method_params"],
+        #         self.text_chunks_storage,
+        #         self.progress_bar,
+        #     )
+        # elif mode == "aggregated":
+        #     results = await traverse_graph_for_aggregated(
+        #         self.synthesizer_llm_client,
+        #         self.tokenizer_instance,
+        #         self.graph_storage,
+        #         partition_config["method_params"],
+        #         self.text_chunks_storage,
+        #         self.progress_bar,
+        #     )
+        # elif mode == "cot":
+        #     results = await generate_cot(
+        #         self.graph_storage,
+        #         self.synthesizer_llm_client,
+        #         method_params=partition_config["method_params"],
+        #     )
+        # else:
+        #     raise ValueError(f"Unknown generation mode: {mode}")
         # Step 2： generate QA pairs
         # TODO
 
         # Step 3: format
-        results = format_generation_results(
-            results, output_data_format=generate_config["data_format"]
-        )
-
-        await self.qa_storage.upsert(results)
-        await self.qa_storage.index_done_callback()
+        # results = format_generation_results(
+        #     results, output_data_format=generate_config["data_format"]
+        # )
+        #
+        # await self.qa_storage.upsert(results)
+        # await self.qa_storage.index_done_callback()
 
     @async_to_sync_method
     async def clear(self):

graphgen/models/__init__.py

@@ -1,11 +1,13 @@
-from .community.community_detector import CommunityDetector
-from .evaluate.length_evaluator import LengthEvaluator
-from .evaluate.mtld_evaluator import MTLDEvaluator
-from .evaluate.reward_evaluator import RewardEvaluator
-from .evaluate.uni_evaluator import UniEvaluator
-from .kg_builder.light_rag_kg_builder import LightRAGKGBuilder
+from .evaluator import LengthEvaluator, MTLDEvaluator, RewardEvaluator, UniEvaluator
+from .kg_builder import LightRAGKGBuilder
 from .llm.openai_client import OpenAIClient
 from .llm.topk_token_model import TopkTokenModel
+from .partitioner import (
+    BFSPartitioner,
+    DFSPartitioner,
+    ECEPartitioner,
+    LeidenPartitioner,
+)
 from .reader import CsvReader, JsonlReader, JsonReader, TxtReader
 from .search.db.uniprot_search import UniProtSearch
 from .search.kg.wiki_search import WikiSearch