post_merge fix with reset option

post_merge.py

@@ -3,9 +3,10 @@
 import io
 import json
 import logging
+import os
 from datetime import datetime, timezone
 
-from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub import HfApi, hf_hub_download, snapshot_download
 from huggingface_hub.utils import EntryNotFoundError
 
 from dedup import DATASET_REPO_ID, compute_fingerprint, compute_sha256
@@ -131,54 +132,54 @@ def update_dataset_card(api: HfApi) -> None:
 def full_rebuild() -> dict:
     """Rebuild manifest.json from scratch and regenerate the dataset card.
 
-    Scans ALL data files on main, recomputes every hash/fingerprint,
-    and writes a fresh manifest (replacing the old one entirely).
+    Downloads the entire dataset at once via snapshot_download, then walks
+    the local directory to compute hashes/fingerprints for all data files.
     """
     logger.info("Starting full rebuild of manifest + dataset card")
 
-    all_data_files: list[str] = []
-    for entry in api.list_repo_tree(
+    # Download entire dataset in one shot
+    local_dir = snapshot_download(
         repo_id=DATASET_REPO_ID,
         repo_type="dataset",
         revision="main",
-        recursive=True,
-    ):
-        if not hasattr(entry, "rfilename"):
-            continue
-        path = entry.rfilename
-        if path.startswith("data/") and (path.endswith(".json") or path.endswith(".jsonl")):
-            all_data_files.append(path)
-
-    logger.info("Found %d data files to index", len(all_data_files))
+    )
+    logger.info("Downloaded dataset snapshot to %s", local_dir)
 
+    # Walk local data/ directory to find all data files
+    data_root = os.path.join(local_dir, "data")
     now = datetime.now(timezone.utc).isoformat()
     manifest = {"files": {}}
 
-    for file_path in all_data_files:
-        try:
-            local_path = hf_hub_download(
-                repo_id=DATASET_REPO_ID,
-                filename=file_path,
-                repo_type="dataset",
-                revision="main",
-            )
-            with open(local_path, "rb") as f:
-                content = f.read()
-
-            sha256 = compute_sha256(content)
-            if file_path.endswith(".json"):
-                fingerprint = compute_fingerprint(content)
-            else:
-                fingerprint = sha256
-
-            manifest["files"][file_path] = {
-                "sha256": sha256,
-                "fingerprint": fingerprint,
-                "added_at": now,
-            }
-            logger.info("Indexed %s", file_path)
-        except Exception:
-            logger.exception("Failed to index %s", file_path)
+    if not os.path.isdir(data_root):
+        logger.warning("No data/ directory found in snapshot")
+    else:
+        for dirpath, _, filenames in os.walk(data_root):
+            for filename in filenames:
+                if not (filename.endswith(".json") or filename.endswith(".jsonl")):
+                    continue
+                local_path = os.path.join(dirpath, filename)
+                # Convert to repo-relative path (data/...)
+                repo_path = os.path.relpath(local_path, local_dir)
+                try:
+                    with open(local_path, "rb") as f:
+                        content = f.read()
+
+                    sha256 = compute_sha256(content)
+                    if filename.endswith(".json"):
+                        fingerprint = compute_fingerprint(content)
+                    else:
+                        fingerprint = sha256
+
+                    manifest["files"][repo_path] = {
+                        "sha256": sha256,
+                        "fingerprint": fingerprint,
+                        "added_at": now,
+                    }
+                    logger.info("Indexed %s", repo_path)
+                except Exception:
+                    logger.exception("Failed to index %s", repo_path)
+
+    logger.info("Indexed %d data files total", len(manifest["files"]))
 
     # Upload fresh manifest
     manifest_bytes = json.dumps(manifest, indent=2, sort_keys=True).encode()