Use ASTERIZER/Luna_Dataset: auto-extract zip, update commands

README.md

@@ -14,9 +14,11 @@ pip install -q -r requirements.txt
 
 ### 2. Get Dataset + Train (one command)
 
+The dataset (~4.5B tokens) is hosted as a zip at [ASTERIZER/Luna_Dataset](https://huggingface.co/datasets/ASTERIZER/Luna_Dataset). The script downloads, extracts, and starts training automatically.
+
 **From HuggingFace (recommended):**
 ```bash
-bash setup_and_train.sh huggingface ASTERIZER/LUNA-pretrain-data
+bash setup_and_train.sh huggingface ASTERIZER/Luna_Dataset
 ```
 
 **From Google Drive:**
@@ -26,7 +28,7 @@ bash setup_and_train.sh gdrive YOUR_GDRIVE_FOLDER_ID
 
 **Smoke test (10M tokens only):**
 ```bash
-bash setup_and_train.sh huggingface ASTERIZER/LUNA-pretrain-data 10000000
+bash setup_and_train.sh huggingface ASTERIZER/Luna_Dataset 10000000
 ```
 
 That's it. The script auto-detects your GPU, VRAM, RAM, CPU cores and configures everything for maximum utilization.

fetch_data.py

@@ -25,7 +25,7 @@ from pathlib import Path
 
 def download_huggingface(repo_id: str, out_dir: Path, hf_token: str = None):
     try:
-        from huggingface_hub import snapshot_download, hf_hub_download
+        from huggingface_hub import snapshot_download
     except ImportError:
         print("  Installing huggingface_hub...")
         os.system(f"{sys.executable} -m pip install -q huggingface_hub")
@@ -41,6 +41,13 @@ def download_huggingface(repo_id: str, out_dir: Path, hf_token: str = None):
         ignore_patterns=["*.md", ".gitattributes"],
     )
     print(f"  Downloaded to: {out_dir}")
+
+    # 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)
+
     _verify(out_dir)
 
 
@@ -93,6 +100,43 @@ def copy_local(local_path: str, out_dir: Path):
     _verify(out_dir)
 
 
+# ─── Zip Extraction & Flattening ──────────────────────────────────────────────
+
+def _extract_zips(data_dir: Path):
+    """Find and extract all .zip files in data_dir, then delete the zips."""
+    import zipfile
+    zips = list(data_dir.glob("*.zip"))
+    if not zips:
+        return
+    for zf in zips:
+        print(f"  Extracting {zf.name} ...")
+        with zipfile.ZipFile(zf) as z:
+            z.extractall(data_dir)
+        zf.unlink()
+        print(f"  Removed {zf.name}")
+
+
+def _flatten_to_root(data_dir: Path):
+    """If index.json is one level deep (e.g. data_dir/subfolder/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"))
+    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}/ ...")
+    for item in sub.iterdir():
+        dest = data_dir / item.name
+        if dest.exists():
+            if dest.is_dir():
+                shutil.rmtree(dest)
+            else:
+                dest.unlink()
+        shutil.move(str(item), str(dest))
+    sub.rmdir()
+
+
 # ─── Verify ───────────────────────────────────────────────────────────────────
 
 def _verify(data_dir: Path):

setup_and_train.sh

@@ -11,10 +11,10 @@
 #    bash setup_and_train.sh gdrive 1AbCdEfGhIjKlMnOpQrStUvWx
 #
 #    # Full dataset from HuggingFace:
-#    bash setup_and_train.sh huggingface YourName/LUNA-pretrain-data
+#    bash setup_and_train.sh huggingface ASTERIZER/Luna_Dataset
 #
 #    # Quick smoke test (10M tokens only):
-#    bash setup_and_train.sh gdrive 1AbCdEfGhIjKlMnOpQrStUvWx 10000000
+#    bash setup_and_train.sh huggingface ASTERIZER/Luna_Dataset 10000000
 #
 #    # Dataset already on disk:
 #    bash setup_and_train.sh local /workspace/data/litdata_pretrain_final