feat: add cot data generation pipeline

Author: ChenZiHong-Gavin

README.md

@@ -176,7 +176,8 @@ See [analysis](https://deepwiki.com/open-sciencelab/GraphGen) by deepwiki for a
 ## 🍀 Acknowledgements
 - [SiliconFlow](https://siliconflow.cn) Abundant LLM API, some models are free
 - [LightRAG](https://github.com/HKUDS/LightRAG) Simple and efficient graph retrieval solution
-- [ROGRAG](https://github.com/tpoisonooo/ROGRAG) ROGRAG: A Robustly Optimized GraphRAG Framework
+- [ROGRAG](https://github.com/tpoisonooo/ROGRAG) A robustly optimized GraphRAG framework
+- [DB-GPT](https://github.com/eosphoros-ai/DB-GPT) An AI native data app development framework
 
 
 ## 📚 Citation

graphgen/models/__init__.py

@@ -1,3 +1,4 @@
+from .community.community_detector import CommunityDetector
 from .evaluate.length_evaluator import LengthEvaluator
 from .evaluate.mtld_evaluator import MTLDEvaluator
 from .evaluate.reward_evaluator import RewardEvaluator
@@ -38,4 +39,6 @@
     "UniEvaluator",
     # strategy models
     "TraverseStrategy",
+    # community models
+    "CommunityDetector",
 ]

graphgen/models/community/__init__.py

@@ -0,0 +1,67 @@
+from dataclasses import dataclass
+from typing import Any, Dict
+
+from graphgen.models.storage.networkx_storage import NetworkXStorage
+
+
+@dataclass
+class CommunityDetector:
+    """Class for community detection algorithms."""
+
+    graph_storage: NetworkXStorage = None
+    method: str = "leiden"
+    method_params: Dict[str, Any] = None
+
+    async def detect_communities(self) -> Dict[str, int]:
+        """
+        Detect communities based on the chosen method.
+        """
+        if self.method == "leiden":
+            return await self._leiden_communities(**self.method_params or {})
+        raise ValueError(f"Unknown community detection method: {self.method}")
+
+    async def get_graph(self):
+        """
+        Asynchronously get the graph from the storage.
+        """
+        return await self.graph_storage.get_graph()
+
+    async def _leiden_communities(self, **kwargs) -> Dict[str, int]:
+        """
+        Detect communities using the Leiden algorithm.
+        """
+        import igraph as ig
+        import networkx as nx
+        from leidenalg import ModularityVertexPartition, find_partition
+
+        graph = await self.get_graph()
+        # Filter out isolated nodes
+        graph.remove_nodes_from(list(nx.isolates(graph)))
+
+        # Convert NetworkX graph to igraph graph
+        ig_graph = ig.Graph.TupleList(graph.edges(), directed=False)
+
+        random_seed = kwargs.get("random_seed", 42)
+        use_lcc = kwargs.get("use_lcc", False)
+
+        communities = {}
+        if use_lcc:
+            # Use the largest connected component
+            lcc = ig_graph.components().giant()
+            partition = find_partition(lcc, ModularityVertexPartition, seed=random_seed)
+            for part, cluster in enumerate(partition):
+                for v in cluster:
+                    communities[v] = part
+        else:
+            offset = 0
+            for component in ig_graph.components():
+                subgraph = ig_graph.induced_subgraph(component)
+                partition = find_partition(
+                    subgraph, ModularityVertexPartition, seed=random_seed
+                )
+                for part, cluster in enumerate(partition):
+                    for v in cluster:
+                        original_node = subgraph.vs[v]["name"]
+                        communities[original_node] = part + offset
+                offset += len(partition)
+        return communities

graphgen/models/community/community_detector.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass
+from typing import Dict
+
+import matplotlib.pyplot as plt
+import networkx as nx
+
+
+@dataclass
+class Visualizer:
+    """
+    Class for visualizing graphs using NetworkX and Matplotlib.
+    """
+
+    graph: nx.Graph = None
+    communities: Dict[str, int] = None
+    layout: str = "spring"
+    max_nodes: int = 1000
+    node_size: int = 10
+    alpha: float = 0.6
+
+    def visualize(self, save_path: str = None):
+        n = self.graph.number_of_nodes()
+        print(f"Loaded graph: {n} nodes, {self.graph.number_of_edges()} edges")
+
+        if self.layout == "spring":
+            k = max(0.1, 1.0 / (n**0.5))
+            pos = nx.spring_layout(self.graph, k=k, seed=42)
+        else:
+            raise ValueError(f"Unknown layout: {self.layout}")
+
+        plt.figure(figsize=(10, 10))
+
+        node_colors = [self.communities.get(node, 0) for node in self.graph.nodes()]
+        print(node_colors)
+
+        nx.draw_networkx_nodes(
+            self.graph,
+            pos,
+            node_size=self.node_size,
+            node_color=node_colors,
+            cmap="viridis",
+            alpha=self.alpha,
+        )
+        nx.draw_networkx_edges(self.graph, pos, alpha=0.3, width=0.2)
+        plt.axis("off")
+
+        if save_path:
+            plt.savefig(save_path, dpi=300, bbox_inches="tight")
+            print("Saved to", save_path)
+        else:
+            plt.show()

graphgen/models/vis/__init__.py

@@ -0,0 +1,110 @@
+import asyncio
+from typing import Dict, List
+
+from tqdm.asyncio import tqdm as tqdm_async
+
+from graphgen.models import CommunityDetector, NetworkXStorage, OpenAIModel
+from graphgen.templates import COT_GENERATION_PROMPT, COT_TEMPLATE_DESIGN_PROMPT
+from graphgen.utils import detect_main_language
+
+
+async def generate_cot(
+    graph_storage: NetworkXStorage,
+    synthesizer_llm_client: OpenAIModel,
+    method: str = "leiden",
+):
+    detector = CommunityDetector(graph_storage=graph_storage, method=method)
+
+    results = await detector.detect_communities()
+
+    # Convert results to a format suitable for summarization
+    communities = {}
+    for node, community_id in results.items():
+        if community_id not in communities:
+            communities[community_id] = []
+        communities[community_id].append(node)
+
+    if not communities:
+        return {}
+
+    semaphore = asyncio.Semaphore(value=1000)
+
+    async def _generate_from_single_community(
+        community_id: int, nodes: List[str]
+    ) -> tuple[int, tuple]:
+        """Summarize a single community."""
+        async with semaphore:
+            entities: List[str] = []
+            relationships: List[str] = []
+
+            for node in nodes:
+                node_data = await graph_storage.get_node(node)
+                if node_data is not None:
+                    entities.append(f"({node}: {node_data.get('description')})")
+
+                edges = await graph_storage.get_node_edges(node)
+                for edge in edges:
+                    target = edge[1]
+                    if target in nodes:
+                        edge_data = await graph_storage.get_edge(node, target)
+                        relationships.append(
+                            f"({node}) - [{edge_data['description']}] -> ({target})"
+                        )
+
+            entities_str = "\n".join(entities)
+            relationships_str = "\n".join(relationships)
+
+            language = (
+                "English"
+                if detect_main_language(entities_str + relationships_str) == "en"
+                else "Chinese"
+            )
+
+            prompt = COT_TEMPLATE_DESIGN_PROMPT[language]["TEMPLATE"].format(
+                entities=entities_str,
+                relationships=relationships_str,
+            )
+
+            cot_template = await synthesizer_llm_client.generate_answer(prompt)
+
+            if "问题：" in cot_template and "推理路径设计：" in cot_template:
+                question = cot_template.split("问题：")[1].split("推理路径设计：")[0].strip()
+                reasoning_path = cot_template.split("推理路径设计：")[1].strip()
+            elif (
+                "Question:" in cot_template and "Reasoning-Path Design:" in cot_template
+            ):
+                question = (
+                    cot_template.split("Question:")[1]
+                    .split("Reasoning-Path Design:")[0]
+                    .strip()
+                )
+                reasoning_path = cot_template.split("Reasoning-Path Design:")[1].strip()
+            else:
+                raise ValueError("COT template format is incorrect.")
+
+            prompt = COT_GENERATION_PROMPT[language]["TEMPLATE"].format(
+                entities=entities_str,
+                relationships=relationships_str,
+                question=question,
+                reasoning_template=reasoning_path,
+            )
+
+            cot_answer = await synthesizer_llm_client.generate_answer(prompt)
+
+            return community_id, (question, reasoning_path, cot_answer)
+
+    cid_nodes = list(communities.items())
+
+    templates: Dict[int, (str, str)] = {}
+    async for coro in tqdm_async(
+        asyncio.as_completed(
+            [_generate_from_single_community(cid, nodes) for cid, nodes in cid_nodes]
+        ),
+        total=len(cid_nodes),
+        desc="[Generate COT] Generating COT templates for communities",
+        unit="community",
+    ):
+        cid, (q, r, a) = await coro
+        templates[cid] = (q, r, a)
+
+    return templates

graphgen/models/vis/community_visualizer.py

@@ -1,9 +1,10 @@
+from .answer_rephrasing import ANSWER_REPHRASING_PROMPT
+from .community import COT_GENERATION_PROMPT, COT_TEMPLATE_DESIGN_PROMPT
+from .coreference_resolution import COREFERENCE_RESOLUTION_PROMPT
+from .description_rephrasing import DESCRIPTION_REPHRASING_PROMPT
 from .kg_extraction import KG_EXTRACTION_PROMPT
 from .kg_summarization import KG_SUMMARIZATION_PROMPT
+from .multi_hop_generation import MULTI_HOP_GENERATION_PROMPT
+from .question_generation import QUESTION_GENERATION_PROMPT
 from .search_judgement import SEARCH_JUDGEMENT_PROMPT
-from .description_rephrasing import DESCRIPTION_REPHRASING_PROMPT
 from .statement_judgement import STATEMENT_JUDGEMENT_PROMPT
-from .answer_rephrasing import ANSWER_REPHRASING_PROMPT
-from .question_generation import QUESTION_GENERATION_PROMPT
-from .multi_hop_generation import MULTI_HOP_GENERATION_PROMPT
-from .coreference_resolution import COREFERENCE_RESOLUTION_TEMPLATE

graphgen/operators/community/__init__.py

@@ -1,5 +1,4 @@
 TEMPLATE_CONTEXT_EN: str = """---Role---
-
 You are an NLP expert responsible for generating a logically structured and coherent rephrased version of the TEXT based on ENTITIES and RELATIONSHIPS provided below. You may refer to the original text to assist in generating the rephrased version, but ensure that the final output text meets the requirements.
 Use {language} as output language.
 
@@ -51,12 +50,10 @@
 """
 
 TEMPLATE_CONTEXT_ZH: str = """---角色---
-
 你是一位NLP专家，负责根据下面提供的实体和关系生成逻辑结构清晰且连贯的文本重述版本。你可以参考原始文本辅助生成，但需要确保最终输出的文本符合要求。
 使用{language}作为输出语言。
 
 ---目标---
-
 生成文本的重述版本，使其传达与原始实体和关系描述相同的含义，同时：
 1. 遵循清晰的逻辑流和结构
 2. 建立适当的因果关系
@@ -101,7 +98,6 @@
 """
 
 TEMPLATE_EN: str = """---Role---
-
 You are an NLP expert responsible for generating a logically structured and coherent rephrased version of the TEXT based on ENTITIES and RELATIONSHIPS provided below.
 Use {language} as output language.
 
@@ -148,12 +144,10 @@
 """
 
 TEMPLATE_ZH: str = """---角色---
-
 你是一位NLP专家，负责根据下面提供的实体和关系生成逻辑结构清晰且连贯的文本重述版本。
 使用{language}作为输出语言。
 
 ---目标---
-
 生成文本的重述版本，使其传达与原始实体和关系描述相同的含义，同时：
 1. 遵循清晰的逻辑流和结构
 2. 建立适当的因果关系
@@ -207,13 +201,13 @@
 """
 
 
-ANSWER_REPHRASING_PROMPT= {
+ANSWER_REPHRASING_PROMPT = {
     "English": {
         "TEMPLATE": TEMPLATE_EN + REQUIREMENT_EN,
-        "CONTEXT_TEMPLATE": TEMPLATE_CONTEXT_EN + REQUIREMENT_EN
+        "CONTEXT_TEMPLATE": TEMPLATE_CONTEXT_EN + REQUIREMENT_EN,
     },
     "Chinese": {
         "TEMPLATE": TEMPLATE_ZH + REQUIREMENT_ZH,
-        "CONTEXT_TEMPLATE": TEMPLATE_CONTEXT_ZH + REQUIREMENT_ZH
-    }
+        "CONTEXT_TEMPLATE": TEMPLATE_CONTEXT_ZH + REQUIREMENT_ZH,
+    },
 }