fix: fix progress_bar in run_concurrent

Author: ChenZiHong-Gavin

graphgen/graphgen.py

@@ -237,7 +237,10 @@ async def generate(self, partition_config: Dict, generate_config: Dict):
 
         # Step 2： generate QA pairs
         results = await generate_qas(
-            self.synthesizer_llm_client, batches, generate_config
+            self.synthesizer_llm_client,
+            batches,
+            generate_config,
+            progress_bar=self.progress_bar,
         )
 
         if not results:

graphgen/models/kg_builder/light_rag_kg_builder.py

@@ -42,7 +42,7 @@ async def extract(
 
         # step 2: initial glean
         final_result = await self.llm_client.generate_answer(hint_prompt)
-        logger.debug("First extraction result: %s", final_result)
+        logger.info("First extraction result: %s", final_result)
 
         # step3: iterative refinement
         history = pack_history_conversations(hint_prompt, final_result)
@@ -57,7 +57,7 @@ async def extract(
             glean_result = await self.llm_client.generate_answer(
                 text=KG_EXTRACTION_PROMPT[language]["CONTINUE"], history=history
             )
-            logger.debug("Loop %s glean: %s", loop_idx + 1, glean_result)
+            logger.info("Loop %s glean: %s", loop_idx + 1, glean_result)
 
             history += pack_history_conversations(
                 KG_EXTRACTION_PROMPT[language]["CONTINUE"], glean_result

graphgen/operators/generate/generate_qas.py

@@ -18,12 +18,14 @@ async def generate_qas(
         ]
     ],
     generation_config: dict,
+    progress_bar=None,
 ) -> list[dict[str, Any]]:
     """
     Generate question-answer pairs based on nodes and edges.
     :param llm_client: LLM client
     :param batches
     :param generation_config
+    :param progress_bar
     :return: QA pairs
     """
     mode = generation_config["mode"]
@@ -45,6 +47,7 @@ async def generate_qas(
         batches,
         desc="[4/4]Generating QAs",
         unit="batch",
+        progress_bar=progress_bar,
     )
 
     # format

graphgen/utils/run_concurrent.py

@@ -10,6 +10,77 @@
 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 = []
+#     async for future in tqdm_async(
+#         tasks, desc=desc, unit=unit
+#     ):
+#         try:
+#             result = await future
+#             results.append(result)
+#         except Exception as e: # pylint: disable=broad-except
+#             logger.exception("Task failed: %s", e)
+#
+#         if progress_bar is not None:
+#             progress_bar((len(results)) / len(items), desc=desc)
+#
+#     if progress_bar is not None:
+#         progress_bar(1.0, desc=desc)
+#     return results
+
+#     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
+
+# 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 = []
+#     # 使用同步方式更新进度条，避免异步冲突
+#     for i, task in enumerate(asyncio.as_completed(tasks)):
+#         try:
+#             result = await task
+#             results.append(result)
+#             # 同步更新进度条
+#             if progress_bar is not None:
+#                 # 在同步上下文中更新进度
+#                 progress_bar((i + 1) / len(items), desc=desc)
+#         except Exception as e:
+#             logger.exception("Task failed: %s", e)
+#             results.append(e)
+#
+#     return results
+
+
 async def run_concurrent(
     coro_fn: Callable[[T], Awaitable[R]],
     items: List[T],
@@ -20,19 +91,36 @@ async def run_concurrent(
 ) -> 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
+    completed_count = 0
+    results = []
+
+    pbar = tqdm_async(total=len(items), desc=desc, unit=unit)
+
+    if progress_bar is not None:
+        progress_bar(0.0, desc=f"{desc} (0/{len(items)})")
+
+    for future in asyncio.as_completed(tasks):
+        try:
+            result = await future
+            results.append(result)
+        except Exception as e:  # pylint: disable=broad-except
+            logger.exception("Task failed: %s", e)
+            # even if failed, record it to keep results consistent with tasks
+            results.append(e)
+
+        completed_count += 1
+        pbar.update(1)
+
+        if progress_bar is not None:
+            progress = completed_count / len(items)
+            progress_bar(progress, desc=f"{desc} ({completed_count}/{len(items)})")
+
+    pbar.close()
+
+    if progress_bar is not None:
+        progress_bar(1.0, desc=f"{desc} (completed)")
+
+    # filter out exceptions
+    results = [res for res in results if not isinstance(res, Exception)]
+
+    return results