Auto-sync from demo at Tue Sep 30 03:30:14 UTC 2025

app.py

@@ -9,7 +9,7 @@ import pandas as pd
 from dotenv import load_dotenv
 
 from graphgen.graphgen import GraphGen
-from graphgen.models import OpenAIModel, Tokenizer
+from graphgen.models import OpenAIClient, Tokenizer
 from graphgen.models.llm.limitter import RPM, TPM
 from graphgen.utils import set_logger
 from webui.base import WebuiParams
@@ -41,7 +41,7 @@ def init_graph_gen(config: dict, env: dict) -> GraphGen:
 
     graph_gen = GraphGen(working_dir=working_dir, config=config)
     # Set up LLM clients
-    graph_gen.synthesizer_llm_client = OpenAIModel(
+    graph_gen.synthesizer_llm_client = OpenAIClient(
         model_name=env.get("SYNTHESIZER_MODEL", ""),
         base_url=env.get("SYNTHESIZER_BASE_URL", ""),
         api_key=env.get("SYNTHESIZER_API_KEY", ""),
@@ -50,7 +50,7 @@ def init_graph_gen(config: dict, env: dict) -> GraphGen:
         tpm=TPM(env.get("TPM", 50000)),
     )
 
-    graph_gen.trainee_llm_client = OpenAIModel(
+    graph_gen.trainee_llm_client = OpenAIClient(
         model_name=env.get("TRAINEE_MODEL", ""),
         base_url=env.get("TRAINEE_BASE_URL", ""),
         api_key=env.get("TRAINEE_API_KEY", ""),

graphgen/bases/__init__.py

@@ -0,0 +1,12 @@
+from .base_kg_builder import BaseKGBuilder
+from .base_llm_client import BaseLLMClient
+from .base_reader import BaseReader
+from .base_splitter import BaseSplitter
+from .base_storage import (
+    BaseGraphStorage,
+    BaseKVStorage,
+    BaseListStorage,
+    StorageNameSpace,
+)
+from .base_tokenizer import BaseTokenizer
+from .datatypes import Chunk, QAPair, Token

graphgen/bases/base_kg_builder.py

@@ -0,0 +1,41 @@
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import Dict, List, Tuple
+
+from graphgen.bases.base_llm_client import BaseLLMClient
+from graphgen.bases.base_storage import BaseGraphStorage
+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))
+    _edges: Dict[Tuple[str, str], List[dict]] = field(
+        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
+    ) -> Tuple[Dict[str, List[dict]], Dict[Tuple[str, str], List[dict]]]:
+        """Extract nodes and edges from a single chunk."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def merge_nodes(
+        self, nodes_data: Dict[str, List[dict]], kg_instance: BaseGraphStorage, llm
+    ) -> None:
+        """Merge extracted nodes into the knowledge graph."""
+        raise NotImplementedError

graphgen/bases/base_llm_client.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import abc
+import re
+from typing import Any, List, Optional
+
+from graphgen.bases.base_tokenizer import BaseTokenizer
+from graphgen.bases.datatypes import Token
+
+
+class BaseLLMClient(abc.ABC):
+    """
+    LLM client base class, agnostic to specific backends (OpenAI / Ollama / ...).
+    """
+
+    def __init__(
+        self,
+        *,
+        system_prompt: str = "",
+        temperature: float = 0.0,
+        max_tokens: int = 4096,
+        repetition_penalty: float = 1.05,
+        top_p: float = 0.95,
+        top_k: int = 50,
+        tokenizer: Optional[BaseTokenizer] = None,
+        **kwargs: Any,
+    ):
+        self.system_prompt = system_prompt
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.repetition_penalty = repetition_penalty
+        self.top_p = top_p
+        self.top_k = top_k
+        self.tokenizer = tokenizer
+
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+    @abc.abstractmethod
+    async def generate_answer(
+        self, text: str, history: Optional[List[str]] = None, **extra: Any
+    ) -> str:
+        """Generate answer from the model."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def generate_topk_per_token(
+        self, text: str, history: Optional[List[str]] = None, **extra: Any
+    ) -> List[Token]:
+        """Generate top-k tokens for the next token prediction."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def generate_inputs_prob(
+        self, text: str, history: Optional[List[str]] = None, **extra: Any
+    ) -> List[Token]:
+        """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:
+        """
+        Remove <think> tags from the text.
+        If the text contains <think> and </think>, it removes everything between them and the tags themselves.
+        """
+        think_pattern = re.compile(rf"<{think_tag}>.*?</{think_tag}>", re.DOTALL)
+        filtered_text = think_pattern.sub("", text).strip()
+        return filtered_text if filtered_text else text.strip()

