Auto-sync from demo at Tue Sep 30 07:59:12 UTC 2025

app.py

@@ -39,27 +39,32 @@ def init_graph_gen(config: dict, env: dict) -> GraphGen:
     set_logger(log_file, if_stream=True)
     os.environ.update({k: str(v) for k, v in env.items()})
 
-    graph_gen = GraphGen(working_dir=working_dir, config=config)
-    # Set up LLM clients
-    graph_gen.synthesizer_llm_client = OpenAIClient(
+    tokenizer_instance = Tokenizer(config.get("tokenizer", "cl100k_base"))
+    synthesizer_llm_client = OpenAIClient(
         model_name=env.get("SYNTHESIZER_MODEL", ""),
         base_url=env.get("SYNTHESIZER_BASE_URL", ""),
         api_key=env.get("SYNTHESIZER_API_KEY", ""),
         request_limit=True,
         rpm=RPM(env.get("RPM", 1000)),
         tpm=TPM(env.get("TPM", 50000)),
+        tokenizer=tokenizer_instance,
     )
-
-    graph_gen.trainee_llm_client = OpenAIClient(
+    trainee_llm_client = OpenAIClient(
         model_name=env.get("TRAINEE_MODEL", ""),
         base_url=env.get("TRAINEE_BASE_URL", ""),
         api_key=env.get("TRAINEE_API_KEY", ""),
         request_limit=True,
         rpm=RPM(env.get("RPM", 1000)),
         tpm=TPM(env.get("TPM", 50000)),
+        tokenizer=tokenizer_instance,
     )
 
-    graph_gen.tokenizer_instance = Tokenizer(config.get("tokenizer", "cl100k_base"))
+    graph_gen = GraphGen(
+        working_dir=working_dir,
+        tokenizer_instance=tokenizer_instance,
+        synthesizer_llm_client=synthesizer_llm_client,
+        trainee_llm_client=trainee_llm_client,
+    )
 
     return graph_gen
 
@@ -78,27 +83,32 @@ def run_graphgen(params: WebuiParams, progress=gr.Progress()):
             "chunk_size": params.chunk_size,
             "chunk_overlap": params.chunk_overlap,
         },
-        "output_data_type": params.output_data_type,
-        "output_data_format": params.output_data_format,
-        "tokenizer": params.tokenizer,
         "search": {"enabled": False},
-        "quiz_and_judge_strategy": {
+        "quiz_and_judge": {
             "enabled": params.if_trainee_model,
             "quiz_samples": params.quiz_samples,
         },
-        "traverse_strategy": {
-            "bidirectional": params.bidirectional,
-            "expand_method": params.expand_method,
-            "max_extra_edges": params.max_extra_edges,
-            "max_tokens": params.max_tokens,
-            "max_depth": params.max_depth,
-            "edge_sampling": params.edge_sampling,
-            "isolated_node_strategy": params.isolated_node_strategy,
-            "loss_strategy": params.loss_strategy,
+        "partition": {
+            "method": "ece",
+            "method_params": {
+                "bidirectional": params.bidirectional,
+                "expand_method": params.expand_method,
+                "max_extra_edges": params.max_extra_edges,
+                "max_tokens": params.max_tokens,
+                "max_depth": params.max_depth,
+                "edge_sampling": params.edge_sampling,
+                "isolated_node_strategy": params.isolated_node_strategy,
+                "loss_strategy": params.loss_strategy,
+            },
+        },
+        "generate": {
+            "mode": params.output_data_type,
+            "data_format": params.output_data_format,
         },
     }
 
     env = {
+        "TOKENIZER_MODEL": params.tokenizer,
         "SYNTHESIZER_BASE_URL": params.synthesizer_url,
         "SYNTHESIZER_MODEL": params.synthesizer_model,
         "TRAINEE_BASE_URL": params.trainee_url,
@@ -128,19 +138,18 @@ def run_graphgen(params: WebuiParams, progress=gr.Progress()):
 
     try:
         # Process the data
-        graph_gen.insert()
+        graph_gen.insert(read_config=config["read"], split_config=config["split"])
 
         if config["if_trainee_model"]:
-            # Generate quiz
-            graph_gen.quiz()
-
-            # Judge statements
-            graph_gen.judge()
+            # Quiz and Judge
+            graph_gen.quiz_and_judge(quiz_and_judge_config=config["quiz_and_judge"])
         else:
-            graph_gen.traverse_strategy.edge_sampling = "random"
+            config["partition"]["method_params"]["edge_sampling"] = "random"
 
-        # Traverse graph
-        graph_gen.traverse()
+        graph_gen.generate(
+            partition_config=config["partition"],
+            generate_config=config["generate"],
+        )
 
         # Save output
         output_data = graph_gen.qa_storage.data
@@ -249,6 +258,9 @@ with gr.Blocks(title="GraphGen Demo", theme=gr.themes.Glass(), css=css) as demo:
         )
 
         with gr.Accordion(label=_("Model Config"), open=False):
+            tokenizer = gr.Textbox(
+                label="Tokenizer", value="cl100k_base", interactive=True
+            )
             synthesizer_url = gr.Textbox(
                 label="Synthesizer URL",
                 value="https://api.siliconflow.cn/v1",
@@ -300,9 +312,6 @@ with gr.Blocks(title="GraphGen Demo", theme=gr.themes.Glass(), css=css) as demo:
                 step=100,
                 interactive=True,
             )
-            tokenizer = gr.Textbox(
-                label="Tokenizer", value="cl100k_base", interactive=True
-            )
             output_data_type = gr.Radio(
                 choices=["atomic", "multi_hop", "aggregated"],
                 label="Output Data Type",

graphgen/configs/aggregated_config.yaml

@@ -6,19 +6,21 @@ split:
 search: # web search configuration
   enabled: false # whether to enable web search
   search_types: ["google"] # search engine types, support: google, bing, uniprot, wikipedia
-output_data_type: aggregated # atomic, aggregated, multi_hop, cot
-output_data_format: ChatML # Alpaca, Sharegpt, ChatML
-tokenizer: cl100k_base # tokenizer for counting tokens, support tiktoken tokenizer names and local tokenizer path
-quiz_and_judge_strategy: # quiz and test whether the LLM masters the knowledge points
+quiz_and_judge: # quiz and test whether the LLM masters the knowledge points
   enabled: true
   quiz_samples: 2 # number of quiz samples to generate
   re_judge: false # whether to re-judge the existing quiz samples
-traverse_strategy: # strategy for clustering sub-graphs using comprehension loss
-  bidirectional: true # whether to traverse the graph in both directions
-  edge_sampling: max_loss # edge sampling strategy, support: random, max_loss, min_loss
-  expand_method: max_width # expand method, support: max_width, max_depth
-  isolated_node_strategy: ignore # strategy for isolated nodes, support: ignore, add
-  max_depth: 5 # maximum depth for graph traversal
-  max_extra_edges: 20 # max edges per direction (if expand_method="max_width")
-  max_tokens: 256 # restricts input length (if expand_method="max_tokens")
-  loss_strategy: only_edge # defines loss computation focus, support: only_edge, both
+partition: # graph partition configuration
+  method: ece # ece is a custom partition method based on comprehension loss
+  method_params:
+    bidirectional: true # whether to traverse the graph in both directions
+    edge_sampling: max_loss # edge sampling strategy, support: random, max_loss, min_loss
+    expand_method: max_width # expand method, support: max_width, max_depth
+    isolated_node_strategy: ignore # strategy for isolated nodes, support: ignore, add
+    max_depth: 5 # maximum depth for graph traversal
+    max_extra_edges: 20 # max edges per direction (if expand_method="max_width")
+    max_tokens: 256 # restricts input length (if expand_method="max_tokens")
+    loss_strategy: only_edge # defines loss computation focus, support: only_edge, both
+generate:
+  mode: aggregated # atomic, aggregated, multi_hop, cot
+  data_format: ChatML # Alpaca, Sharegpt, ChatML

graphgen/configs/atomic_config.yaml

@@ -6,19 +6,21 @@ split:
 search: # web search configuration
   enabled: false # whether to enable web search
   search_types: ["google"] # search engine types, support: google, bing, uniprot, wikipedia
-output_data_type: atomic # atomic, aggregated, multi_hop, cot
-output_data_format: Alpaca # Alpaca, Sharegpt, ChatML
-tokenizer: cl100k_base # tokenizer for counting tokens, support tiktoken tokenizer names and local tokenizer path
-quiz_and_judge_strategy: # quiz and test whether the LLM masters the knowledge points
+quiz_and_judge: # quiz and test whether the LLM masters the knowledge points
   enabled: true
   quiz_samples: 2 # number of quiz samples to generate
   re_judge: false # whether to re-judge the existing quiz samples
-traverse_strategy: # strategy for clustering sub-graphs using comprehension loss
-  bidirectional: true # whether to traverse the graph in both directions
-  edge_sampling: max_loss # edge sampling strategy, support: random, max_loss, min_loss
-  expand_method: max_width # expand method, support: max_width, max_depth
-  isolated_node_strategy: ignore # strategy for isolated nodes, support: ignore, add
-  max_depth: 3 # maximum depth for graph traversal
-  max_extra_edges: 5 # max edges per direction (if expand_method="max_width")
-  max_tokens: 256 # restricts input length (if expand_method="max_tokens")
-  loss_strategy: only_edge # defines loss computation focus, support: only_edge, both
+partition: # graph partition configuration
+  method: ece # ece is a custom partition method based on comprehension loss
+  method_params:
+    bidirectional: true # whether to traverse the graph in both directions
+    edge_sampling: max_loss # edge sampling strategy, support: random, max_loss, min_loss
+    expand_method: max_width # expand method, support: max_width, max_depth
+    isolated_node_strategy: ignore # strategy for isolated nodes, support: ignore, add
+    max_depth: 3 # maximum depth for graph traversal
+    max_extra_edges: 5 # max edges per direction (if expand_method="max_width")
+    max_tokens: 256 # restricts input length (if expand_method="max_tokens")
+    loss_strategy: only_edge # defines loss computation focus, support: only_edge, both
+generate:
+  mode: atomic # atomic, aggregated, multi_hop, cot
+  data_format: Alpaca # Alpaca, Sharegpt, ChatML

graphgen/configs/cot_config.yaml

@@ -6,11 +6,14 @@ split:
 search: # web search configuration
   enabled: false # whether to enable web search
   search_types: ["google"] # search engine types, support: google, bing, uniprot, wikipedia
-output_data_type: cot # atomic, aggregated, multi_hop, cot
-output_data_format: Sharegpt # Alpaca, Sharegpt, ChatML
-tokenizer: cl100k_base # tokenizer for counting tokens, support tiktoken tokenizer names and local tokenizer path
-method_params:
-  method: leiden
-  max_size: 20 # Maximum size of communities
-  use_lcc: false
-  random_seed: 42
+quiz_and_judge: # quiz and test whether the LLM masters the knowledge points
+  enabled: false
+partition: # graph partition configuration
+  method: leiden # leiden is a community detection algorithm
+  method_params:
+    max_size: 20 # Maximum size of communities
+    use_lcc: false
+    random_seed: 42
+generate:
+  mode: cot # atomic, aggregated, multi_hop, cot
+  data_format: Sharegpt # Alpaca, Sharegpt, ChatML

graphgen/configs/multi_hop_config.yaml

@@ -6,19 +6,21 @@ split:
 search: # web search configuration
   enabled: false # whether to enable web search
   search_types: ["google"] # search engine types, support: google, bing, uniprot, wikipedia
-output_data_type: multi_hop # atomic, aggregated, multi_hop, cot
-output_data_format: ChatML # Alpaca, Sharegpt, ChatML
-tokenizer: cl100k_base # tokenizer for counting tokens, support tiktoken tokenizer names and local tokenizer path
-quiz_and_judge_strategy: # quiz and test whether the LLM masters the knowledge points
+quiz_and_judge: # quiz and test whether the LLM masters the knowledge points
   enabled: false
   quiz_samples: 2 # number of quiz samples to generate
   re_judge: false # whether to re-judge the existing quiz samples
-traverse_strategy: # strategy for clustering sub-graphs using comprehension loss
-  bidirectional: true # whether to traverse the graph in both directions
-  edge_sampling: max_loss # edge sampling strategy, support: random, max_loss, min_loss
-  expand_method: max_width # expand method, support: max_width, max_depth
-  isolated_node_strategy: ignore # strategy for isolated nodes, support: ignore, add
-  max_depth: 1 # maximum depth for graph traversal
-  max_extra_edges: 2 # max edges per direction (if expand_method="max_width")
-  max_tokens: 256 # restricts input length (if expand_method="max_tokens")
-  loss_strategy: only_edge # defines loss computation focus, support: only_edge, both
+partition: # graph partition configuration
+  method: ece # ece is a custom partition method based on comprehension loss
+  method_params:
+    bidirectional: true # whether to traverse the graph in both directions
+    edge_sampling: max_loss # edge sampling strategy, support: random, max_loss, min_loss
+    expand_method: max_width # expand method, support: max_width, max_depth
+    isolated_node_strategy: ignore # strategy for isolated nodes, support: ignore, add
+    max_depth: 1 # maximum depth for graph traversal
+    max_extra_edges: 2 # max edges per direction (if expand_method="max_width")
+    max_tokens: 256 # restricts input length (if expand_method="max_tokens")
+    loss_strategy: only_edge # defines loss computation focus, support: only_edge, both
+generate:
+  mode: multi_hop # strategy for generating multi-hop QA pairs
+  data_format: ChatML # Alpaca, Sharegpt, ChatML

graphgen/generate.py

@@ -6,8 +6,8 @@ from importlib.resources import files
 import yaml
 from dotenv import load_dotenv
 
-from .graphgen import GraphGen
-from .utils import logger, set_logger
+from graphgen.graphgen import GraphGen
+from graphgen.utils import logger, set_logger
 
 sys_path = os.path.abspath(os.path.dirname(__file__))
 
@@ -50,50 +50,51 @@ def main():
     with open(args.config_file, "r", encoding="utf-8") as f:
         config = yaml.load(f, Loader=yaml.FullLoader)
 
-    output_data_type = config["output_data_type"]
+    mode = config["generate"]["mode"]
     unique_id = int(time.time())
 
-    output_path = os.path.join(
-        working_dir, "data", "graphgen", f"{unique_id}_{output_data_type}"
-    )
+    output_path = os.path.join(working_dir, "data", "graphgen", f"{unique_id}")
     set_working_dir(output_path)
 
     set_logger(
-        os.path.join(output_path, f"{unique_id}.log"),
+        os.path.join(output_path, f"{unique_id}_{mode}.log"),
         if_stream=True,
     )
     logger.info(
         "GraphGen with unique ID %s logging to %s",
         unique_id,
-        os.path.join(
-            working_dir, "logs", f"{unique_id}_graphgen_{output_data_type}.log"
-        ),
+        os.path.join(working_dir, f"{unique_id}_{mode}.log"),
     )
 
-    graph_gen = GraphGen(working_dir=working_dir, unique_id=unique_id, config=config)
+    graph_gen = GraphGen(unique_id=unique_id, working_dir=working_dir)
 
-    graph_gen.insert()
+    graph_gen.insert(read_config=config["read"], split_config=config["split"])
 
-    if config["search"]["enabled"]:
-        graph_gen.search()
+    graph_gen.search(search_config=config["search"])
 
     # Use pipeline according to the output data type
-    if output_data_type in ["atomic", "aggregated", "multi_hop"]:
-        if "quiz_and_judge_strategy" in config and config[
-            "quiz_and_judge_strategy"
-        ].get("enabled", False):
-            graph_gen.quiz()
-            graph_gen.judge()
+    if mode in ["atomic", "aggregated", "multi_hop"]:
+        logger.info("Generation mode set to '%s'. Start generation.", mode)
+        if "quiz_and_judge" in config and config["quiz_and_judge"]["enabled"]:
+            graph_gen.quiz_and_judge(quiz_and_judge_config=config["quiz_and_judge"])
         else:
             logger.warning(
                 "Quiz and Judge strategy is disabled. Edge sampling falls back to random."
             )
-            graph_gen.traverse_strategy.edge_sampling = "random"
-        graph_gen.traverse()
-    elif output_data_type == "cot":
-        graph_gen.generate_reasoning(method_params=config["method_params"])
+            assert (
+                config["partition"]["method"] == "ece"
+                and "ece_params" in config["partition"]
+            ), "Only ECE partition with edge sampling is supported."
+            config["partition"]["method_params"]["edge_sampling"] = "random"
+    elif mode == "cot":
+        logger.info("Generation mode set to 'cot'. Start generation.")
     else:
-        raise ValueError(f"Unsupported output data type: {output_data_type}")
+        raise ValueError(f"Unsupported output data type: {mode}")
+
+    graph_gen.generate(
+        partition_config=config["partition"],
+        generate_config=config["generate"],
+    )
 
     save_config(os.path.join(output_path, "config.yaml"), config)
     logger.info("GraphGen completed successfully. Data saved to %s", output_path)

graphgen/graphgen.py

@@ -1,7 +1,7 @@
 import asyncio
 import os
 import time
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Dict, cast
 
 import gradio as gr
@@ -14,7 +14,6 @@ from graphgen.models import (
     NetworkXStorage,
     OpenAIClient,
     Tokenizer,
-    TraverseStrategy,
 )
 from graphgen.operators import (
     chunk_documents,
@@ -42,46 +41,36 @@ sys_path = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
 class GraphGen:
     unique_id: int = int(time.time())
     working_dir: str = os.path.join(sys_path, "cache")
-    config: Dict = field(default_factory=dict)
 
     # llm
     tokenizer_instance: Tokenizer = None
     synthesizer_llm_client: OpenAIClient = None
     trainee_llm_client: OpenAIClient = None
 
-    # search
-    search_config: dict = field(
-        default_factory=lambda: {"enabled": False, "search_types": ["wikipedia"]}
-    )
-
-    # traversal
-    traverse_strategy: TraverseStrategy = None
-
     # webui
     progress_bar: gr.Progress = None
 
     def __post_init__(self):
-        self.tokenizer_instance: Tokenizer = Tokenizer(
-            model_name=self.config["tokenizer"]
+        self.tokenizer_instance: Tokenizer = self.tokenizer_instance or Tokenizer(
+            model_name=os.getenv("TOKENIZER_MODEL")
         )
-        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=self.tokenizer_instance,
+
+        self.synthesizer_llm_client: OpenAIClient = (
+            self.synthesizer_llm_client
+            or OpenAIClient(
+                model_name=os.getenv("SYNTHESIZER_MODEL"),
+                api_key=os.getenv("SYNTHESIZER_API_KEY"),
+                base_url=os.getenv("SYNTHESIZER_BASE_URL"),
+                tokenizer=self.tokenizer_instance,
+            )
         )
-        self.trainee_llm_client: OpenAIClient = OpenAIClient(
+
+        self.trainee_llm_client: OpenAIClient = self.trainee_llm_client or OpenAIClient(
             model_name=os.getenv("TRAINEE_MODEL"),
             api_key=os.getenv("TRAINEE_API_KEY"),
             base_url=os.getenv("TRAINEE_BASE_URL"),
             tokenizer=self.tokenizer_instance,
         )
-        self.search_config = self.config["search"]
-
-        if "traverse_strategy" in self.config:
-            self.traverse_strategy = TraverseStrategy(
-                **self.config["traverse_strategy"]
-            )
 
         self.full_docs_storage: JsonKVStorage = JsonKVStorage(
             self.working_dir, namespace="full_docs"
@@ -99,24 +88,17 @@ class GraphGen:
             self.working_dir, namespace="rephrase"
         )
         self.qa_storage: JsonListStorage = JsonListStorage(
-            os.path.join(
-                self.working_dir,
-                "data",
-                "graphgen",
-                f"{self.unique_id}_{self.config['output_data_type']}",
-            ),
+            os.path.join(self.working_dir, "data", "graphgen", f"{self.unique_id}"),
             namespace="qa",
         )
 
     @async_to_sync_method
-    async def insert(self):
+    async def insert(self, read_config: Dict, split_config: Dict):
         """
         insert chunks into the graph
         """
-        input_file = self.config["read"]["input_file"]
-
         # Step 1: Read files
-        data = read_files(input_file)
+        data = read_files(read_config["input_file"])
         if len(data) == 0:
             logger.warning("No data to process")
             return
@@ -141,8 +123,8 @@ class GraphGen:
 
         inserting_chunks = await chunk_documents(
             new_docs,
-            self.config["split"]["chunk_size"],
-            self.config["split"]["chunk_overlap"],
+            split_config["chunk_size"],
+            split_config["chunk_overlap"],
             self.tokenizer_instance,
             self.progress_bar,
         )
@@ -178,6 +160,7 @@ class GraphGen:
             return
 
         await self._insert_done()
+        return _add_entities_and_relations
 
     async def _insert_done(self):
         tasks = []
@@ -193,14 +176,12 @@ class GraphGen:
         await asyncio.gather(*tasks)
 
     @async_to_sync_method
-    async def search(self):
+    async def search(self, search_config: Dict):
         logger.info(
-            "Search is %s", "enabled" if self.search_config["enabled"] else "disabled"
+            "Search is %s", "enabled" if search_config["enabled"] else "disabled"
         )
-        if self.search_config["enabled"]:
-            logger.info(
-                "[Search] %s ...", ", ".join(self.search_config["search_types"])
-            )
+        if search_config["enabled"]:
+            logger.info("[Search] %s ...", ", ".join(search_config["search_types"]))
             all_nodes = await self.graph_storage.get_all_nodes()
             all_nodes_names = [node[0] for node in all_nodes]
             new_search_entities = await self.full_docs_storage.filter_keys(
@@ -210,7 +191,7 @@ class GraphGen:
                 "[Search] Found %d entities to search", len(new_search_entities)
             )
             _add_search_data = await search_all(
-                search_types=self.search_config["search_types"],
+                search_types=search_config["search_types"],
                 search_entities=new_search_entities,
             )
             if _add_search_data:
@@ -230,78 +211,77 @@ class GraphGen:
                 await self.insert()
 
     @async_to_sync_method
-    async def quiz(self):
-        max_samples = self.config["quiz_and_judge_strategy"]["quiz_samples"]
+    async def quiz_and_judge(self, quiz_and_judge_config: Dict):
+        if quiz_and_judge_config is None or not quiz_and_judge_config.get(
+            "enabled", False
+        ):
+            logger.warning("Quiz and Judge is not used in this pipeline.")
+            return
+        max_samples = quiz_and_judge_config["quiz_samples"]
         await quiz(
             self.synthesizer_llm_client,
             self.graph_storage,
             self.rephrase_storage,
             max_samples,
         )
-        await self.rephrase_storage.index_done_callback()
 
-    @async_to_sync_method
-    async def judge(self):
-        re_judge = self.config["quiz_and_judge_strategy"]["re_judge"]
+        # TODO： assert trainee_llm_client is valid before judge
+        re_judge = quiz_and_judge_config["re_judge"]
         _update_relations = await judge_statement(
             self.trainee_llm_client,
             self.graph_storage,
             self.rephrase_storage,
             re_judge,
         )
+        await self.rephrase_storage.index_done_callback()
         await _update_relations.index_done_callback()
 
     @async_to_sync_method
-    async def traverse(self):
-        output_data_type = self.config["output_data_type"]
-
-        if output_data_type == "atomic":
+    async def generate(self, partition_config: Dict, generate_config: Dict):
+        # Step 1: partition the graph
+        # TODO: implement graph partitioning, e.g. Partitioner().partition(self.graph_storage)
+        mode = generate_config["mode"]
+        if mode == "atomic":
             results = await traverse_graph_for_atomic(
                 self.synthesizer_llm_client,
                 self.tokenizer_instance,
                 self.graph_storage,
-                self.traverse_strategy,
+                partition_config["method_params"],
                 self.text_chunks_storage,
                 self.progress_bar,
             )
-        elif output_data_type == "multi_hop":
+        elif mode == "multi_hop":
             results = await traverse_graph_for_multi_hop(
                 self.synthesizer_llm_client,
                 self.tokenizer_instance,
                 self.graph_storage,
-                self.traverse_strategy,
+                partition_config["method_params"],
                 self.text_chunks_storage,
                 self.progress_bar,
             )
-        elif output_data_type == "aggregated":
+        elif mode == "aggregated":
             results = await traverse_graph_for_aggregated(
                 self.synthesizer_llm_client,
                 self.tokenizer_instance,
                 self.graph_storage,
-                self.traverse_strategy,
+                partition_config["method_params"],
                 self.text_chunks_storage,
                 self.progress_bar,
             )
+        elif mode == "cot":
+            results = await generate_cot(
+                self.graph_storage,
+                self.synthesizer_llm_client,
+                method_params=partition_config["method_params"],
+            )
         else:
-            raise ValueError(f"Unknown qa_form: {output_data_type}")
-
-        results = format_generation_results(
-            results, output_data_format=self.config["output_data_format"]
-        )
-
-        await self.qa_storage.upsert(results)
-        await self.qa_storage.index_done_callback()
-
-    @async_to_sync_method
-    async def generate_reasoning(self, method_params):
-        results = await generate_cot(
-            self.graph_storage,
-            self.synthesizer_llm_client,
-            method_params=method_params,
-        )
+            raise ValueError(f"Unknown generation mode: {mode}")
+        # Step 2： generate QA pairs
+        # TODO
 
+        # Step 3: format
         results = format_generation_results(
-            results, output_data_format=self.config["output_data_format"]
+            results, output_data_format=generate_config["data_format"]
         )
 
         await self.qa_storage.upsert(results)

graphgen/models/__init__.py

@@ -13,5 +13,4 @@ from .search.web.google_search import GoogleSearch
 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/strategy/travserse_strategy.py

@@ -1,28 +0,0 @@
-from dataclasses import dataclass, fields
-
-
-@dataclass
-class TraverseStrategy:
-    # 生成的QA形式：原子、多跳、聚合型
-    qa_form: str = "atomic"  # "atomic" or "multi_hop" or "aggregated"
-    # 最大边数和最大token数方法中选择一个生效
-    expand_method: str = "max_tokens"  # "max_width" or "max_tokens"
-    # 单向拓展还是双向拓展
-    bidirectional: bool = True
-    # 每个方向拓展的最大边数
-    max_extra_edges: int = 5
-    # 最长token数
-    max_tokens: int = 256
-    # 每个方向拓展的最大深度
-    max_depth: int = 2
-    # 同一层中选边的策略（如果是双向拓展，同一层指的是两边连接的边的集合）
-    edge_sampling: str = "max_loss"  # "max_loss" or "min_loss" or "random"
-    # 孤立节点的处理策略
-    isolated_node_strategy: str = "add"  # "add" or "ignore"
-    loss_strategy: str = "only_edge"  # only_edge, both
-
-    def to_yaml(self):
-        strategy_dict = {}
-        for f in fields(self):
-            strategy_dict[f.name] = getattr(self, f.name)
-        return {"traverse_strategy": strategy_dict}

graphgen/models/tokenizer/__init__.py

@@ -39,6 +39,8 @@ class Tokenizer(BaseTokenizer):
     _impl: BaseTokenizer = field(init=False, repr=False)
 
     def __post_init__(self):
+        if not self.model_name:
+            raise ValueError("TOKENIZER_MODEL must be specified in the ENV variables.")
         self._impl = get_tokenizer_impl(self.model_name)
 
     def encode(self, text: str) -> List[int]:

graphgen/operators/build_kg/split_kg.py

@@ -1,9 +1,10 @@
 import random
 from collections import defaultdict
+from typing import Dict
 
 from tqdm.asyncio import tqdm as tqdm_async
 
-from graphgen.models import NetworkXStorage, TraverseStrategy
+from graphgen.models import NetworkXStorage
 from graphgen.utils import logger
 
 
@@ -247,9 +248,9 @@ async def get_batches_with_strategy(  # pylint: disable=too-many-branches
     nodes: list,
     edges: list,
     graph_storage: NetworkXStorage,
-    traverse_strategy: TraverseStrategy,
+    traverse_strategy: Dict,
 ):
-    expand_method = traverse_strategy.expand_method
+    expand_method = traverse_strategy["expand_method"]
     if expand_method == "max_width":
         logger.info("Using max width strategy")
     elif expand_method == "max_tokens":
@@ -257,8 +258,8 @@ async def get_batches_with_strategy(  # pylint: disable=too-many-branches
     else:
         raise ValueError(f"Invalid expand method: {expand_method}")
 
-    max_depth = traverse_strategy.max_depth
-    edge_sampling = traverse_strategy.edge_sampling
+    max_depth = traverse_strategy["max_depth"]
+    edge_sampling = traverse_strategy["edge_sampling"]
 
     # 构建临接矩阵
     edge_adj_list = defaultdict(list)
@@ -275,16 +276,16 @@ async def get_batches_with_strategy(  # pylint: disable=too-many-branches
     for i, (node_name, _) in enumerate(nodes):
         node_dict[node_name] = i
 
-    if traverse_strategy.loss_strategy == "both":
+    if traverse_strategy["loss_strategy"] == "both":
         er_tuples = [
             ([nodes[node_dict[edge[0]]], nodes[node_dict[edge[1]]]], edge)
             for edge in edges
         ]
         edges = _sort_tuples(er_tuples, edge_sampling)
-    elif traverse_strategy.loss_strategy == "only_edge":
+    elif traverse_strategy["loss_strategy"] == "only_edge":
         edges = _sort_edges(edges, edge_sampling)
     else:
-        raise ValueError(f"Invalid loss strategy: {traverse_strategy.loss_strategy}")
+        raise ValueError(f"Invalid loss strategy: {traverse_strategy['loss_strategy']}")
 
     for i, (src, tgt, _) in enumerate(edges):
         edge_adj_list[src].append(i)
@@ -315,10 +316,10 @@ async def get_batches_with_strategy(  # pylint: disable=too-many-branches
                 nodes,
                 edge,
                 max_depth,
-                traverse_strategy.bidirectional,
-                traverse_strategy.max_extra_edges,
+                traverse_strategy["bidirectional"],
+                traverse_strategy["max_extra_edges"],
                 edge_sampling,
-                traverse_strategy.loss_strategy,
+                traverse_strategy["loss_strategy"],
             )
         else:
             level_n_edges = _get_level_n_edges_by_max_tokens(
@@ -328,10 +329,10 @@ async def get_batches_with_strategy(  # pylint: disable=too-many-branches
                 nodes,
                 edge,
                 max_depth,
-                traverse_strategy.bidirectional,
-                traverse_strategy.max_tokens,
+                traverse_strategy["bidirectional"],
+                traverse_strategy["max_tokens"],
                 edge_sampling,
-                traverse_strategy.loss_strategy,
+                traverse_strategy["loss_strategy"],
             )
 
         for _edge in level_n_edges:
@@ -352,7 +353,7 @@ async def get_batches_with_strategy(  # pylint: disable=too-many-branches
     logger.info("Processing batches: %d", len(processing_batches))
 
     # isolate nodes
-    isolated_node_strategy = traverse_strategy.isolated_node_strategy
+    isolated_node_strategy = traverse_strategy["isolated_node_strategy"]
     if isolated_node_strategy == "add":
         processing_batches = await _add_isolated_nodes(
             nodes, processing_batches, graph_storage

graphgen/operators/traverse_graph.py

@@ -1,15 +1,10 @@
 import asyncio
+from typing import Dict
 
 import gradio as gr
 from tqdm.asyncio import tqdm as tqdm_async
 
-from graphgen.models import (
-    JsonKVStorage,
-    NetworkXStorage,
-    OpenAIClient,
-    Tokenizer,
-    TraverseStrategy,
-)
+from graphgen.models import JsonKVStorage, NetworkXStorage, OpenAIClient, Tokenizer
 from graphgen.operators.build_kg.split_kg import get_batches_with_strategy
 from graphgen.templates import (
     ANSWER_REPHRASING_PROMPT,
@@ -164,7 +159,7 @@ async def traverse_graph_for_aggregated(
     llm_client: OpenAIClient,
     tokenizer: Tokenizer,
     graph_storage: NetworkXStorage,
-    traverse_strategy: TraverseStrategy,
+    traverse_strategy: Dict,
     text_chunks_storage: JsonKVStorage,
     progress_bar: gr.Progress = None,
     max_concurrent: int = 1000,
@@ -240,7 +235,7 @@ async def traverse_graph_for_aggregated(
                         "question": question,
                         "answer": context,
                         "loss": get_average_loss(
-                            _process_batch, traverse_strategy.loss_strategy
+                            _process_batch, traverse_strategy["loss_strategy"]
                         ),
                     }
                 }
@@ -272,7 +267,7 @@ async def traverse_graph_for_aggregated(
                     "question": qa["question"],
                     "answer": qa["answer"],
                     "loss": get_average_loss(
-                        _process_batch, traverse_strategy.loss_strategy
+                        _process_batch, traverse_strategy["loss_strategy"]
                     ),
                 }
             return final_results
@@ -313,7 +308,7 @@ async def traverse_graph_for_atomic(
     llm_client: OpenAIClient,
     tokenizer: Tokenizer,
     graph_storage: NetworkXStorage,
-    traverse_strategy: TraverseStrategy,
+    traverse_strategy: Dict,
     text_chunks_storage: JsonKVStorage,
     progress_bar: gr.Progress = None,
     max_concurrent: int = 1000,
@@ -331,7 +326,6 @@ async def traverse_graph_for_atomic(
     :return: question and answer
     """
 
-    assert traverse_strategy.qa_form == "atomic"
     semaphore = asyncio.Semaphore(max_concurrent)
 
     def _parse_qa(qa: str) -> tuple:
@@ -429,7 +423,7 @@ async def traverse_graph_for_multi_hop(
     llm_client: OpenAIClient,
     tokenizer: Tokenizer,
     graph_storage: NetworkXStorage,
-    traverse_strategy: TraverseStrategy,
+    traverse_strategy: Dict,
     text_chunks_storage: JsonKVStorage,
     progress_bar: gr.Progress = None,
     max_concurrent: int = 1000,
@@ -517,7 +511,7 @@ async def traverse_graph_for_multi_hop(
                         "question": question,
                         "answer": answer,
                         "loss": get_average_loss(
-                            _process_batch, traverse_strategy.loss_strategy
+                            _process_batch, traverse_strategy["loss_strategy"]
                         ),
                     }
                 }

webui/app.py

@@ -39,27 +39,32 @@ def init_graph_gen(config: dict, env: dict) -> GraphGen:
     set_logger(log_file, if_stream=True)
     os.environ.update({k: str(v) for k, v in env.items()})
 
-    graph_gen = GraphGen(working_dir=working_dir, config=config)
-    # Set up LLM clients
-    graph_gen.synthesizer_llm_client = OpenAIClient(
+    tokenizer_instance = Tokenizer(config.get("tokenizer", "cl100k_base"))
+    synthesizer_llm_client = OpenAIClient(
         model_name=env.get("SYNTHESIZER_MODEL", ""),
         base_url=env.get("SYNTHESIZER_BASE_URL", ""),
         api_key=env.get("SYNTHESIZER_API_KEY", ""),
         request_limit=True,
         rpm=RPM(env.get("RPM", 1000)),
         tpm=TPM(env.get("TPM", 50000)),
+        tokenizer=tokenizer_instance,
     )
-
-    graph_gen.trainee_llm_client = OpenAIClient(
+    trainee_llm_client = OpenAIClient(
         model_name=env.get("TRAINEE_MODEL", ""),
         base_url=env.get("TRAINEE_BASE_URL", ""),
         api_key=env.get("TRAINEE_API_KEY", ""),
         request_limit=True,
         rpm=RPM(env.get("RPM", 1000)),
         tpm=TPM(env.get("TPM", 50000)),
+        tokenizer=tokenizer_instance,
     )
 
-    graph_gen.tokenizer_instance = Tokenizer(config.get("tokenizer", "cl100k_base"))
+    graph_gen = GraphGen(
+        working_dir=working_dir,
+        tokenizer_instance=tokenizer_instance,
+        synthesizer_llm_client=synthesizer_llm_client,
+        trainee_llm_client=trainee_llm_client,
+    )
 
     return graph_gen
 
@@ -78,27 +83,32 @@ def run_graphgen(params: WebuiParams, progress=gr.Progress()):
             "chunk_size": params.chunk_size,
             "chunk_overlap": params.chunk_overlap,
         },
-        "output_data_type": params.output_data_type,
-        "output_data_format": params.output_data_format,
-        "tokenizer": params.tokenizer,
         "search": {"enabled": False},
-        "quiz_and_judge_strategy": {
+        "quiz_and_judge": {
             "enabled": params.if_trainee_model,
             "quiz_samples": params.quiz_samples,
         },
-        "traverse_strategy": {
-            "bidirectional": params.bidirectional,
-            "expand_method": params.expand_method,
-            "max_extra_edges": params.max_extra_edges,
-            "max_tokens": params.max_tokens,
-            "max_depth": params.max_depth,
-            "edge_sampling": params.edge_sampling,
-            "isolated_node_strategy": params.isolated_node_strategy,
-            "loss_strategy": params.loss_strategy,
+        "partition": {
+            "method": "ece",
+            "method_params": {
+                "bidirectional": params.bidirectional,
+                "expand_method": params.expand_method,
+                "max_extra_edges": params.max_extra_edges,
+                "max_tokens": params.max_tokens,
+                "max_depth": params.max_depth,
+                "edge_sampling": params.edge_sampling,
+                "isolated_node_strategy": params.isolated_node_strategy,
+                "loss_strategy": params.loss_strategy,
+            },
+        },
+        "generate": {
+            "mode": params.output_data_type,
+            "data_format": params.output_data_format,
         },
     }
 
     env = {
+        "TOKENIZER_MODEL": params.tokenizer,
         "SYNTHESIZER_BASE_URL": params.synthesizer_url,
         "SYNTHESIZER_MODEL": params.synthesizer_model,
         "TRAINEE_BASE_URL": params.trainee_url,
@@ -128,19 +138,18 @@ def run_graphgen(params: WebuiParams, progress=gr.Progress()):
 
     try:
         # Process the data
-        graph_gen.insert()
+        graph_gen.insert(read_config=config["read"], split_config=config["split"])
 
         if config["if_trainee_model"]:
-            # Generate quiz
-            graph_gen.quiz()
-
-            # Judge statements
-            graph_gen.judge()
+            # Quiz and Judge
+            graph_gen.quiz_and_judge(quiz_and_judge_config=config["quiz_and_judge"])
         else:
-            graph_gen.traverse_strategy.edge_sampling = "random"
+            config["partition"]["method_params"]["edge_sampling"] = "random"
 
-        # Traverse graph
-        graph_gen.traverse()
+        graph_gen.generate(
+            partition_config=config["partition"],
+            generate_config=config["generate"],
+        )
 
         # Save output
         output_data = graph_gen.qa_storage.data
@@ -249,6 +258,9 @@ with gr.Blocks(title="GraphGen Demo", theme=gr.themes.Glass(), css=css) as demo:
         )
 
         with gr.Accordion(label=_("Model Config"), open=False):
+            tokenizer = gr.Textbox(
+                label="Tokenizer", value="cl100k_base", interactive=True
+            )
             synthesizer_url = gr.Textbox(
                 label="Synthesizer URL",
                 value="https://api.siliconflow.cn/v1",
@@ -300,9 +312,6 @@ with gr.Blocks(title="GraphGen Demo", theme=gr.themes.Glass(), css=css) as demo:
                 step=100,
                 interactive=True,
             )
-            tokenizer = gr.Textbox(
-                label="Tokenizer", value="cl100k_base", interactive=True
-            )
             output_data_type = gr.Radio(
                 choices=["atomic", "multi_hop", "aggregated"],
                 label="Output Data Type",