Upload fetch_data.py with huggingface_hub

fetch_data.py

@@ -45,8 +45,8 @@ def download_huggingface(repo_id: str, out_dir: Path, hf_token: str = None):
     # Auto-extract any zip files found in the download
     _extract_zips(out_dir)
 
-    # If index.json landed in a subdirectory, move contents up
-    _flatten_to_root(out_dir)
+    # If index.json landed in a subdirectory, move the intended dataset up
+    _flatten_to_root(out_dir, preferred_name=out_dir.name)
 
     _verify(out_dir)
 
@@ -79,6 +79,7 @@ def download_gdrive(gdrive_id: str, out_dir: Path):
                 z.extractall(out_dir)
             dest.unlink()
     print(f"  Downloaded to: {out_dir}")
+    _flatten_to_root(out_dir, preferred_name=out_dir.name)
     _verify(out_dir)
 
 
@@ -90,6 +91,7 @@ def copy_local(local_path: str, out_dir: Path):
         raise FileNotFoundError(f"Local path not found: {src}")
     if out_dir.resolve() == src.resolve():
         print(f"  Source == destination, no copy needed.")
+        _flatten_to_root(out_dir, preferred_name=out_dir.name)
         _verify(out_dir)
         return
     print(f"  Copying {src} → {out_dir}")
@@ -97,6 +99,7 @@ def copy_local(local_path: str, out_dir: Path):
         shutil.rmtree(out_dir)
     shutil.copytree(src, out_dir)
     print(f"  Copied to: {out_dir}")
+    _flatten_to_root(out_dir, preferred_name=out_dir.name)
     _verify(out_dir)
 
 
@@ -116,15 +119,62 @@ def _extract_zips(data_dir: Path):
         print(f"  Removed {zf.name}")
 
 
-def _flatten_to_root(data_dir: Path):
-    """If index.json is nested (e.g. data_dir/a/b/index.json),
-    move everything from that subfolder up to data_dir."""
+def _read_index_summary(index_path: Path):
+    try:
+        with open(index_path, encoding="utf-8") as f:
+            idx = json.load(f)
+    except Exception:
+        return {"chunks": 0, "tokens": 0}
+    chunks = idx.get("chunks", [])
+    total_tokens = sum(c.get("dim", 0) for c in chunks)
+    return {"chunks": len(chunks), "tokens": total_tokens}
+
+
+def _pick_dataset_subdir(candidates, preferred_name: str | None = None):
+    """Pick the most likely dataset subdir from multiple nested index.json files.
+
+    Priority:
+    1. Parent folder name exactly matches preferred_name
+    2. Parent folder name contains preferred_name
+    3. Highest token count
+    """
+    ranked = []
+    preferred_name = (preferred_name or "").lower()
+
+    for index_path in candidates:
+        parent = index_path.parent
+        summary = _read_index_summary(index_path)
+        name = parent.name.lower()
+        exact = int(bool(preferred_name and name == preferred_name))
+        contains = int(bool(preferred_name and preferred_name in name))
+        ranked.append((exact, contains, summary["tokens"], summary["chunks"], parent, summary))
+
+    ranked.sort(reverse=True, key=lambda item: (item[0], item[1], item[2], item[3]))
+    return ranked[0][4], ranked[0][5]
+
+
+def _flatten_to_root(data_dir: Path, preferred_name: str | None = None):
+    """If index.json is nested, move the intended dataset subfolder up to data_dir."""
     if (data_dir / "index.json").exists():
-        return  # already at root
+        return
+
     candidates = list(data_dir.glob("**/index.json"))
-    if len(candidates) != 1:
-        return  # ambiguous or not found
-    sub = candidates[0].parent
+    if not candidates:
+        return
+
+    if len(candidates) == 1:
+        sub = candidates[0].parent
+        summary = _read_index_summary(candidates[0])
+    else:
+        sub, summary = _pick_dataset_subdir(candidates, preferred_name=preferred_name)
+        print(
+            f"  Multiple nested datasets found; selected {sub.relative_to(data_dir)}/ "
+            f"({summary['chunks']} chunks, {summary['tokens']:,} tokens)"
+        )
+
+    if sub == data_dir:
+        return
+
     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