graphgen/bases/base_tokenizer.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import List
+
+
+@dataclass
+class BaseTokenizer(ABC):
+    model_name: str = "cl100k_base"
+
+    @abstractmethod
+    def encode(self, text: str) -> List[int]:
+        """Encode text -> token ids."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def decode(self, token_ids: List[int]) -> str:
+        """Decode token ids -> text."""
+        raise NotImplementedError
+
+    def count_tokens(self, text: str) -> int:
+        return len(self.encode(text))
+
+    def chunk_by_token_size(
+        self,
+        content: str,
+        *,
+        overlap_token_size: int = 128,
+        max_token_size: int = 1024,
+    ) -> List[dict]:
+        tokens = self.encode(content)
+        results = []
+        step = max_token_size - overlap_token_size
+        for index, start in enumerate(range(0, len(tokens), step)):
+            chunk_ids = tokens[start : start + max_token_size]
+            results.append(
+                {
+                    "tokens": len(chunk_ids),
+                    "content": self.decode(chunk_ids).strip(),
+                    "chunk_order_index": index,
+                }
+            )
+        return results

graphgen/bases/datatypes.py

@@ -1,4 +1,6 @@
+import math
 from dataclasses import dataclass, field
+from typing import List, Union
 
 
 @dataclass
@@ -16,3 +18,15 @@ class QAPair:
 
     question: str
     answer: str
+
+
+@dataclass
+class Token:
+    text: str
+    prob: float
+    top_candidates: List = field(default_factory=list)
+    ppl: Union[float, None] = field(default=None)
+
+    @property
+    def logprob(self) -> float:
+        return math.log(self.prob)

graphgen/graphgen.py

@@ -2,10 +2,9 @@ import asyncio
 import os
 import time
 from dataclasses import dataclass, field
-from typing import Dict, List, Union, cast
+from typing import Dict, cast
 
 import gradio as gr
-from tqdm.asyncio import tqdm as tqdm_async
 
 from graphgen.bases.base_storage import StorageNameSpace
 from graphgen.bases.datatypes import Chunk
@@ -13,27 +12,25 @@ from graphgen.models import (
     JsonKVStorage,
     JsonListStorage,
     NetworkXStorage,
-    OpenAIModel,
+    OpenAIClient,
     Tokenizer,
     TraverseStrategy,
-    read_file,
-    split_chunks,
 )
-
-from .operators import (
+from graphgen.operators import (
+    chunk_documents,
     extract_kg,
     generate_cot,
     judge_statement,
     quiz,
+    read_files,
     search_all,
     traverse_graph_for_aggregated,
     traverse_graph_for_atomic,
     traverse_graph_for_multi_hop,
 )
-from .utils import (
+from graphgen.utils import (
+    async_to_sync_method,
     compute_content_hash,
-    create_event_loop,
-    detect_main_language,
     format_generation_results,
     logger,
 )
@@ -49,8 +46,8 @@ class GraphGen:
 
     # llm
     tokenizer_instance: Tokenizer = None
-    synthesizer_llm_client: OpenAIModel = None
-    trainee_llm_client: OpenAIModel = None
+    synthesizer_llm_client: OpenAIClient = None
+    trainee_llm_client: OpenAIClient = None
 
     # search
     search_config: dict = field(
@@ -67,17 +64,17 @@ class GraphGen:
         self.tokenizer_instance: Tokenizer = Tokenizer(
             model_name=self.config["tokenizer"]
         )
-        self.synthesizer_llm_client: OpenAIModel = OpenAIModel(
+        self.synthesizer_llm_client: OpenAIClient = OpenAIClient(
             model_name=os.getenv("SYNTHESIZER_MODEL"),
             api_key=os.getenv("SYNTHESIZER_API_KEY"),
             base_url=os.getenv("SYNTHESIZER_BASE_URL"),
-            tokenizer_instance=self.tokenizer_instance,
+            tokenizer=self.tokenizer_instance,
         )
-        self.trainee_llm_client: OpenAIModel = OpenAIModel(
+        self.trainee_llm_client: OpenAIClient = OpenAIClient(
             model_name=os.getenv("TRAINEE_MODEL"),
             api_key=os.getenv("TRAINEE_API_KEY"),
             base_url=os.getenv("TRAINEE_BASE_URL"),
-            tokenizer_instance=self.tokenizer_instance,
+            tokenizer=self.tokenizer_instance,
         )
         self.search_config = self.config["search"]
 
@@ -111,15 +108,23 @@ class GraphGen:
             namespace="qa",
         )
 
-    async def async_split_chunks(self, data: List[Union[List, Dict]]) -> dict:
-        # TODO: configurable whether to use coreference resolution
+    @async_to_sync_method
+    async def insert(self):
+        """
+        insert chunks into the graph
+        """
+        input_file = self.config["read"]["input_file"]
+
+        # Step 1: Read files
+        data = read_files(input_file)
         if len(data) == 0:
-            return {}
+            logger.warning("No data to process")
+            return
 
-        inserting_chunks = {}
-        assert isinstance(data, list) and isinstance(data[0], dict)
+        # TODO: configurable whether to use coreference resolution
 
-        # compute hash for each document
+        # Step 2: Split chunks and filter existing ones
+        assert isinstance(data, list) and isinstance(data[0], dict)
         new_docs = {
             compute_content_hash(doc["content"], prefix="doc-"): {
                 "content": doc["content"]
@@ -128,38 +133,19 @@ class GraphGen:
         }
         _add_doc_keys = await self.full_docs_storage.filter_keys(list(new_docs.keys()))
         new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys}
+
         if len(new_docs) == 0:
             logger.warning("All docs are already in the storage")
-            return {}
+            return
         logger.info("[New Docs] inserting %d docs", len(new_docs))
 
-        cur_index = 1
-        doc_number = len(new_docs)
-        async for doc_key, doc in tqdm_async(
-            new_docs.items(), desc="[1/4]Chunking documents", unit="doc"
-        ):
-            doc_language = detect_main_language(doc["content"])
-            text_chunks = split_chunks(
-                doc["content"],
-                language=doc_language,
-                chunk_size=self.config["split"]["chunk_size"],
-                chunk_overlap=self.config["split"]["chunk_overlap"],
-            )
-
-            chunks = {
-                compute_content_hash(txt, prefix="chunk-"): {
-                    "content": txt,
-                    "full_doc_id": doc_key,
-                    "length": len(self.tokenizer_instance.encode_string(txt)),
-                    "language": doc_language,
-                }
-                for txt in text_chunks
-            }
-            inserting_chunks.update(chunks)
-
-            if self.progress_bar is not None:
-                self.progress_bar(cur_index / doc_number, f"Chunking {doc_key}")
-                cur_index += 1
+        inserting_chunks = await chunk_documents(
+            new_docs,
+            self.config["split"]["chunk_size"],
+            self.config["split"]["chunk_overlap"],
+            self.tokenizer_instance,
+            self.progress_bar,
+        )
 
         _add_chunk_keys = await self.text_chunks_storage.filter_keys(
             list(inserting_chunks.keys())
@@ -167,29 +153,16 @@ class GraphGen:
         inserting_chunks = {
             k: v for k, v in inserting_chunks.items() if k in _add_chunk_keys
         }
-        await self.full_docs_storage.upsert(new_docs)
-        await self.text_chunks_storage.upsert(inserting_chunks)
-
-        return inserting_chunks
-
-    def insert(self):
-        loop = create_event_loop()
-        loop.run_until_complete(self.async_insert())
-
-    async def async_insert(self):
-        """
-        insert chunks into the graph
-        """
-
-        input_file = self.config["read"]["input_file"]
-        data = read_file(input_file)
-        inserting_chunks = await self.async_split_chunks(data)
 
         if len(inserting_chunks) == 0:
             logger.warning("All chunks are already in the storage")
             return
+
         logger.info("[New Chunks] inserting %d chunks", len(inserting_chunks))
+        await self.full_docs_storage.upsert(new_docs)
+        await self.text_chunks_storage.upsert(inserting_chunks)
 
+        # Step 3: Extract entities and relations from chunks
         logger.info("[Entity and Relation Extraction]...")
         _add_entities_and_relations = await extract_kg(
             llm_client=self.synthesizer_llm_client,
@@ -219,11 +192,8 @@ class GraphGen:
             tasks.append(cast(StorageNameSpace, storage_instance).index_done_callback())
         await asyncio.gather(*tasks)
 
-    def search(self):
-        loop = create_event_loop()
-        loop.run_until_complete(self.async_search())
-
-    async def async_search(self):
+    @async_to_sync_method
+    async def search(self):
         logger.info(
             "Search is %s", "enabled" if self.search_config["enabled"] else "disabled"
         )
@@ -257,13 +227,10 @@ class GraphGen:
                         ]
                     )
                 # TODO: fix insert after search
-                await self.async_insert()
-
-    def quiz(self):
-        loop = create_event_loop()
-        loop.run_until_complete(self.async_quiz())
+                await self.insert()
 
-    async def async_quiz(self):
+    @async_to_sync_method
+    async def quiz(self):
         max_samples = self.config["quiz_and_judge_strategy"]["quiz_samples"]
         await quiz(
             self.synthesizer_llm_client,
@@ -273,11 +240,8 @@ class GraphGen:
         )
         await self.rephrase_storage.index_done_callback()
 
-    def judge(self):
-        loop = create_event_loop()
-        loop.run_until_complete(self.async_judge())
-
-    async def async_judge(self):
+    @async_to_sync_method
+    async def judge(self):
         re_judge = self.config["quiz_and_judge_strategy"]["re_judge"]
         _update_relations = await judge_statement(
             self.trainee_llm_client,
@@ -287,11 +251,8 @@ class GraphGen:
         )
         await _update_relations.index_done_callback()
 
-    def traverse(self):
-        loop = create_event_loop()
-        loop.run_until_complete(self.async_traverse())
-
-    async def async_traverse(self):
+    @async_to_sync_method
+    async def traverse(self):
         output_data_type = self.config["output_data_type"]
 
         if output_data_type == "atomic":
@@ -331,11 +292,8 @@ class GraphGen:
         await self.qa_storage.upsert(results)
         await self.qa_storage.index_done_callback()
 
-    def generate_reasoning(self, method_params):
-        loop = create_event_loop()
-        loop.run_until_complete(self.async_generate_reasoning(method_params))
-
-    async def async_generate_reasoning(self, method_params):
+    @async_to_sync_method
+    async def generate_reasoning(self, method_params):
         results = await generate_cot(
             self.graph_storage,
             self.synthesizer_llm_client,
@@ -349,11 +307,8 @@ class GraphGen:
         await self.qa_storage.upsert(results)
         await self.qa_storage.index_done_callback()
 
-    def clear(self):
-        loop = create_event_loop()
-        loop.run_until_complete(self.async_clear())
-
-    async def async_clear(self):
+    @async_to_sync_method
+    async def clear(self):
         await self.full_docs_storage.drop()
         await self.text_chunks_storage.drop()
         await self.search_storage.drop()

graphgen/models/__init__.py

@@ -3,15 +3,15 @@ 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 .llm.openai_model import OpenAIModel
-from .llm.tokenizer import Tokenizer
-from .llm.topk_token_model import Token, TopkTokenModel
-from .reader import read_file
+from .llm.openai_client import OpenAIClient
+from .llm.topk_token_model import TopkTokenModel
+from .reader import CsvReader, JsonlReader, JsonReader, TxtReader
 from .search.db.uniprot_search import UniProtSearch
 from .search.kg.wiki_search import WikiSearch
 from .search.web.bing_search import BingSearch
 from .search.web.google_search import GoogleSearch
-from .splitter import split_chunks
+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/evaluate/length_evaluator.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass
 
 from graphgen.bases.datatypes import QAPair
 from graphgen.models.evaluate.base_evaluator import BaseEvaluator
-from graphgen.models.llm.tokenizer import Tokenizer
+from graphgen.models.tokenizer import Tokenizer
 from graphgen.utils import create_event_loop
 
 
@@ -18,5 +18,5 @@ class LengthEvaluator(BaseEvaluator):
         return await loop.run_in_executor(None, self._calculate_length, pair.answer)
 
     def _calculate_length(self, text: str) -> float:
-        tokens = self.tokenizer.encode_string(text)
+        tokens = self.tokenizer.encode(text)
         return len(tokens)

graphgen/models/kg_builder/NetworkXKGBuilder.py

@@ -0,0 +1,18 @@
+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/llm/ollama_client.py

@@ -0,0 +1,21 @@
+# TODO: implement ollama client
+from typing import Any, List, Optional
+
+from graphgen.bases import BaseLLMClient, Token
+
+
+class OllamaClient(BaseLLMClient):
+    async def generate_answer(
+        self, text: str, history: Optional[List[str]] = None, **extra: Any
+    ) -> str:
+        pass
+
+    async def generate_topk_per_token(
+        self, text: str, history: Optional[List[str]] = None, **extra: Any
+    ) -> List[Token]:
+        pass
+
+    async def generate_inputs_prob(
+        self, text: str, history: Optional[List[str]] = None, **extra: Any
+    ) -> List[Token]:
+        pass

graphgen/models/llm/{openai_model.py → openai_client.py}

@@ -1,7 +1,5 @@
 import math
-import re
-from dataclasses import dataclass, field
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 import openai
 from openai import APIConnectionError, APITimeoutError, AsyncOpenAI, RateLimitError
@@ -12,9 +10,9 @@ from tenacity import (
     wait_exponential,
 )
 
+from graphgen.bases.base_llm_client import BaseLLMClient
+from graphgen.bases.datatypes import Token
 from graphgen.models.llm.limitter import RPM, TPM
-from graphgen.models.llm.tokenizer import Tokenizer
-from graphgen.models.llm.topk_token_model import Token, TopkTokenModel
 
 
 def get_top_response_tokens(response: openai.ChatCompletion) -> List[Token]:
@@ -30,32 +28,33 @@ def get_top_response_tokens(response: openai.ChatCompletion) -> List[Token]:
     return tokens
 
 
-def filter_think_tags(text: str) -> str:
-    """
-    Remove <think> tags from the text.
-    If the text contains <think> and </think>, it removes everything between them and the tags themselves.
-    """
-    think_pattern = re.compile(r"<think>.*?</think>", re.DOTALL)
-    filtered_text = think_pattern.sub("", text).strip()
-    return filtered_text if filtered_text else text.strip()
-
-
-@dataclass
-class OpenAIModel(TopkTokenModel):
-    model_name: str = "gpt-4o-mini"
-    api_key: str = None
-    base_url: str = None
-
-    system_prompt: str = ""
-    json_mode: bool = False
-    seed: int = None
-
-    token_usage: list = field(default_factory=list)
-    request_limit: bool = False
-    rpm: RPM = field(default_factory=lambda: RPM(rpm=1000))
-    tpm: TPM = field(default_factory=lambda: TPM(tpm=50000))
-
-    tokenizer_instance: Tokenizer = field(default_factory=Tokenizer)
+class OpenAIClient(BaseLLMClient):
+    def __init__(
+        self,
+        *,
+        model_name: str = "gpt-4o-mini",
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        json_mode: bool = False,
+        seed: Optional[int] = None,
+        topk_per_token: int = 5,  # number of topk tokens to generate for each token
+        request_limit: bool = False,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self.model_name = model_name
+        self.api_key = api_key
+        self.base_url = base_url
+        self.json_mode = json_mode
+        self.seed = seed
+        self.topk_per_token = topk_per_token
+
+        self.token_usage: list = []
+        self.request_limit = request_limit
+        self.rpm = RPM(rpm=1000)
+        self.tpm = TPM(tpm=50000)
+
+        self.__post_init__()
 
     def __post_init__(self):
         assert self.api_key is not None, "Please provide api key to access openai api."
@@ -66,7 +65,7 @@ class OpenAIModel(TopkTokenModel):
     def _pre_generate(self, text: str, history: List[str]) -> Dict:
         kwargs = {
             "temperature": self.temperature,
-            "top_p": self.topp,
+            "top_p": self.top_p,
             "max_tokens": self.max_tokens,
         }
         if self.seed:
@@ -94,7 +93,10 @@ class OpenAIModel(TopkTokenModel):
         ),
     )
     async def generate_topk_per_token(
-        self, text: str, history: Optional[List[str]] = None
+        self,
+        text: str,
+        history: Optional[List[str]] = None,
+        **extra: Any,
     ) -> List[Token]:
         kwargs = self._pre_generate(text, history)
         if self.topk_per_token > 0:
@@ -120,16 +122,16 @@ class OpenAIModel(TopkTokenModel):
         ),
     )
     async def generate_answer(
-        self, text: str, history: Optional[List[str]] = None, temperature: int = 0
+        self,
+        text: str,
+        history: Optional[List[str]] = None,
+        **extra: Any,
     ) -> str:
         kwargs = self._pre_generate(text, history)
-        kwargs["temperature"] = temperature
 
         prompt_tokens = 0
         for message in kwargs["messages"]:
-            prompt_tokens += len(
-                self.tokenizer_instance.encode_string(message["content"])
-            )
+            prompt_tokens += len(self.tokenizer.encode(message["content"]))
         estimated_tokens = prompt_tokens + kwargs["max_tokens"]
 
         if self.request_limit:
@@ -147,9 +149,10 @@ class OpenAIModel(TopkTokenModel):
                     "total_tokens": completion.usage.total_tokens,
                 }
             )
-        return filter_think_tags(completion.choices[0].message.content)
+        return self.filter_think_tags(completion.choices[0].message.content)
 
     async def generate_inputs_prob(
-        self, text: str, history: Optional[List[str]] = None
+        self, text: str, history: Optional[List[str]] = None, **extra: Any
     ) -> List[Token]:
+        """Generate probabilities for each token in the input."""
         raise NotImplementedError

graphgen/models/llm/tokenizer.py

@@ -1,73 +0,0 @@
-from dataclasses import dataclass
-from typing import List
-import tiktoken
-
-try:
-    from transformers import AutoTokenizer
-    TRANSFORMERS_AVAILABLE = True
-except ImportError:
-    AutoTokenizer = None
-    TRANSFORMERS_AVAILABLE = False
-
-
-def get_tokenizer(tokenizer_name: str = "cl100k_base"):
-    """
-    Get a tokenizer instance by name.
-
-    :param tokenizer_name: tokenizer name, tiktoken encoding name or Hugging Face model name
-    :return: tokenizer instance
-    """
-    if tokenizer_name in tiktoken.list_encoding_names():
-        return tiktoken.get_encoding(tokenizer_name)
-    if TRANSFORMERS_AVAILABLE:
-        try:
-            return AutoTokenizer.from_pretrained(tokenizer_name)
-        except Exception as e:
-            raise ValueError(f"Failed to load tokenizer from Hugging Face: {e}") from e
-    else:
-        raise ValueError("Hugging Face Transformers is not available, please install it first.")
-
-@dataclass
-class Tokenizer:
-    model_name: str = "cl100k_base"
-
-    def __post_init__(self):
-        self.tokenizer = get_tokenizer(self.model_name)
-
-    def encode_string(self, text: str) -> List[int]:
-        """
-        Encode text to tokens
-
-        :param text
-        :return: tokens
-        """
-        return self.tokenizer.encode(text)
-
-    def decode_tokens(self, tokens: List[int]) -> str:
-        """
-        Decode tokens to text
-
-        :param tokens
-        :return: text
-        """
-        return self.tokenizer.decode(tokens)
-
-    def chunk_by_token_size(
-        self, content: str, overlap_token_size=128, max_token_size=1024
-    ):
-        tokens = self.encode_string(content)
-        results = []
-        for index, start in enumerate(
-            range(0, len(tokens), max_token_size - overlap_token_size)
-        ):
-            chunk_content = self.decode_tokens(
-                tokens[start : start + max_token_size]
-            )
-            results.append(
-                {
-                    "tokens": min(max_token_size, len(tokens) - start),
-                    "content": chunk_content.strip(),
-                    "chunk_order_index": index,
-                }
-            )
-        return results

graphgen/models/llm/topk_token_model.py

@@ -1,18 +1,7 @@
-import math
-from dataclasses import dataclass, field
-from typing import List, Union, Optional
+from dataclasses import dataclass
+from typing import List, Optional
 
-
-@dataclass
-class Token:
-    text: str
-    prob: float
-    top_candidates: List = field(default_factory=list)
-    ppl: Union[float, None] = field(default=None)
-
-    @property
-    def logprob(self) -> float:
-        return math.log(self.prob)
+from graphgen.bases import Token
 
 
 @dataclass
@@ -34,14 +23,18 @@ class TopkTokenModel:
         """
         raise NotImplementedError
 
