refactor(webui): restructure GraphGen parameter handling

Author: ChenZiHong-Gavin

webui/app.py

@@ -8,6 +8,7 @@
 
 from gradio_i18n import Translate, gettext as _
 
+from base import GraphGenParams
 from test_api import test_api_connection
 from cache_utils import setup_workspace, cleanup_workspace
 from count_tokens import count_tokens
@@ -30,6 +31,7 @@
 }
 """
 
+
 def init_graph_gen(config: dict, env: dict) -> GraphGen:
     # Set up working directory
     log_file, working_dir = setup_workspace(os.path.join(root_dir, "cache"))
@@ -77,40 +79,39 @@ def init_graph_gen(config: dict, env: dict) -> GraphGen:
     return graph_gen
 
 # pylint: disable=too-many-statements
-def run_graphgen(*arguments: list, progress=gr.Progress()):
+def run_graphgen(params, progress=gr.Progress()):
     def sum_tokens(client):
         return sum(u["total_tokens"] for u in client.token_usage)
 
-    # Unpack arguments
     config = {
-        "if_trainee_model": arguments[0],
-        "input_file": arguments[1],
-        "tokenizer": arguments[2],
-        "qa_form": arguments[3],
+        "if_trainee_model": params.if_trainee_model,
+        "input_file": params.input_file,
+        "tokenizer": params.tokenizer,
+        "qa_form": params.qa_form,
         "web_search": False,
-        "quiz_samples": arguments[19],
+        "quiz_samples": params.quiz_samples,
         "traverse_strategy": {
-            "bidirectional": arguments[4],
-            "expand_method": arguments[5],
-            "max_extra_edges": arguments[6],
-            "max_tokens": arguments[7],
-            "max_depth": arguments[8],
-            "edge_sampling": arguments[9],
-            "isolated_node_strategy": arguments[10],
-            "loss_strategy": arguments[11]
+            "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
         },
-        "chunk_size": arguments[16],
+        "chunk_size": params.chunk_size,
     }
 
     env = {
-        "SYNTHESIZER_BASE_URL": arguments[12],
-        "SYNTHESIZER_MODEL": arguments[13],
-        "TRAINEE_BASE_URL": arguments[20],
-        "TRAINEE_MODEL": arguments[14],
-        "SYNTHESIZER_API_KEY": arguments[15],
-        "TRAINEE_API_KEY": arguments[21],
-        "RPM": arguments[17],
-        "TPM": arguments[18],
+        "SYNTHESIZER_BASE_URL": params.synthesizer_url,
+        "SYNTHESIZER_MODEL": params.synthesizer_model,
+        "TRAINEE_BASE_URL": params.trainee_url,
+        "TRAINEE_MODEL": params.trainee_model,
+        "SYNTHESIZER_API_KEY": params.api_key,
+        "TRAINEE_API_KEY": params.trainee_api_key,
+        "RPM": params.rpm,
+        "TPM": params.tpm,
     }
 
     # Test API connection
@@ -189,7 +190,7 @@ def sum_tokens(client):
         trainee_tokens = sum_tokens(graph_gen.trainee_llm_client) if config['if_trainee_model'] else 0
         total_tokens = synthesizer_tokens + trainee_tokens
 
-        data_frame = arguments[-1]
+        data_frame = params.token_counter
         try:
             _update_data = [
                 [
@@ -460,7 +461,6 @@ def sum_tokens(client):
             inputs=if_trainee_model,
             outputs=[trainee_url, trainee_model, quiz_samples, edge_sampling, trainee_api_key])
 
-        # 计算上传文件的token数
         upload_file.change(
             lambda x: (gr.update(visible=True)),
             inputs=[upload_file],
@@ -476,8 +476,34 @@ def sum_tokens(client):
             lambda x: (gr.update(visible=False)),
             inputs=[token_counter],
             outputs=[token_counter],
-        ).then(
-            run_graphgen,
+        )
+
+        submit_btn.click(
+            lambda *args: run_graphgen(GraphGenParams(
+                if_trainee_model=args[0],
+                input_file=args[1],
+                tokenizer=args[2],
+                qa_form=args[3],
+                bidirectional=args[4],
+                expand_method=args[5],
+                max_extra_edges=args[6],
+                max_tokens=args[7],
+                max_depth=args[8],
+                edge_sampling=args[9],
+                isolated_node_strategy=args[10],
+                loss_strategy=args[11],
+                synthesizer_url=args[12],
+                synthesizer_model=args[13],
+                trainee_model=args[14],
+                api_key=args[15],
+                chunk_size=args[16],
+                rpm=args[17],
+                tpm=args[18],
+                quiz_samples=args[19],
+                trainee_url=args[20],
+                trainee_api_key=args[21],
+                token_counter=args[22],
+            )),
             inputs=[
                 if_trainee_model, upload_file, tokenizer, qa_form,
                 bidirectional, expand_method, max_extra_edges, max_tokens,

webui/base.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+from typing import Any
+
+@dataclass
+class GraphGenParams:
+    """
+    GraphGen parameters
+    """
+    if_trainee_model: bool
+    input_file: str
+    tokenizer: str
+    qa_form: str
+    bidirectional: bool
+    expand_method: str
+    max_extra_edges: int
+    max_tokens: int
+    max_depth: int
+    edge_sampling: str
+    isolated_node_strategy: str
+    loss_strategy: str
+    synthesizer_url: str
+    synthesizer_model: str
+    trainee_model: str
+    api_key: str
+    chunk_size: int
+    rpm: int
+    tpm: int
+    quiz_samples: int
+    trainee_url: str
+    trainee_api_key: str
+    token_counter: Any