fix(graphgen): optimize quizzing

Author: ChenZiHong-Gavin

graphgen/operators/quiz_relations.py

@@ -1,4 +1,5 @@
 import asyncio
+from collections import defaultdict
 
 from tqdm.asyncio import tqdm as tqdm_async
 from models import JsonKVStorage, OpenAIModel, NetworkXStorage
@@ -27,67 +28,67 @@ async def quiz_relations(
 
     async def _quiz_single_relation(
         edge: tuple,
+        des: str,
+        prompt: str,
+        gt: str
     ):
         async with semaphore:
             source_id = edge[0]
             target_id = edge[1]
-            edge_data = edge[2]
-
-            description = edge_data["description"]
-            language = "English" if detect_main_language(description) == "en" else "Chinese"
 
             try:
                 # 如果在rephrase_storage中已经存在，直接取出
-                descriptions = await rephrase_storage.get_by_id(description)
-                if not descriptions:
-                    # 多次采样，取平均
-                    descriptions = [(description, 'yes')]
-
-                    new_description_tasks = []
-                    new_anti_description_tasks = []
-                    for i in range(max_samples):
-                        if i > 0:
-                            new_description_tasks.append(
-                                teacher_llm_client.generate_answer(
-                                    DESCRIPTION_REPHRASING_PROMPT[language]['TEMPLATE'].format(
-                                        input_sentence=description),
-                                    temperature=1
-                                )
-                            )
-                        new_anti_description_tasks.append(
-                            teacher_llm_client.generate_answer(
-                                DESCRIPTION_REPHRASING_PROMPT[language]['ANTI_TEMPLATE'].format(
-                                    input_sentence=description),
-                                temperature=1
-                            )
-                        )
-
-                    new_descriptions = await asyncio.gather(*new_description_tasks)
-                    new_anti_descriptions = await asyncio.gather(*new_anti_description_tasks)
-
-                    for new_description in new_descriptions:
-                        descriptions.append((new_description, 'yes'))
-                    for new_anti_description in new_anti_descriptions:
-                        descriptions.append((new_anti_description, 'no'))
-
-                    descriptions = list(set(descriptions))
+                descriptions = await rephrase_storage.get_by_id(des)
+                if descriptions:
+                    return None
+
+                new_description = await teacher_llm_client.generate_answer(
+                    prompt,
+                    temperature=1
+                )
+                return  {des: [(new_description, gt)]}
+
             except Exception as e: # pylint: disable=broad-except
                 logger.error("Error when quizzing edge %s -> %s: %s", source_id, target_id, e)
-                descriptions = [(description, 'yes')]
+                return None
+
 
-            await rephrase_storage.upsert({description: descriptions})
+    edges = await graph_storage.get_all_edges()
 
-            return {description: descriptions}
+    results = defaultdict(list)
+    tasks = []
+    for edge in edges:
+        edge_data = edge[2]
 
+        description = edge_data["description"]
+        language = "English" if detect_main_language(description) == "en" else "Chinese"
+
+        results[description] = [(description, 'yes')]
+
+        for i in range(max_samples):
+            if i > 0:
+                tasks.append(
+                    _quiz_single_relation(edge, description,
+                                          DESCRIPTION_REPHRASING_PROMPT[language]['TEMPLATE'].format(
+                                              input_sentence=description), 'yes')
+                )
+            tasks.append(_quiz_single_relation(edge, description,
+                                              DESCRIPTION_REPHRASING_PROMPT[language]['ANTI_TEMPLATE'].format(
+                                                  input_sentence=description), 'no'))
 
-    edges = await graph_storage.get_all_edges()
 
-    results = []
     for result in tqdm_async(
-            asyncio.as_completed([_quiz_single_relation(edge) for edge in edges]),
-            total=len(edges),
+            asyncio.as_completed(tasks),
+            total=len(tasks),
             desc="Quizzing relations"
     ):
-        results.append(await result)
+        new_result = await result
+        if new_result:
+            for key, value in new_result.items():
+                results[key].extend(value)
+
+    for key, value in results.items():
+        results[key] = list(set(value))
+        await rephrase_storage.upsert({key: results[key]})
 
     return rephrase_storage