-    async def generate_inputs_prob(self, text: str, history: Optional[List[str]] = None) -> List[Token]:
+    async def generate_inputs_prob(
+        self, text: str, history: Optional[List[str]] = None
+    ) -> List[Token]:
         """
         Generate prob and text for each token of the input text.
         This function is used to visualize the ppl.
         """
         raise NotImplementedError
 
-    async def generate_answer(self, text: str, history: Optional[List[str]] = None) -> str:
+    async def generate_answer(
+        self, text: str, history: Optional[List[str]] = None
+    ) -> str:
         """
         Generate answer from the model.
         """

graphgen/models/reader/__init__.py

@@ -2,21 +2,3 @@ from .csv_reader import CsvReader
 from .json_reader import JsonReader
 from .jsonl_reader import JsonlReader
 from .txt_reader import TxtReader
-
-_MAPPING = {
-    "jsonl": JsonlReader,
-    "json": JsonReader,
-    "txt": TxtReader,
-    "csv": CsvReader,
-}
-
-
-def read_file(file_path: str):
-    suffix = file_path.split(".")[-1]
-    if suffix in _MAPPING:
-        reader = _MAPPING[suffix]()
-    else:
-        raise ValueError(
-            f"Unsupported file format: {suffix}. Supported formats are: {list(_MAPPING.keys())}"
-        )
-    return reader.read(file_path)

