Fix: flatten deeply nested zip extracts (data/litdata_pretrain_final/)

fetch_data.py

@@ -117,15 +117,15 @@ def _extract_zips(data_dir: Path):
 
 
 def _flatten_to_root(data_dir: Path):
-    """If index.json is one level deep (e.g. data_dir/subfolder/index.json),
+    """If index.json is nested (e.g. data_dir/a/b/index.json),
     move everything from that subfolder up to data_dir."""
     if (data_dir / "index.json").exists():
         return  # already at root
-    candidates = list(data_dir.glob("*/index.json"))
+    candidates = list(data_dir.glob("**/index.json"))
     if len(candidates) != 1:
         return  # ambiguous or not found
     sub = candidates[0].parent
-    print(f"  Moving contents from {sub.name}/ up to {data_dir.name}/ ...")
+    print(f"  Moving contents from {sub.relative_to(data_dir)}/ up to {data_dir.name}/ ...")
     for item in sub.iterdir():
         dest = data_dir / item.name
         if dest.exists():
@@ -134,7 +134,14 @@ def _flatten_to_root(data_dir: Path):
             else:
                 dest.unlink()
         shutil.move(str(item), str(dest))
-    sub.rmdir()
+    # Remove the now-empty nested directories
+    # Walk up from sub to data_dir, removing empty dirs
+    while sub != data_dir:
+        try:
+            sub.rmdir()
+        except OSError:
+            break
+        sub = sub.parent
 
 
 # ─── Verify ───────────────────────────────────────────────────────────────────