Merge pull request #98 from CHERRY-ui8/feat/add-progress-bar-and-refactor-concurrent

ChenZiHong-Gavin · web-flow · commit 66b1e66d36a4 · 2025-11-21T14:05:41.000+08:00
refactor: refactor concurrent and add progress bar
diff --git a/graphgen/graphgen.py b/graphgen/graphgen.py
@@ -237,6 +237,7 @@ async def quiz_and_judge(self, quiz_and_judge_config: Dict):
             self.graph_storage,
             self.rephrase_storage,
             re_judge,
+            progress_bar=self.progress_bar,
         )
 
         await self.rephrase_storage.index_done_callback()
diff --git a/graphgen/operators/partition/pre_tokenize.py b/graphgen/operators/partition/pre_tokenize.py
@@ -1,6 +1,8 @@
 import asyncio
 from typing import List, Tuple
 
+import gradio as gr
+
 from graphgen.bases import BaseGraphStorage, BaseTokenizer
 from graphgen.utils import run_concurrent
 
@@ -10,9 +12,11 @@ async def pre_tokenize(
     tokenizer: BaseTokenizer,
     edges: List[Tuple],
     nodes: List[Tuple],
+    progress_bar: gr.Progress = None,
+    max_concurrent: int = 1000,
 ) -> Tuple[List, List]:
     """为 edges/nodes 补 token-length 并回写存储，并发 1000，带进度条。"""
-    sem = asyncio.Semaphore(1000)
+    sem = asyncio.Semaphore(max_concurrent)
 
     async def _patch_and_write(obj: Tuple, *, is_node: bool) -> Tuple:
         async with sem:
@@ -35,11 +39,15 @@ async def _patch_and_write(obj: Tuple, *, is_node: bool) -> Tuple:
             lambda e: _patch_and_write(e, is_node=False),
             edges,
             desc="Pre-tokenizing edges",
+            unit="edge",
+            progress_bar=progress_bar,
         ),
         run_concurrent(
             lambda n: _patch_and_write(n, is_node=True),
             nodes,
             desc="Pre-tokenizing nodes",
+            unit="node",
+            progress_bar=progress_bar,
         ),
     )
 
diff --git a/graphgen/operators/quiz_and_judge/judge.py b/graphgen/operators/quiz_and_judge/judge.py
@@ -1,20 +1,19 @@
-import asyncio
 import math
 
-from tqdm.asyncio import tqdm as tqdm_async
+import gradio as gr
 
 from graphgen.bases import BaseLLMWrapper
 from graphgen.models import JsonKVStorage, NetworkXStorage
 from graphgen.templates import STATEMENT_JUDGEMENT_PROMPT
-from graphgen.utils import logger, yes_no_loss_entropy
+from graphgen.utils import logger, run_concurrent, yes_no_loss_entropy
 
 
 async def judge_statement(  # pylint: disable=too-many-statements
     trainee_llm_client: BaseLLMWrapper,
     graph_storage: NetworkXStorage,
     rephrase_storage: JsonKVStorage,
     re_judge: bool = False,
-    max_concurrent: int = 1000,
+    progress_bar: gr.Progress = None,
 ) -> NetworkXStorage:
     """
     Get all edges and nodes and judge them
@@ -23,128 +22,124 @@ async def judge_statement(  # pylint: disable=too-many-statements
     :param graph_storage: graph storage instance
     :param rephrase_storage: rephrase storage instance
     :param re_judge: re-judge the relations
-    :param max_concurrent: max concurrent
+    :param progress_bar
     :return:
     """
 
-    semaphore = asyncio.Semaphore(max_concurrent)
-
     async def _judge_single_relation(
         edge: tuple,
     ):
-        async with semaphore:
-            source_id = edge[0]
-            target_id = edge[1]
-            edge_data = edge[2]
-
-            if (not re_judge) and "loss" in edge_data and edge_data["loss"] is not None:
-                logger.debug(
-                    "Edge %s -> %s already judged, loss: %s, skip",
-                    source_id,
-                    target_id,
-                    edge_data["loss"],
-                )
-                return source_id, target_id, edge_data
+        source_id = edge[0]
+        target_id = edge[1]
+        edge_data = edge[2]
+
+        if (not re_judge) and "loss" in edge_data and edge_data["loss"] is not None:
+            logger.debug(
+                "Edge %s -> %s already judged, loss: %s, skip",
+                source_id,
+                target_id,
+                edge_data["loss"],
+            )
+            return source_id, target_id, edge_data
 
-            description = edge_data["description"]
+        description = edge_data["description"]
 
-            try:
-                descriptions = await rephrase_storage.get_by_id(description)
-                assert descriptions is not None
+        try:
+            descriptions = await rephrase_storage.get_by_id(description)
+            assert descriptions is not None
 