graphgen/models/splitter/__init__.py

@@ -1,31 +1,4 @@
-from functools import lru_cache
-from typing import Union
-
 from .recursive_character_splitter import (
     ChineseRecursiveTextSplitter,
     RecursiveCharacterSplitter,
 )
-
-_MAPPING = {
-    "en": RecursiveCharacterSplitter,
-    "zh": ChineseRecursiveTextSplitter,
-}
-
-SplitterT = Union[RecursiveCharacterSplitter, ChineseRecursiveTextSplitter]
-
-
-@lru_cache(maxsize=None)
-def _get_splitter(language: str, frozen_kwargs: frozenset) -> SplitterT:
-    cls = _MAPPING[language]
-    kwargs = dict(frozen_kwargs)
-    return cls(**kwargs)
-
-
-def split_chunks(text: str, language: str = "en", **kwargs) -> list:
-    if language not in _MAPPING:
-        raise ValueError(
-            f"Unsupported language: {language}. "
-            f"Supported languages are: {list(_MAPPING.keys())}"
-        )
-    splitter = _get_splitter(language, frozenset(kwargs.items()))
-    return splitter.split_text(text)

graphgen/models/tokenizer/__init__.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass, field
+from typing import List
+
+from graphgen.bases import BaseTokenizer
+
+from .hf_tokenizer import HFTokenizer
+from .tiktoken_tokenizer import TiktokenTokenizer
+
+try:
+    from transformers import AutoTokenizer
+
+    _HF_AVAILABLE = True
+except ImportError:
+    _HF_AVAILABLE = False
+
+
+def get_tokenizer_impl(tokenizer_name: str = "cl100k_base") -> BaseTokenizer:
+    import tiktoken
+
+    if tokenizer_name in tiktoken.list_encoding_names():
+        return TiktokenTokenizer(model_name=tokenizer_name)
+
+    # 2. HuggingFace
+    if _HF_AVAILABLE:
+        return HFTokenizer(model_name=tokenizer_name)
+
+    raise ValueError(
+        f"Unknown tokenizer {tokenizer_name} and HuggingFace not available."
+    )
+
+
+@dataclass
+class Tokenizer(BaseTokenizer):
+    """
+    Encapsulates different tokenization implementations based on the specified model name.
+    """
+
+    model_name: str = "cl100k_base"
+    _impl: BaseTokenizer = field(init=False, repr=False)
+
+    def __post_init__(self):
+        self._impl = get_tokenizer_impl(self.model_name)
+
+    def encode(self, text: str) -> List[int]:
+        return self._impl.encode(text)
+
+    def decode(self, token_ids: List[int]) -> str:
+        return self._impl.decode(token_ids)
+
+    def count_tokens(self, text: str) -> int:
+        return self._impl.count_tokens(text)

