fix: relocate progress bar

Author: ChenZiHong-Gavin

graphgen/operators/extract_kg.py

@@ -110,9 +110,11 @@ async def _process_single_content(chunk: Chunk, max_loop: int = 3):
         unit="chunk",
     ):
         try:
-            results.append(await result)
             if progress_bar is not None:
                 progress_bar(len(results) / chunk_number, desc="[3/4]Extracting entities and relationships from chunks")
+            results.append(await result)
+            if progress_bar is not None and len(results) == chunk_number:
+                progress_bar(1, desc="[3/4]Extracting entities and relationships from chunks")
         except Exception as e: # pylint: disable=broad-except
             logger.error("Error occurred while extracting entities and relationships from chunks: %s", e)
 

graphgen/operators/traverse_graph.py

@@ -294,9 +294,11 @@ async def _process_single_batch(
         [_process_single_batch(batch) for batch in processing_batches]
     ), total=len(processing_batches), desc="[4/4]Generating QAs"):
         try:
-            results.update(await result)
             if progress_bar is not None:
                 progress_bar(len(results) / len(processing_batches), desc="[4/4]Generating QAs")
+            results.update(await result)
+            if progress_bar is not None and len(results) == len(processing_batches):
+                progress_bar(1, desc="[4/4]Generating QAs")
         except Exception as e: # pylint: disable=broad-except
             logger.error("Error occurred while generating QA: %s", e)
 
@@ -401,9 +403,11 @@ async def _generate_question(
         desc="[4/4]Generating QAs"
     ):
         try:
-            results.update(await result)
             if progress_bar is not None:
                 progress_bar(len(results) / len(tasks), desc="[4/4]Generating QAs")
+            results.update(await result)
+            if progress_bar is not None and len(results) == len(tasks):
+                progress_bar(1, desc="[4/4]Generating QAs")
         except Exception as e: # pylint: disable=broad-except
             logger.error("Error occurred while generating QA: %s", e)
     return results
@@ -513,9 +517,11 @@ async def _process_single_batch(
         desc="[4/4]Generating QAs"
     ):
         try:
-            results.update(await result)
             if progress_bar is not None:
                 progress_bar(len(results) / len(processing_batches), desc="[4/4]Generating QAs")
+            results.update(await result)
+            if progress_bar is not None and len(results) == len(processing_batches):
+                progress_bar(1, desc="[4/4]Generating QAs")
         except Exception as e: # pylint: disable=broad-except
             logger.error("Error occurred while generating QA: %s", e)
     return results