Auto-sync from demo at Fri Oct 10 04:18:04 UTC 2025

graphgen/bases/base_kg_builder.py

@@ -10,7 +10,6 @@ from graphgen.bases.datatypes import Chunk
 
 @dataclass
 class BaseKGBuilder(ABC):
-    kg_instance: BaseGraphStorage
     llm_client: BaseLLMClient
 
     _nodes: Dict[str, List[dict]] = field(default_factory=lambda: defaultdict(list))
@@ -18,14 +17,6 @@ class BaseKGBuilder(ABC):
         default_factory=lambda: defaultdict(list)
     )
 
-    def build(self, chunks: List[Chunk]) -> None:
-        pass
-
-    @abstractmethod
-    async def extract_all(self, chunks: List[Chunk]) -> None:
-        """Extract nodes and edges from all chunks."""
-        raise NotImplementedError
-
     @abstractmethod
     async def extract(
         self, chunk: Chunk
@@ -35,7 +26,18 @@ class BaseKGBuilder(ABC):
 
     @abstractmethod
     async def merge_nodes(
-        self, nodes_data: Dict[str, List[dict]], kg_instance: BaseGraphStorage, llm
+        self,
+        node_data: tuple[str, List[dict]],
+        kg_instance: BaseGraphStorage,
     ) -> None:
         """Merge extracted nodes into the knowledge graph."""
         raise NotImplementedError
+
+    @abstractmethod
+    async def merge_edges(
+        self,
+        edges_data: tuple[Tuple[str, str], List[dict]],
+        kg_instance: BaseGraphStorage,
+    ) -> None:
+        """Merge extracted edges into the knowledge graph."""
+        raise NotImplementedError

graphgen/bases/base_llm_client.py

@@ -57,12 +57,6 @@ class BaseLLMClient(abc.ABC):
         """Generate probabilities for each token in the input."""
         raise NotImplementedError
 
-    def count_tokens(self, text: str) -> int:
-        """Count the number of tokens in the text."""
-        if self.tokenizer is None:
-            raise ValueError("Tokenizer is not set. Please provide a tokenizer to use count_tokens.")
-        return len(self.tokenizer.encode(text))
-
     @staticmethod
     def filter_think_tags(text: str, think_tag: str = "think") -> str:
         """

graphgen/graphgen.py

@@ -16,8 +16,8 @@ from graphgen.models import (
     Tokenizer,
 )
 from graphgen.operators import (
+    build_kg,
     chunk_documents,
-    extract_kg,
     generate_cot,
     judge_statement,
     quiz,
@@ -146,10 +146,9 @@ class GraphGen:
 
         # Step 3: Extract entities and relations from chunks
         logger.info("[Entity and Relation Extraction]...")
-        _add_entities_and_relations = await extract_kg(
+        _add_entities_and_relations = await build_kg(
             llm_client=self.synthesizer_llm_client,
             kg_instance=self.graph_storage,
-            tokenizer_instance=self.tokenizer_instance,
             chunks=[
                 Chunk(id=k, content=v["content"]) for k, v in inserting_chunks.items()
             ],

graphgen/models/__init__.py

@@ -3,6 +3,7 @@ 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 .llm.openai_client import OpenAIClient
 from .llm.topk_token_model import TopkTokenModel
 from .reader import CsvReader, JsonlReader, JsonReader, TxtReader

graphgen/models/kg_builder/NetworkXKGBuilder.py

@@ -1,18 +0,0 @@
-from dataclasses import dataclass
-
-from graphgen.bases import BaseKGBuilder
-
-
-@dataclass
-class NetworkXKGBuilder(BaseKGBuilder):
-    def build(self, chunks):
-        pass
-
-    async def extract_all(self, chunks):
-        pass
-
-    async def extract(self, chunk):
-        pass
-
-    async def merge_nodes(self, nodes_data, kg_instance, llm):
-        pass

graphgen/models/kg_builder/light_rag_kg_builder.py

@@ -0,0 +1,226 @@
+import re
+from collections import Counter, defaultdict
+from dataclasses import dataclass
+from typing import Dict, List, Tuple
+
+from graphgen.bases import BaseGraphStorage, BaseKGBuilder, BaseLLMClient, Chunk
+from graphgen.templates import KG_EXTRACTION_PROMPT, KG_SUMMARIZATION_PROMPT
+from graphgen.utils import (
+    detect_if_chinese,
+    detect_main_language,
+    handle_single_entity_extraction,
+    handle_single_relationship_extraction,
+    logger,
+    pack_history_conversations,
+    split_string_by_multi_markers,
+)
+
+
+@dataclass
+class LightRAGKGBuilder(BaseKGBuilder):
+    llm_client: BaseLLMClient = None
+    max_loop: int = 3
+
+    async def extract(
+        self, chunk: Chunk
+    ) -> Tuple[Dict[str, List[dict]], Dict[Tuple[str, str], List[dict]]]:
+        """
+        Extract entities and relationships from a single chunk using the LLM client.
+        :param chunk
+        :return: (nodes_data, edges_data)
+        """
+        chunk_id = chunk.id
+        content = chunk.content
+
+        # step 1: language_detection
+        language = "Chinese" if detect_if_chinese(content) else "English"
+        KG_EXTRACTION_PROMPT["FORMAT"]["language"] = language
+
+        hint_prompt = KG_EXTRACTION_PROMPT[language]["TEMPLATE"].format(
+            **KG_EXTRACTION_PROMPT["FORMAT"], input_text=content
+        )
+
+        # step 2: initial glean
+        final_result = await self.llm_client.generate_answer(hint_prompt)
+        logger.debug("First extraction result: %s", final_result)
+
+        # step3: iterative refinement
+        history = pack_history_conversations(hint_prompt, final_result)
+        for loop_idx in range(self.max_loop):
+            if_loop_result = await self.llm_client.generate_answer(
+                text=KG_EXTRACTION_PROMPT[language]["IF_LOOP"], history=history
+            )
+            if_loop_result = if_loop_result.strip().strip('"').strip("'").lower()
+            if if_loop_result != "yes":
+                break
+
+            glean_result = await self.llm_client.generate_answer(
+                text=KG_EXTRACTION_PROMPT[language]["CONTINUE"], history=history
+            )
+            logger.debug("Loop %s glean: %s", loop_idx + 1, glean_result)
+
+            history += pack_history_conversations(
+                KG_EXTRACTION_PROMPT[language]["CONTINUE"], glean_result
+            )
+            final_result += glean_result
+
+        # step 4: parse the final result
+        records = split_string_by_multi_markers(
+            final_result,
+            [
+                KG_EXTRACTION_PROMPT["FORMAT"]["record_delimiter"],
+                KG_EXTRACTION_PROMPT["FORMAT"]["completion_delimiter"],
+            ],
+        )
+
+        nodes = defaultdict(list)
+        edges = defaultdict(list)
+
+        for record in records:
+            match = re.search(r"\((.*)\)", record)
+            if not match:
+                continue
+            inner = match.group(1)
+
+            attributes = split_string_by_multi_markers(
+                inner, [KG_EXTRACTION_PROMPT["FORMAT"]["tuple_delimiter"]]
+            )
+
+            entity = await handle_single_entity_extraction(attributes, chunk_id)
+            if entity is not None:
+                nodes[entity["entity_name"]].append(entity)
+                continue
+
+            relation = await handle_single_relationship_extraction(attributes, chunk_id)
+            if relation is not None:
+                key = (relation["src_id"], relation["tgt_id"])
+                edges[key].append(relation)
+
+        return dict(nodes), dict(edges)
+
+    async def merge_nodes(
+        self,
+        node_data: tuple[str, List[dict]],
+        kg_instance: BaseGraphStorage,
+    ) -> None:
+        entity_name, node_data = node_data
+        entity_types = []
+        source_ids = []
+        descriptions = []
+
+        node = await kg_instance.get_node(entity_name)
+        if node is not None:
+            entity_types.append(node["entity_type"])
+            source_ids.extend(
+                split_string_by_multi_markers(node["source_id"], ["<SEP>"])
+            )
+            descriptions.append(node["description"])
+
+        # take the most frequent entity_type
+        entity_type = sorted(
+            Counter([dp["entity_type"] for dp in node_data] + entity_types).items(),
+            key=lambda x: x[1],
+            reverse=True,
+        )[0][0]
+
+        description = "<SEP>".join(
+            sorted(set([dp["description"] for dp in node_data] + descriptions))
+        )
+        description = await self._handle_kg_summary(entity_name, description)
+
+        source_id = "<SEP>".join(
+            set([dp["source_id"] for dp in node_data] + source_ids)
+        )
+
+        node_data = {
+            "entity_type": entity_type,
+            "description": description,
+            "source_id": source_id,
+        }
+        await kg_instance.upsert_node(entity_name, node_data=node_data)
+
+    async def merge_edges(
+        self,
+        edges_data: tuple[Tuple[str, str], List[dict]],
+        kg_instance: BaseGraphStorage,
+    ) -> None:
+        (src_id, tgt_id), edge_data = edges_data
+
+        source_ids = []
+        descriptions = []
+
+        edge = await kg_instance.get_edge(src_id, tgt_id)
+        if edge is not None:
+            source_ids.extend(
+                split_string_by_multi_markers(edge["source_id"], ["<SEP>"])
+            )
+            descriptions.append(edge["description"])
+
+        description = "<SEP>".join(
+            sorted(set([dp["description"] for dp in edge_data] + descriptions))
+        )
+        source_id = "<SEP>".join(
+            set([dp["source_id"] for dp in edge_data] + source_ids)
+        )
+
+        for insert_id in [src_id, tgt_id]:
+            if not await kg_instance.has_node(insert_id):
+                await kg_instance.upsert_node(
+                    insert_id,
+                    node_data={
+                        "source_id": source_id,
+                        "description": description,
+                        "entity_type": "UNKNOWN",
+                    },
+                )
+
+        description = await self._handle_kg_summary(
+            f"({src_id}, {tgt_id})", description
+        )
+
+        await kg_instance.upsert_edge(
+            src_id,
+            tgt_id,
+            edge_data={"source_id": source_id, "description": description},
+        )
+
+    async def _handle_kg_summary(
+        self,
+        entity_or_relation_name: str,
+        description: str,
+        max_summary_tokens: int = 200,
+    ) -> str:
+        """
+        Handle knowledge graph summary
+
+        :param entity_or_relation_name
+        :param description
+        :param max_summary_tokens
+        :return summary
+        """
+
+        tokenizer_instance = self.llm_client.tokenizer
+        language = detect_main_language(description)
+        if language == "en":
+            language = "English"
+        else:
+            language = "Chinese"
+        KG_EXTRACTION_PROMPT["FORMAT"]["language"] = language
+
+        tokens = tokenizer_instance.encode(description)
+        if len(tokens) < max_summary_tokens:
+            return description
+
+        use_description = tokenizer_instance.decode(tokens[:max_summary_tokens])
+        prompt = KG_SUMMARIZATION_PROMPT[language]["TEMPLATE"].format(
+            entity_name=entity_or_relation_name,
+            description_list=use_description.split("<SEP>"),
+            **KG_SUMMARIZATION_PROMPT["FORMAT"],
+        )
+        new_description = await self.llm_client.generate_answer(prompt)
+        logger.info(
+            "Entity or relation %s summary: %s",
+            entity_or_relation_name,
+            new_description,
+        )
+        return new_description

graphgen/operators/__init__.py

@@ -1,4 +1,4 @@
-from graphgen.operators.build_kg.extract_kg import extract_kg
+from graphgen.operators.build_kg.build_kg import build_kg
 from graphgen.operators.generate.generate_cot import generate_cot
 from graphgen.operators.search.search_all import search_all
 

graphgen/operators/build_kg/build_kg.py

@@ -0,0 +1,56 @@
+from collections import defaultdict
+from typing import List
+
+import gradio as gr
+
+from graphgen.bases.base_storage import BaseGraphStorage
+from graphgen.bases.datatypes import Chunk
+from graphgen.models import LightRAGKGBuilder, OpenAIClient
+from graphgen.utils import run_concurrent
+
+
+async def build_kg(
+    llm_client: OpenAIClient,
+    kg_instance: BaseGraphStorage,
+    chunks: List[Chunk],
+    progress_bar: gr.Progress = None,
+):
+    """
+    :param llm_client: Synthesizer LLM model to extract entities and relationships
+    :param kg_instance
+    :param chunks
+    :param progress_bar: Gradio progress bar to show the progress of the extraction
+    :return:
+    """
+
+    kg_builder = LightRAGKGBuilder(llm_client=llm_client, max_loop=3)
+
+    results = await run_concurrent(
+        kg_builder.extract,
+        chunks,
+        desc="[2/4]Extracting entities and relationships from chunks",
+        unit="chunk",
+        progress_bar=progress_bar,
+    )
+
+    nodes = defaultdict(list)
+    edges = defaultdict(list)
+    for n, e in results:
+        for k, v in n.items():
+            nodes[k].extend(v)
+        for k, v in e.items():
+            edges[tuple(sorted(k))].extend(v)
+
+    await run_concurrent(
+        lambda kv: kg_builder.merge_nodes(kv, kg_instance=kg_instance),
+        list(nodes.items()),
+        desc="Inserting entities into storage",
+    )
+
+    await run_concurrent(
+        lambda kv: kg_builder.merge_edges(kv, kg_instance=kg_instance),
+        list(edges.items()),
+        desc="Inserting relationships into storage",
+    )
+
+    return kg_instance

graphgen/operators/build_kg/extract_kg.py

@@ -1,127 +0,0 @@
-import re
-from collections import defaultdict
-from typing import List
-
-import gradio as gr
-
-from graphgen.bases.base_storage import BaseGraphStorage
-from graphgen.bases.datatypes import Chunk
-from graphgen.models import OpenAIClient, Tokenizer
-from graphgen.operators.build_kg.merge_kg import merge_edges, merge_nodes
-from graphgen.templates import KG_EXTRACTION_PROMPT
-from graphgen.utils import (
-    detect_if_chinese,
-    handle_single_entity_extraction,
-    handle_single_relationship_extraction,
-    logger,
-    pack_history_conversations,
-    run_concurrent,
-    split_string_by_multi_markers,
-)
-
-
-# pylint: disable=too-many-statements
-async def extract_kg(
-    llm_client: OpenAIClient,
-    kg_instance: BaseGraphStorage,
-    tokenizer_instance: Tokenizer,
-    chunks: List[Chunk],
-    progress_bar: gr.Progress = None,
-):
-    """
-    :param llm_client: Synthesizer LLM model to extract entities and relationships
-    :param kg_instance
-    :param tokenizer_instance
-    :param chunks
-    :param progress_bar: Gradio progress bar to show the progress of the extraction
-    :return:
-    """
-
-    async def _process_single_content(chunk: Chunk, max_loop: int = 3):
-        chunk_id = chunk.id
-        content = chunk.content
-        if detect_if_chinese(content):
-            language = "Chinese"
-        else:
-            language = "English"
-        KG_EXTRACTION_PROMPT["FORMAT"]["language"] = language
-
-        hint_prompt = KG_EXTRACTION_PROMPT[language]["TEMPLATE"].format(
-            **KG_EXTRACTION_PROMPT["FORMAT"], input_text=content
-        )
-
-        final_result = await llm_client.generate_answer(hint_prompt)
-        logger.info("First result: %s", final_result)
-
-        history = pack_history_conversations(hint_prompt, final_result)
-        for loop_index in range(max_loop):
-            if_loop_result = await llm_client.generate_answer(
-                text=KG_EXTRACTION_PROMPT[language]["IF_LOOP"], history=history
-            )
-            if_loop_result = if_loop_result.strip().strip('"').strip("'").lower()
-            if if_loop_result != "yes":
-                break
-
-            glean_result = await llm_client.generate_answer(
-                text=KG_EXTRACTION_PROMPT[language]["CONTINUE"], history=history
-            )
-            logger.info("Loop %s glean: %s", loop_index, glean_result)
-
-            history += pack_history_conversations(
-                KG_EXTRACTION_PROMPT[language]["CONTINUE"], glean_result
-            )
-            final_result += glean_result
-            if loop_index == max_loop - 1:
-                break
-
-        records = split_string_by_multi_markers(
-            final_result,
-            [
-                KG_EXTRACTION_PROMPT["FORMAT"]["record_delimiter"],
-                KG_EXTRACTION_PROMPT["FORMAT"]["completion_delimiter"],
-            ],
-        )
-
-        nodes = defaultdict(list)
-        edges = defaultdict(list)
-
-        for record in records:
-            record = re.search(r"\((.*)\)", record)
-            if record is None:
-                continue
-            record = record.group(1)  # 提取括号内的内容
-            record_attributes = split_string_by_multi_markers(
-                record, [KG_EXTRACTION_PROMPT["FORMAT"]["tuple_delimiter"]]
-            )
-
-            entity = await handle_single_entity_extraction(record_attributes, chunk_id)
-            if entity is not None:
-                nodes[entity["entity_name"]].append(entity)
-                continue
-            relation = await handle_single_relationship_extraction(
-                record_attributes, chunk_id
-            )
-            if relation is not None:
-                edges[(relation["src_id"], relation["tgt_id"])].append(relation)
-        return dict(nodes), dict(edges)
-
-    results = await run_concurrent(
-        _process_single_content,
-        chunks,
-        desc="[2/4]Extracting entities and relationships from chunks",
-        unit="chunk",
-        progress_bar=progress_bar,
-    )
-
-    nodes = defaultdict(list)
-    edges = defaultdict(list)
-    for n, e in results:
-        for k, v in n.items():
-            nodes[k].extend(v)
-        for k, v in e.items():
-            edges[tuple(sorted(k))].extend(v)
-
-    await merge_nodes(nodes, kg_instance, llm_client, tokenizer_instance)
-    await merge_edges(edges, kg_instance, llm_client, tokenizer_instance)
-
-    return kg_instance

graphgen/operators/build_kg/merge_kg.py

@@ -1,212 +0,0 @@
-import asyncio
-from collections import Counter
-
-from tqdm.asyncio import tqdm as tqdm_async
-
-from graphgen.bases import BaseGraphStorage, BaseLLMClient
-from graphgen.models import Tokenizer
-from graphgen.templates import KG_EXTRACTION_PROMPT, KG_SUMMARIZATION_PROMPT
-from graphgen.utils import detect_main_language, logger
-from graphgen.utils.format import split_string_by_multi_markers
-
-
-async def _handle_kg_summary(
-    entity_or_relation_name: str,
-    description: str,
-    llm_client: BaseLLMClient,
-    tokenizer_instance: Tokenizer,
-    max_summary_tokens: int = 200,
-) -> str:
-    """
-    处理实体或关系的描述信息
-
-    :param entity_or_relation_name
-    :param description
-    :param llm_client
-    :param tokenizer_instance
-    :param max_summary_tokens
-    :return: new description
-    """
-    language = detect_main_language(description)
-    if language == "en":
-        language = "English"
-    else:
-        language = "Chinese"
-    KG_EXTRACTION_PROMPT["FORMAT"]["language"] = language
-
-    tokens = tokenizer_instance.encode(description)
-    if len(tokens) < max_summary_tokens:
-        return description
-
-    use_description = tokenizer_instance.decode(tokens[:max_summary_tokens])
-    prompt = KG_SUMMARIZATION_PROMPT[language]["TEMPLATE"].format(
-        entity_name=entity_or_relation_name,
-        description_list=use_description.split("<SEP>"),
-        **KG_SUMMARIZATION_PROMPT["FORMAT"],
-    )
-    new_description = await llm_client.generate_answer(prompt)
-    logger.info(
-        "Entity or relation %s summary: %s", entity_or_relation_name, new_description
-    )
-    return new_description
-
-
-async def merge_nodes(
-    nodes_data: dict,
-    kg_instance: BaseGraphStorage,
-    llm_client: BaseLLMClient,
-    tokenizer_instance: Tokenizer,
-    max_concurrent: int = 1000,
-):
-    """
-    Merge nodes
-
-    :param nodes_data
-    :param kg_instance
-    :param llm_client
-    :param tokenizer_instance
-    :param max_concurrent
-    :return
-    """
-
-    semaphore = asyncio.Semaphore(max_concurrent)
-
-    async def process_single_node(entity_name: str, node_data: list[dict]):
-        async with semaphore:
-            entity_types = []
-            source_ids = []
-            descriptions = []
-
-            node = await kg_instance.get_node(entity_name)
-            if node is not None:
-                entity_types.append(node["entity_type"])
-                source_ids.extend(
-                    split_string_by_multi_markers(node["source_id"], ["<SEP>"])
-                )
-                descriptions.append(node["description"])
-
-            # 统计当前节点数据和已有节点数据的entity_type出现次数，取出现次数最多的entity_type
-            entity_type = sorted(
-                Counter([dp["entity_type"] for dp in node_data] + entity_types).items(),
-                key=lambda x: x[1],
-                reverse=True,
-            )[0][0]
-
-            description = "<SEP>".join(
-                sorted(set([dp["description"] for dp in node_data] + descriptions))
-            )
-            description = await _handle_kg_summary(
-                entity_name, description, llm_client, tokenizer_instance
-            )
-
-            source_id = "<SEP>".join(
-                set([dp["source_id"] for dp in node_data] + source_ids)
-            )
-
-            node_data = {
-                "entity_type": entity_type,
-                "description": description,
-                "source_id": source_id,
-            }
-            await kg_instance.upsert_node(entity_name, node_data=node_data)
-            node_data["entity_name"] = entity_name
-            return node_data
-
-    logger.info("Inserting entities into storage...")
-    entities_data = []
-    for result in tqdm_async(
-        asyncio.as_completed(
-            [process_single_node(k, v) for k, v in nodes_data.items()]
-        ),
-        total=len(nodes_data),
-        desc="Inserting entities into storage",
-        unit="entity",
-    ):
-        try:
-            entities_data.append(await result)
-        except Exception as e:  # pylint: disable=broad-except
-            logger.error("Error occurred while inserting entities into storage: %s", e)
-
-
-async def merge_edges(
-    edges_data: dict,
-    kg_instance: BaseGraphStorage,
-    llm_client: BaseLLMClient,
-    tokenizer_instance: Tokenizer,
-    max_concurrent: int = 1000,
-):
-    """
-    Merge edges
-
-    :param edges_data
-    :param kg_instance
-    :param llm_client
-    :param tokenizer_instance
-    :param max_concurrent
-    :return
-    """
-
-    semaphore = asyncio.Semaphore(max_concurrent)
-
-    async def process_single_edge(src_id: str, tgt_id: str, edge_data: list[dict]):
-        async with semaphore:
-            source_ids = []
-            descriptions = []
-
-            edge = await kg_instance.get_edge(src_id, tgt_id)
-            if edge is not None:
-                source_ids.extend(
-                    split_string_by_multi_markers(edge["source_id"], ["<SEP>"])
-                )
-                descriptions.append(edge["description"])
-
-            description = "<SEP>".join(
-                sorted(set([dp["description"] for dp in edge_data] + descriptions))
-            )
-            source_id = "<SEP>".join(
-                set([dp["source_id"] for dp in edge_data] + source_ids)
-            )
-
-            for insert_id in [src_id, tgt_id]:
-                if not await kg_instance.has_node(insert_id):
-                    await kg_instance.upsert_node(
-                        insert_id,
-                        node_data={
-                            "source_id": source_id,
-                            "description": description,
-                            "entity_type": "UNKNOWN",
-                        },
-                    )
-
-            description = await _handle_kg_summary(
-                f"({src_id}, {tgt_id})", description, llm_client, tokenizer_instance
-            )
-
-            await kg_instance.upsert_edge(
-                src_id,
-                tgt_id,
-                edge_data={"source_id": source_id, "description": description},
-            )
-
-            edge_data = {"src_id": src_id, "tgt_id": tgt_id, "description": description}
-            return edge_data
-
-    logger.info("Inserting relationships into storage...")
-    relationships_data = []
-    for result in tqdm_async(
-        asyncio.as_completed(
-            [
-                process_single_edge(src_id, tgt_id, v)
-                for (src_id, tgt_id), v in edges_data.items()
-            ]
-        ),
-        total=len(edges_data),
-        desc="Inserting relationships into storage",
-        unit="relationship",
-    ):
-        try:
-            relationships_data.append(await result)
-        except Exception as e:  # pylint: disable=broad-except
-            logger.error(
-                "Error occurred while inserting relationships into storage: %s", e
-            )