graphgen/models/tokenizer/hf_tokenizer.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+from typing import List
+
+from transformers import AutoTokenizer
+
+from graphgen.bases import BaseTokenizer
+
+
+@dataclass
+class HFTokenizer(BaseTokenizer):
+    def __post_init__(self):
+        self.enc = AutoTokenizer.from_pretrained(self.model_name)
+
+    def encode(self, text: str) -> List[int]:
+        return self.enc.encode(text, add_special_tokens=False)
+
+    def decode(self, token_ids: List[int]) -> str:
+        return self.enc.decode(token_ids, skip_special_tokens=True)

graphgen/models/tokenizer/tiktoken_tokenizer.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+from typing import List
+
+import tiktoken
+
+from graphgen.bases import BaseTokenizer
+
+
+@dataclass
+class TiktokenTokenizer(BaseTokenizer):
+    def __post_init__(self):
+        self.enc = tiktoken.get_encoding(self.model_name)
+
+    def encode(self, text: str) -> List[int]:
+        return self.enc.encode(text)
+
+    def decode(self, token_ids: List[int]) -> str:
+        return self.enc.decode(token_ids)

graphgen/operators/__init__.py

@@ -1,22 +1,13 @@
+from graphgen.operators.build_kg.extract_kg import extract_kg
 from graphgen.operators.generate.generate_cot import generate_cot
-from graphgen.operators.kg.extract_kg import extract_kg
 from graphgen.operators.search.search_all import search_all
 
 from .judge import judge_statement
 from .quiz import quiz
+from .read import read_files
+from .split import chunk_documents
 from .traverse_graph import (
     traverse_graph_for_aggregated,
     traverse_graph_for_atomic,
     traverse_graph_for_multi_hop,
 )
-
-__all__ = [
-    "extract_kg",
-    "quiz",
-    "judge_statement",
-    "search_all",
-    "traverse_graph_for_aggregated",
-    "traverse_graph_for_atomic",
-    "traverse_graph_for_multi_hop",
-    "generate_cot",
-]

graphgen/operators/build_kg/extract_kg.py

@@ -0,0 +1,127 @@
+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/{kg → build_kg}/merge_kg.py

@@ -3,8 +3,8 @@ from collections import Counter
 
 from tqdm.asyncio import tqdm as tqdm_async
 
-from graphgen.bases.base_storage import BaseGraphStorage
-from graphgen.models import Tokenizer, TopkTokenModel
+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
@@ -13,7 +13,7 @@ from graphgen.utils.format import split_string_by_multi_markers
 async def _handle_kg_summary(
     entity_or_relation_name: str,
     description: str,
-    llm_client: TopkTokenModel,
+    llm_client: BaseLLMClient,
     tokenizer_instance: Tokenizer,
     max_summary_tokens: int = 200,
 ) -> str:
