refactor: use run_concurrent in quiz operator and fix language codes

- Replace manual asyncio with run_concurrent utility
- Fix language codes: English/Chinese -> en/zh
- Add progress_bar support

Author: CHERRY-ui8

graphgen/graphgen.py

@@ -221,6 +221,7 @@ async def quiz_and_judge(self, quiz_and_judge_config: Dict):
             self.graph_storage,
             self.rephrase_storage,
             max_samples,
+            progress_bar=self.progress_bar,
         )
 
         # TODO： assert trainee_llm_client is valid before judge

graphgen/models/generator/quiz_generator.py

@@ -43,7 +43,7 @@ def build_prompt_for_description(description: str, template_type: str = "TEMPLAT
         :param template_type: Either "TEMPLATE" (same meaning) or "ANTI_TEMPLATE" (opposite meaning)
         :return: Prompt string
         """
-        language = "English" if detect_main_language(description) == "en" else "Chinese"
+        language = "en" if detect_main_language(description) == "en" else "zh"
         prompt = DESCRIPTION_REPHRASING_PROMPT[language][template_type].format(
             input_sentence=description
         )

graphgen/operators/generate/generate_qas.py

@@ -1,5 +1,7 @@
 from typing import Any
 
+import gradio as gr
+
 from graphgen.bases import BaseLLMWrapper
 from graphgen.models import (
     AggregatedGenerator,
@@ -19,7 +21,7 @@ async def generate_qas(
         ]
     ],
     generation_config: dict,
-    progress_bar=None,
+    progress_bar: gr.Progress = None,
 ) -> list[dict[str, Any]]:
     """
     Generate question-answer pairs based on nodes and edges.

graphgen/operators/quiz_and_judge/quiz.py

@@ -1,19 +1,18 @@
-import asyncio
 from collections import defaultdict
 
-from tqdm.asyncio import tqdm as tqdm_async
+import gradio as gr
 
 from graphgen.bases import BaseLLMWrapper
 from graphgen.models import JsonKVStorage, NetworkXStorage, QuizGenerator
-from graphgen.utils import logger
+from graphgen.utils import logger, run_concurrent
 
 
 async def quiz(
     synth_llm_client: BaseLLMWrapper,
     graph_storage: NetworkXStorage,
     rephrase_storage: JsonKVStorage,
     max_samples: int = 1,
-    max_concurrent: int = 1000,
+    progress_bar: gr.Progress = None,
 ) -> JsonKVStorage:
     """
     Get all edges and quiz them using QuizGenerator.
@@ -22,37 +21,36 @@ async def quiz(
     :param graph_storage: graph storage instance
     :param rephrase_storage: rephrase storage instance
     :param max_samples: max samples for each edge
-    :param max_concurrent: max concurrent
+    :param progress_bar
     :return:
     """
 
-    semaphore = asyncio.Semaphore(max_concurrent)
     generator = QuizGenerator(synth_llm_client)
 
-    async def _process_single_quiz(description: str, template_type: str, gt: str):
-        async with semaphore:
-            try:
-                # if rephrase_storage exists already, directly get it
-                descriptions = await rephrase_storage.get_by_id(description)
-                if descriptions:
-                    return None
-
-                prompt = generator.build_prompt_for_description(description, template_type)
-                new_description = await synth_llm_client.generate_answer(
-                    prompt, temperature=1
-                )
-                rephrased_text = generator.parse_rephrased_text(new_description)
-                return {description: [(rephrased_text, gt)]}
-
-            except Exception as e:  # pylint: disable=broad-except
-                logger.error("Error when quizzing description %s: %s", description, e)
+    async def _process_single_quiz(item: tuple[str, str, str]):
+        description, template_type, gt = item
+        try:
+            # if rephrase_storage exists already, directly get it
+            descriptions = await rephrase_storage.get_by_id(description)
+            if descriptions:
                 return None
 
+            prompt = generator.build_prompt_for_description(description, template_type)
+            new_description = await synth_llm_client.generate_answer(
+                prompt, temperature=1
+            )
+            rephrased_text = generator.parse_rephrased_text(new_description)
+            return {description: [(rephrased_text, gt)]}
+
+        except Exception as e:  # pylint: disable=broad-except
+            logger.error("Error when quizzing description %s: %s", description, e)
+            return None
+
     edges = await graph_storage.get_all_edges()
     nodes = await graph_storage.get_all_nodes()
 
     results = defaultdict(list)
-    tasks = []
+    items = []
     for edge in edges:
         edge_data = edge[2]
         description = edge_data["description"]
@@ -61,12 +59,8 @@ async def _process_single_quiz(description: str, template_type: str, gt: str):
 
         for i in range(max_samples):
             if i > 0:
-                tasks.append(
-                    _process_single_quiz(description, "TEMPLATE", "yes")
-                )
-            tasks.append(
-                _process_single_quiz(description, "ANTI_TEMPLATE", "no")
-            )
+                items.append((description, "TEMPLATE", "yes"))
+            items.append((description, "ANTI_TEMPLATE", "no"))
 
     for node in nodes:
         node_data = node[1]
@@ -76,17 +70,18 @@ async def _process_single_quiz(description: str, template_type: str, gt: str):
 
         for i in range(max_samples):
             if i > 0:
-                tasks.append(
-                    _process_single_quiz(description, "TEMPLATE", "yes")
-                )
-            tasks.append(
-                _process_single_quiz(description, "ANTI_TEMPLATE", "no")
-            )
-
-    for result in tqdm_async(
-        asyncio.as_completed(tasks), total=len(tasks), desc="Quizzing descriptions"
-    ):
-        new_result = await result
+                items.append((description, "TEMPLATE", "yes"))
+            items.append((description, "ANTI_TEMPLATE", "no"))
+
+    quiz_results = await run_concurrent(
+        _process_single_quiz,
+        items,
+        desc="Quizzing descriptions",
+        unit="description",
+        progress_bar=progress_bar,
+    )
+
+    for new_result in quiz_results:
         if new_result:
             for key, value in new_result.items():
                 results[key].extend(value)

graphgen/templates/description_rephrasing.py

@@ -110,11 +110,11 @@
 
 
 DESCRIPTION_REPHRASING_PROMPT= {
-    "English": {
+    "en": {
         "ANTI_TEMPLATE": ANTI_TEMPLATE_EN,
         "TEMPLATE": TEMPLATE_EN
     },
-    "Chinese": {
+    "zh": {
         "ANTI_TEMPLATE": ANTI_TEMPLATE_ZH,
         "TEMPLATE": TEMPLATE_ZH
     }