From 3b6f7af360e886890d971d752d12a84fb585f00c Mon Sep 17 00:00:00 2001
From: yichuan520030910320 <yichuan_wang@berkeley.edu>
Date: Sat, 27 Jun 2026 06:12:00 -0700
Subject: [PATCH] feat(embed): add tqdm progress bar to chunk stage, clean up
 embed logging
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- chunk.py: add tqdm progress bar for article-level chunking so users
  see real-time progress during Stage 2 (e.g. "Chunking: 50%|█████ | 4/8")
- embed_cpu.py: remove redundant "Embedded X/Y" logger.info every 10
  items — tqdm already shows this and the log lines break the progress
  bar display

Tested with 8 Wikipedia pages (64 chunks produced).
---
 embed/src/pixelrag_embed/chunk.py     | 35 ++++++++++++++++-----------
 embed/src/pixelrag_embed/embed_cpu.py |  3 ---
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/embed/src/pixelrag_embed/chunk.py b/embed/src/pixelrag_embed/chunk.py
index c8f9f018..bd019702 100644
--- a/embed/src/pixelrag_embed/chunk.py
+++ b/embed/src/pixelrag_embed/chunk.py
@@ -36,6 +36,7 @@
 from pathlib import Path
 
 from PIL import Image
+from tqdm import tqdm
 
 Image.MAX_IMAGE_PIXELS = None  # some tiles exceed default 178M pixel limit
 
@@ -304,21 +305,27 @@ def process_shard(
         # Flat structure — article dirs directly in shard_dir
         sub_dirs = [Path(shard_dir)]
 
-    for sub_dir in sub_dirs:
-        for article_dir in sorted(sub_dir.iterdir()):
-            if not article_dir.is_dir() or not article_dir.name.endswith(".png.tiles"):
-                continue
-            total_articles += 1
-
-            result = chunk_article(str(article_dir), dry_run=dry_run, force=force)
-            if result is None:
-                skipped_articles += 1
-                continue
+    all_article_dirs = [
+        article_dir
+        for sub_dir in sub_dirs
+        for article_dir in sorted(sub_dir.iterdir())
+        if article_dir.is_dir() and article_dir.name.endswith(".png.tiles")
+    ]
+
+    for article_dir in tqdm(
+        all_article_dirs, desc="Chunking", disable=not all_article_dirs
+    ):
+        total_articles += 1
+
+        result = chunk_article(str(article_dir), dry_run=dry_run, force=force)
+        if result is None:
+            skipped_articles += 1
+            continue
 
-            chunked_articles += 1
-            total_tiles += result["num_tiles"]
-            total_chunks += result["num_chunks"]
-            total_files += result["files_written"]
+        chunked_articles += 1
+        total_tiles += result["num_tiles"]
+        total_chunks += result["num_chunks"]
+        total_files += result["files_written"]
 
     # Delete tiles after chunking the whole shard
     tiles_deleted = 0
diff --git a/embed/src/pixelrag_embed/embed_cpu.py b/embed/src/pixelrag_embed/embed_cpu.py
index 3b717af1..39b40248 100644
--- a/embed/src/pixelrag_embed/embed_cpu.py
+++ b/embed/src/pixelrag_embed/embed_cpu.py
@@ -187,9 +187,6 @@ def embed_items(
             pooled = pooled / pooled.norm()
             embeddings[i] = pooled.cpu().numpy().astype(np.float16)
 
-        if (i + 1) % 10 == 0:
-            logger.info("Embedded %d/%d", i + 1, len(items))
-
     return embeddings