-                judgements = []
-                gts = [gt for _, gt in descriptions]
-                for description, gt in descriptions:
-                    judgement = await trainee_llm_client.generate_topk_per_token(
-                        STATEMENT_JUDGEMENT_PROMPT["TEMPLATE"].format(
-                            statement=description
-                        )
+            judgements = []
+            gts = [gt for _, gt in descriptions]
+            for description, gt in descriptions:
+                judgement = await trainee_llm_client.generate_topk_per_token(
+                    STATEMENT_JUDGEMENT_PROMPT["TEMPLATE"].format(
+                        statement=description
                     )
-                    judgements.append(judgement[0].top_candidates)
+                )
+                judgements.append(judgement[0].top_candidates)
 
-                loss = yes_no_loss_entropy(judgements, gts)
+            loss = yes_no_loss_entropy(judgements, gts)
 
-                logger.debug(
-                    "Edge %s -> %s description: %s loss: %s",
-                    source_id,
-                    target_id,
-                    description,
-                    loss,
-                )
+            logger.debug(
+                "Edge %s -> %s description: %s loss: %s",
+                source_id,
+                target_id,
+                description,
+                loss,
+            )
 
-                edge_data["loss"] = loss
-            except Exception as e:  # pylint: disable=broad-except
-                logger.error(
-                    "Error in judging relation %s -> %s: %s", source_id, target_id, e
-                )
-                logger.info("Use default loss 0.1")
-                edge_data["loss"] = -math.log(0.1)
+            edge_data["loss"] = loss
+        except Exception as e:  # pylint: disable=broad-except
+            logger.error(
+                "Error in judging relation %s -> %s: %s", source_id, target_id, e
+            )
+            logger.info("Use default loss 0.1")
+            edge_data["loss"] = -math.log(0.1)
 
-            await graph_storage.update_edge(source_id, target_id, edge_data)
-            return source_id, target_id, edge_data
+        await graph_storage.update_edge(source_id, target_id, edge_data)
+        return source_id, target_id, edge_data
 
     edges = await graph_storage.get_all_edges()
 
-    results = []
-    for result in tqdm_async(
-        asyncio.as_completed([_judge_single_relation(edge) for edge in edges]),
-        total=len(edges),
+    await run_concurrent(
+        _judge_single_relation,
+        edges,
         desc="Judging relations",
-    ):
-        results.append(await result)
+        unit="relation",
+        progress_bar=progress_bar,
+    )
 
     async def _judge_single_entity(
         node: tuple,
     ):
-        async with semaphore:
-            node_id = node[0]
-            node_data = node[1]
+        node_id = node[0]
+        node_data = node[1]
 
-            if (not re_judge) and "loss" in node_data and node_data["loss"] is not None:
-                logger.debug(
-                    "Node %s already judged, loss: %s, skip", node_id, node_data["loss"]
-                )
-                return node_id, node_data
+        if (not re_judge) and "loss" in node_data and node_data["loss"] is not None:
+            logger.debug(
+                "Node %s already judged, loss: %s, skip", node_id, node_data["loss"]
+            )
+            return node_id, node_data
 
-            description = node_data["description"]
+        description = node_data["description"]
 
-            try:
-                descriptions = await rephrase_storage.get_by_id(description)
-                assert descriptions is not None
+        try:
+            descriptions = await rephrase_storage.get_by_id(description)
+            assert descriptions is not None
 
-                judgements = []
-                gts = [gt for _, gt in descriptions]
-                for description, gt in descriptions:
-                    judgement = await trainee_llm_client.generate_topk_per_token(
-                        STATEMENT_JUDGEMENT_PROMPT["TEMPLATE"].format(
-                            statement=description
-                        )
+            judgements = []
+            gts = [gt for _, gt in descriptions]
+            for description, gt in descriptions:
+                judgement = await trainee_llm_client.generate_topk_per_token(
+                    STATEMENT_JUDGEMENT_PROMPT["TEMPLATE"].format(
+                        statement=description
                     )
-                    judgements.append(judgement[0].top_candidates)
+                )
+                judgements.append(judgement[0].top_candidates)
 
-                loss = yes_no_loss_entropy(judgements, gts)
+            loss = yes_no_loss_entropy(judgements, gts)
 
-                logger.debug(
-                    "Node %s description: %s loss: %s", node_id, description, loss
-                )
+            logger.debug(
+                "Node %s description: %s loss: %s", node_id, description, loss
+            )
 
-                node_data["loss"] = loss
-            except Exception as e:  # pylint: disable=broad-except
-                logger.error("Error in judging entity %s: %s", node_id, e)
-                logger.error("Use default loss 0.1")
-                node_data["loss"] = -math.log(0.1)
+            node_data["loss"] = loss
+        except Exception as e:  # pylint: disable=broad-except
+            logger.error("Error in judging entity %s: %s", node_id, e)
+            logger.error("Use default loss 0.1")
+            node_data["loss"] = -math.log(0.1)
 
-            await graph_storage.update_node(node_id, node_data)
-            return node_id, node_data
+        await graph_storage.update_node(node_id, node_data)
+        return node_id, node_data
 
     nodes = await graph_storage.get_all_nodes()
 
-    results = []
-    for result in tqdm_async(
-        asyncio.as_completed([_judge_single_entity(node) for node in nodes]),
-        total=len(nodes),
+    await run_concurrent(
+        _judge_single_entity,
+        nodes,
         desc="Judging entities",
-    ):
-        results.append(await result)
+        unit="entity",
+        progress_bar=progress_bar,
+    )
 
     return graph_storage

Original file line number	Diff line number	Diff line change
`@@ -237,6 +237,7 @@ async def quiz_and_judge(self, quiz_and_judge_config: Dict):`
`237`	`237`	`self.graph_storage,`
`238`	`238`	`self.rephrase_storage,`
`239`	`239`	`re_judge,`
	`240`	`+ progress_bar=self.progress_bar,`
`240`	`241`	`)`
`241`	`242`
`242`	`243`	`await self.rephrase_storage.index_done_callback()`