@@ -34,11 +34,11 @@ async def _handle_kg_summary(
         language = "Chinese"
     KG_EXTRACTION_PROMPT["FORMAT"]["language"] = language
 
-    tokens = tokenizer_instance.encode_string(description)
+    tokens = tokenizer_instance.encode(description)
     if len(tokens) < max_summary_tokens:
         return description
 
-    use_description = tokenizer_instance.decode_tokens(tokens[:max_summary_tokens])
+    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>"),
@@ -54,7 +54,7 @@ async def _handle_kg_summary(
 async def merge_nodes(
     nodes_data: dict,
     kg_instance: BaseGraphStorage,
-    llm_client: TopkTokenModel,
+    llm_client: BaseLLMClient,
     tokenizer_instance: Tokenizer,
     max_concurrent: int = 1000,
 ):
@@ -131,7 +131,7 @@ async def merge_nodes(
 async def merge_edges(
     edges_data: dict,
     kg_instance: BaseGraphStorage,
-    llm_client: TopkTokenModel,
+    llm_client: BaseLLMClient,
     tokenizer_instance: Tokenizer,
     max_concurrent: int = 1000,
 ):

graphgen/operators/generate/generate_cot.py

@@ -3,14 +3,14 @@ from typing import Dict, List, Tuple
 
 from tqdm.asyncio import tqdm as tqdm_async
 
-from graphgen.models import CommunityDetector, NetworkXStorage, OpenAIModel
+from graphgen.models import CommunityDetector, NetworkXStorage, OpenAIClient
 from graphgen.templates import COT_GENERATION_PROMPT, COT_TEMPLATE_DESIGN_PROMPT
 from graphgen.utils import compute_content_hash, detect_main_language
 
 
 async def generate_cot(
     graph_storage: NetworkXStorage,
-    synthesizer_llm_client: OpenAIModel,
+    synthesizer_llm_client: OpenAIClient,
     method_params: Dict = None,
 ):
     method = method_params.get("method", "leiden")

graphgen/operators/judge.py

@@ -3,13 +3,13 @@ import math
 
 from tqdm.asyncio import tqdm as tqdm_async
 
-from graphgen.models import JsonKVStorage, NetworkXStorage, OpenAIModel
+from graphgen.models import JsonKVStorage, NetworkXStorage, OpenAIClient
 from graphgen.templates import STATEMENT_JUDGEMENT_PROMPT
 from graphgen.utils import logger, yes_no_loss_entropy
 
 
 async def judge_statement(  # pylint: disable=too-many-statements
-    trainee_llm_client: OpenAIModel,
+    trainee_llm_client: OpenAIClient,
     graph_storage: NetworkXStorage,
     rephrase_storage: JsonKVStorage,
     re_judge: bool = False,

graphgen/operators/kg/extract_kg.py

@@ -1,152 +0,0 @@
-import asyncio
-import re
-from collections import defaultdict
-from typing import List
-
-import gradio as gr
-from tqdm.asyncio import tqdm as tqdm_async
-
-from graphgen.bases.base_storage import BaseGraphStorage
-from graphgen.bases.datatypes import Chunk
-from graphgen.models import OpenAIModel, Tokenizer
-from graphgen.operators.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,
-    split_string_by_multi_markers,
-)
-
-
-# pylint: disable=too-many-statements
-async def extract_kg(
-    llm_client: OpenAIModel,
-    kg_instance: BaseGraphStorage,
-    tokenizer_instance: Tokenizer,
-    chunks: List[Chunk],
-    progress_bar: gr.Progress = None,
-    max_concurrent: int = 1000,
-):
-    """
-    :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
-    :param max_concurrent
-    :return:
-    """
-
-    semaphore = asyncio.Semaphore(max_concurrent)
-
-    async def _process_single_content(chunk: Chunk, max_loop: int = 3):
-        async with semaphore:
-            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 = []
-    chunk_number = len(chunks)
-    async for result in tqdm_async(
-        asyncio.as_completed([_process_single_content(c) for c in chunks]),
-        total=len(chunks),
-        desc="[2/4]Extracting entities and relationships from chunks",
-        unit="chunk",
-    ):
-        try:
-            if progress_bar is not None:
-                progress_bar(
-                    len(results) / chunk_number,
-                    desc="[3/4]Extracting entities and relationships from chunks",
-                )
-            results.append(await result)
-            if progress_bar is not None and len(results) == chunk_number:
-                progress_bar(
-                    1, desc="[3/4]Extracting entities and relationships from chunks"
-                )
-        except Exception as e:  # pylint: disable=broad-except
-            logger.error(
-                "Error occurred while extracting entities and relationships from chunks: %s",
-                e,
-            )
-
-    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/preprocess/resolute_coreference.py

@@ -1,13 +1,13 @@
 from typing import List
 
 from graphgen.bases.datatypes import Chunk
-from graphgen.models import OpenAIModel
+from graphgen.models import OpenAIClient
 from graphgen.templates import COREFERENCE_RESOLUTION_PROMPT
 from graphgen.utils import detect_main_language
 
 
 async def resolute_coreference(
-    llm_client: OpenAIModel, chunks: List[Chunk]
+    llm_client: OpenAIClient, chunks: List[Chunk]
 ) -> List[Chunk]:
     """
     Resolute conference

graphgen/operators/quiz.py

@@ -2,17 +2,19 @@ import asyncio
 from collections import defaultdict
 
 from tqdm.asyncio import tqdm as tqdm_async
-from graphgen.models import JsonKVStorage, OpenAIModel, NetworkXStorage
-from graphgen.utils import logger, detect_main_language
+
+from graphgen.models import JsonKVStorage, NetworkXStorage, OpenAIClient
 from graphgen.templates import DESCRIPTION_REPHRASING_PROMPT
+from graphgen.utils import detect_main_language, logger
 
 
 async def quiz(
-        synth_llm_client: OpenAIModel,
-        graph_storage: NetworkXStorage,
-        rephrase_storage: JsonKVStorage,
-        max_samples: int = 1,
-        max_concurrent: int = 1000) -> JsonKVStorage:
+    synth_llm_client: OpenAIClient,
+    graph_storage: NetworkXStorage,
+    rephrase_storage: JsonKVStorage,
+    max_samples: int = 1,
+    max_concurrent: int = 1000,
+) -> JsonKVStorage:
     """
     Get all edges and quiz them
 
@@ -26,11 +28,7 @@ async def quiz(
 
     semaphore = asyncio.Semaphore(max_concurrent)
 
-    async def _process_single_quiz(
-        des: str,
-        prompt: str,
-        gt: str
-    ):
+    async def _process_single_quiz(des: str, prompt: str, gt: str):
         async with semaphore:
             try:
                 # 如果在rephrase_storage中已经存在，直接取出
@@ -39,16 +37,14 @@ async def quiz(
                     return None
 
                 new_description = await synth_llm_client.generate_answer(
-                    prompt,
-                    temperature=1
+                    prompt, temperature=1
                 )
-                return  {des: [(new_description, gt)]}
+                return {des: [(new_description, gt)]}
 
-            except Exception as e: # pylint: disable=broad-except
+            except Exception as e:  # pylint: disable=broad-except
                 logger.error("Error when quizzing description %s: %s", des, e)
                 return None
 
-
     edges = await graph_storage.get_all_edges()
     nodes = await graph_storage.get_all_nodes()
 
@@ -60,41 +56,59 @@ async def quiz(
         description = edge_data["description"]
         language = "English" if detect_main_language(description) == "en" else "Chinese"
 
-        results[description] = [(description, 'yes')]
+        results[description] = [(description, "yes")]
 
         for i in range(max_samples):
             if i > 0:
                 tasks.append(
-                    _process_single_quiz(description,
-                                          DESCRIPTION_REPHRASING_PROMPT[language]['TEMPLATE'].format(
-                                              input_sentence=description), 'yes')
+                    _process_single_quiz(
+                        description,
+                        DESCRIPTION_REPHRASING_PROMPT[language]["TEMPLATE"].format(
+                            input_sentence=description
+                        ),
+                        "yes",
+                    )
                 )
-            tasks.append(_process_single_quiz(description,
-                                              DESCRIPTION_REPHRASING_PROMPT[language]['ANTI_TEMPLATE'].format(
-                                                  input_sentence=description), 'no'))
+            tasks.append(
+                _process_single_quiz(
+                    description,
+                    DESCRIPTION_REPHRASING_PROMPT[language]["ANTI_TEMPLATE"].format(
+                        input_sentence=description
+                    ),
+                    "no",
+                )
+            )
 
     for node in nodes:
         node_data = node[1]
         description = node_data["description"]
         language = "English" if detect_main_language(description) == "en" else "Chinese"
 
-        results[description] = [(description, 'yes')]
+        results[description] = [(description, "yes")]
 
         for i in range(max_samples):
             if i > 0:
                 tasks.append(
-                    _process_single_quiz(description,
-                                          DESCRIPTION_REPHRASING_PROMPT[language]['TEMPLATE'].format(
-                                              input_sentence=description), 'yes')
+                    _process_single_quiz(
+                        description,
+                        DESCRIPTION_REPHRASING_PROMPT[language]["TEMPLATE"].format(
+                            input_sentence=description
+                        ),
+                        "yes",
+                    )
+                )
+            tasks.append(
+                _process_single_quiz(
+                    description,
+                    DESCRIPTION_REPHRASING_PROMPT[language]["ANTI_TEMPLATE"].format(
+                        input_sentence=description
+                    ),
+                    "no",
                 )
-            tasks.append(_process_single_quiz(description,
-                                              DESCRIPTION_REPHRASING_PROMPT[language]['ANTI_TEMPLATE'].format(
-                                                  input_sentence=description), 'no'))
+            )
 
     for result in tqdm_async(
-            asyncio.as_completed(tasks),
-            total=len(tasks),
-            desc="Quizzing descriptions"
+        asyncio.as_completed(tasks), total=len(tasks), desc="Quizzing descriptions"
     ):
         new_result = await result
         if new_result:
@@ -105,5 +119,4 @@ async def quiz(
         results[key] = list(set(value))
         await rephrase_storage.upsert({key: results[key]})
 
-
     return rephrase_storage

graphgen/operators/read/__init__.py

@@ -0,0 +1 @@
+from .read_files import read_files

graphgen/operators/read/read_files.py

@@ -0,0 +1,19 @@
+from graphgen.models import CsvReader, JsonlReader, JsonReader, TxtReader
+
+_MAPPING = {
+    "jsonl": JsonlReader,
+    "json": JsonReader,
+    "txt": TxtReader,
+    "csv": CsvReader,
+}
+
+
+def read_files(file_path: str):
+    suffix = file_path.split(".")[-1]
+    if suffix in _MAPPING:
+        reader = _MAPPING[suffix]()
+    else:
+        raise ValueError(
+            f"Unsupported file format: {suffix}. Supported formats are: {list(_MAPPING.keys())}"
+        )
+    return reader.read(file_path)

graphgen/operators/split/__init__.py

@@ -0,0 +1 @@
+from .split_chunks import chunk_documents

graphgen/operators/split/split_chunks.py

@@ -0,0 +1,76 @@
+from functools import lru_cache
+from typing import Union
+
+from tqdm.asyncio import tqdm as tqdm_async
+
+from graphgen.models import (
+    ChineseRecursiveTextSplitter,
+    RecursiveCharacterSplitter,
+    Tokenizer,
+)
+from graphgen.utils import compute_content_hash, detect_main_language
+
+_MAPPING = {
+    "en": RecursiveCharacterSplitter,
+    "zh": ChineseRecursiveTextSplitter,
+}
+
+SplitterT = Union[RecursiveCharacterSplitter, ChineseRecursiveTextSplitter]
+
+
+@lru_cache(maxsize=None)
+def _get_splitter(language: str, frozen_kwargs: frozenset) -> SplitterT:
+    cls = _MAPPING[language]
+    kwargs = dict(frozen_kwargs)
+    return cls(**kwargs)
+
+
+def split_chunks(text: str, language: str = "en", **kwargs) -> list:
+    if language not in _MAPPING:
+        raise ValueError(
+            f"Unsupported language: {language}. "
+            f"Supported languages are: {list(_MAPPING.keys())}"
+        )
+    splitter = _get_splitter(language, frozenset(kwargs.items()))
+    return splitter.split_text(text)
+
+
+async def chunk_documents(
+    new_docs: dict,
+    chunk_size: int = 1024,
+    chunk_overlap: int = 100,
+    tokenizer_instance: Tokenizer = None,
+    progress_bar=None,
+) -> dict:
+    inserting_chunks = {}
+    cur_index = 1
+    doc_number = len(new_docs)
+    async for doc_key, doc in tqdm_async(
+        new_docs.items(), desc="[1/4]Chunking documents", unit="doc"
+    ):
+        doc_language = detect_main_language(doc["content"])
+        text_chunks = split_chunks(
+            doc["content"],
+            language=doc_language,
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+        )
+
+        chunks = {
+            compute_content_hash(txt, prefix="chunk-"): {
+                "content": txt,
+                "full_doc_id": doc_key,
+                "length": len(tokenizer_instance.encode(txt))
+                if tokenizer_instance
+                else len(txt),
+                "language": doc_language,
+            }
+            for txt in text_chunks
+        }
+        inserting_chunks.update(chunks)
+
+        if progress_bar is not None:
+            progress_bar(cur_index / doc_number, f"Chunking {doc_key}")
+            cur_index += 1
+
+    return inserting_chunks

graphgen/operators/traverse_graph.py

@@ -6,11 +6,11 @@ from tqdm.asyncio import tqdm as tqdm_async
 from graphgen.models import (
     JsonKVStorage,
     NetworkXStorage,
-    OpenAIModel,
+    OpenAIClient,
     Tokenizer,
     TraverseStrategy,
 )
-from graphgen.operators.kg.split_kg import get_batches_with_strategy
+from graphgen.operators.build_kg.split_kg import get_batches_with_strategy
 from graphgen.templates import (
     ANSWER_REPHRASING_PROMPT,
     MULTI_HOP_GENERATION_PROMPT,
@@ -30,7 +30,7 @@ async def _pre_tokenize(
             if "length" not in edge[2]:
                 edge[2]["length"] = len(
                     await asyncio.get_event_loop().run_in_executor(
-                        None, tokenizer.encode_string, edge[2]["description"]
+                        None, tokenizer.encode, edge[2]["description"]
                     )
                 )
             return edge
@@ -40,7 +40,7 @@ async def _pre_tokenize(
             if "length" not in node[1]:
                 node[1]["length"] = len(
                     await asyncio.get_event_loop().run_in_executor(
-                        None, tokenizer.encode_string, node[1]["description"]
+                        None, tokenizer.encode, node[1]["description"]
                     )
                 )
             return node
@@ -161,7 +161,7 @@ def _post_process_synthetic_data(data):
 
 
 async def traverse_graph_for_aggregated(
-    llm_client: OpenAIModel,
+    llm_client: OpenAIClient,
     tokenizer: Tokenizer,
     graph_storage: NetworkXStorage,
     traverse_strategy: TraverseStrategy,
@@ -310,7 +310,7 @@ async def traverse_graph_for_aggregated(
 
 # pylint: disable=too-many-branches, too-many-statements
 async def traverse_graph_for_atomic(
-    llm_client: OpenAIModel,
+    llm_client: OpenAIClient,
     tokenizer: Tokenizer,
     graph_storage: NetworkXStorage,
     traverse_strategy: TraverseStrategy,
@@ -426,7 +426,7 @@ async def traverse_graph_for_atomic(
 
 
 async def traverse_graph_for_multi_hop(
-    llm_client: OpenAIModel,
+    llm_client: OpenAIClient,
     tokenizer: Tokenizer,
     graph_storage: NetworkXStorage,
     traverse_strategy: TraverseStrategy,

graphgen/utils/__init__.py

@@ -13,3 +13,5 @@ from .hash import compute_args_hash, compute_content_hash
 from .help_nltk import NLTKHelper
 from .log import logger, parse_log, set_logger
 from .loop import create_event_loop
+from .run_concurrent import run_concurrent
+from .wrap import async_to_sync_method

graphgen/utils/calculate_confidence.py

@@ -1,34 +1,41 @@
 import math
 from typing import List
-from graphgen.models.llm.topk_token_model import Token
+
+from graphgen.bases.datatypes import Token
+
 
 def preprocess_tokens(tokens: List[Token]) -> List[Token]:
     """Preprocess tokens for calculating confidence."""
     tokens = [x for x in tokens if x.prob > 0]
     return tokens
 
+
 def joint_probability(tokens: List[Token]) -> float:
     """Calculate joint probability of a list of tokens."""
     tokens = preprocess_tokens(tokens)
     logprob_sum = sum(x.logprob for x in tokens)
     return math.exp(logprob_sum / len(tokens))
 
+
 def min_prob(tokens: List[Token]) -> float:
     """Calculate the minimum probability of a list of tokens."""
     tokens = preprocess_tokens(tokens)
     return min(x.prob for x in tokens)
 
+
 def average_prob(tokens: List[Token]) -> float:
     """Calculate the average probability of a list of tokens."""
     tokens = preprocess_tokens(tokens)
     return sum(x.prob for x in tokens) / len(tokens)
 
+
 def average_confidence(tokens: List[Token]) -> float:
     """Calculate the average confidence of a list of tokens."""
     tokens = preprocess_tokens(tokens)
     confidence = [x.prob / sum(y.prob for y in x.top_candidates[:5]) for x in tokens]
     return sum(confidence) / len(tokens)
 
+
 def yes_no_loss(tokens_list: List[List[Token]], ground_truth: List[str]) -> float:
     """Calculate the loss for yes/no question."""
     losses = []
@@ -41,7 +48,10 @@ def yes_no_loss(tokens_list: List[List[Token]], ground_truth: List[str]) -> floa
             losses.append(token.prob)
     return sum(losses) / len(losses)
 
-def yes_no_loss_entropy(tokens_list: List[List[Token]], ground_truth: List[str]) -> float:
+
+def yes_no_loss_entropy(
+    tokens_list: List[List[Token]], ground_truth: List[str]
+) -> float:
     """Calculate the loss for yes/no question using entropy."""
     losses = []
     for i, tokens in enumerate(tokens_list):

graphgen/utils/run_concurrent.py

@@ -0,0 +1,38 @@
+import asyncio
+from typing import Awaitable, Callable, List, Optional, TypeVar
+
+import gradio as gr
+from tqdm.asyncio import tqdm as tqdm_async
+
+from graphgen.utils.log import logger
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+async def run_concurrent(
+    coro_fn: Callable[[T], Awaitable[R]],
+    items: List[T],
+    *,
+    desc: str = "processing",
+    unit: str = "item",
+    progress_bar: Optional[gr.Progress] = None,
+) -> List[R]:
+    tasks = [asyncio.create_task(coro_fn(it)) for it in items]
+
+    results = await tqdm_async.gather(*tasks, desc=desc, unit=unit)
+
+    ok_results = []
+    for idx, res in enumerate(results):
+        if isinstance(res, Exception):
+            logger.exception("Task failed: %s", res)
+            if progress_bar:
+                progress_bar((idx + 1) / len(items), desc=desc)
+            continue
+        ok_results.append(res)
+        if progress_bar:
+            progress_bar((idx + 1) / len(items), desc=desc)
+
+    if progress_bar:
+        progress_bar(1.0, desc=desc)
+    return ok_results

graphgen/utils/wrap.py

@@ -0,0 +1,13 @@
+from functools import wraps
+from typing import Any, Callable
+
+from .loop import create_event_loop
+
+
+def async_to_sync_method(func: Callable) -> Callable:
+    @wraps(func)
+    def wrapper(self, *args, **kwargs) -> Any:
+        loop = create_event_loop()
+        return loop.run_until_complete(func(self, *args, **kwargs))
+
+    return wrapper

webui/app.py

@@ -9,7 +9,7 @@ import pandas as pd
 from dotenv import load_dotenv
 
 from graphgen.graphgen import GraphGen
-from graphgen.models import OpenAIModel, Tokenizer
+from graphgen.models import OpenAIClient, Tokenizer
 from graphgen.models.llm.limitter import RPM, TPM
 from graphgen.utils import set_logger
 from webui.base import WebuiParams
@@ -41,7 +41,7 @@ def init_graph_gen(config: dict, env: dict) -> GraphGen:
 
     graph_gen = GraphGen(working_dir=working_dir, config=config)
     # Set up LLM clients
-    graph_gen.synthesizer_llm_client = OpenAIModel(
+    graph_gen.synthesizer_llm_client = OpenAIClient(
         model_name=env.get("SYNTHESIZER_MODEL", ""),
         base_url=env.get("SYNTHESIZER_BASE_URL", ""),
         api_key=env.get("SYNTHESIZER_API_KEY", ""),
@@ -50,7 +50,7 @@ def init_graph_gen(config: dict, env: dict) -> GraphGen:
         tpm=TPM(env.get("TPM", 50000)),
     )
 
-    graph_gen.trainee_llm_client = OpenAIModel(
+    graph_gen.trainee_llm_client = OpenAIClient(
         model_name=env.get("TRAINEE_MODEL", ""),
         base_url=env.get("TRAINEE_BASE_URL", ""),
         api_key=env.get("TRAINEE_API_KEY", ""),

webui/utils/count_tokens.py

@@ -45,7 +45,7 @@ def count_tokens(file, tokenizer_name, data_frame):
             content = item.get("content", "")
         else:
             content = item
-        token_count += len(tokenizer.encode_string(content))
+        token_count += len(tokenizer.encode(content))
 
     _update_data = [[str(token_count), str(token_count * 50), "N/A"]]