| # SAGE — 5 Billion Token Dataset Downloader | |
| Automatically downloads ~5B tokens from free, public Hugging Face datasets and saves them as JSONL files in your `data/raw/` directory, fully compatible with the SAGE training pipeline. | |
| --- | |
| ## Token Budget | |
| | File | Source | Tokens | | |
| |---|---|---| | |
| | `general_web.jsonl` | FineWeb | 2.5B | | |
| | `code.jsonl` | The Stack v2 (Python, JS, Rust, Go, C++ and more) | 1.0B | | |
| | `math_science.jsonl` | OpenWebMath | 0.5B | | |
| | `multilingual.jsonl` | Wikipedia (20+ languages) | 0.5B | | |
| | `synthetic.jsonl` | OpenHermes 2.5 (instruction data) | 0.5B | | |
| | **Total** | | **~5.0B tokens** | | |
| **Estimated disk space:** ~20–25 GB | |
| **Estimated download time:** 2–8 hours depending on your connection | |
| **Cost:** 100% free, no account required | |
| --- | |
| ## Requirements | |
| ### System | |
| - Python 3.9+ | |
| - 25 GB free disk space | |
| - Stable internet connection | |
| ### Python packages | |
| ```bash | |
| pip install datasets huggingface_hub tqdm | |
| ``` | |
| --- | |
| ## Usage | |
| ### Basic — download everything | |
| ```bash | |
| python debug/download_5b_tokens.py --output-dir data/raw | |
| ``` | |
| ### Test run — 1% of data to verify everything works | |
| ```bash | |
| python debug/download_5b_tokens.py --output-dir data/raw --scale 0.01 | |
| ``` | |
| ### Resume — continue after an internet cutout | |
| ```bash | |
| python debug/download_5b_tokens.py --output-dir data/raw --resume | |
| ``` | |
| ### Download only one specific file | |
| ```bash | |
| python debug/download_5b_tokens.py --output-dir data/raw --only code.jsonl | |
| ``` | |
| ### Download multiple specific files | |
| ```bash | |
| python debug/download_5b_tokens.py --output-dir data/raw --only code.jsonl math_science.jsonl | |
| ``` | |
| --- | |
| ## All Flags | |
| | Flag | Default | Description | | |
| |---|---|---| | |
| | `--output-dir` | `data/raw` | Directory where JSONL files are saved | | |
| | `--resume` | off | Skip files that already hit their token target | | |
| | `--only` | all files | Download only the specified file(s) | | |
| | `--scale` | `1.0` | Scale all token targets (e.g. `0.1` = 10% of 5B = 500M tokens) | | |
| --- | |
| ## Output Format | |
| Every record written to disk follows this structure with at minimum a `text` field, making it directly compatible with the SAGE pipeline: | |
| ```json | |
| { "text": "your training sample here", "source": "fineweb", "language": "en" } | |
| ``` | |
| --- | |
| ## Data Sources | |
| ### 1. FineWeb — `general_web.jsonl` | |
| - **Dataset:** `HuggingFaceFW/fineweb` (sample-10BT subset) | |
| - **What it is:** A pre-shuffled, deduplicated 10B-token slice of web-crawl text, one of the cleanest freely available web datasets | |
| - **Why it's used:** Broad general language coverage, essential for fluent text generation | |
| ### 2. The Stack v2 — `code.jsonl` | |
| - **Dataset:** `bigcode/the-stack-v2-train-smol-ids` | |
| - **What it is:** Source code across 10 programming languages: Python, JavaScript, TypeScript, Rust, Go, C++, Java, Bash, SQL, and HTML | |
| - **Why it's used:** Teaches the model programming syntax, logic, and structure | |
| ### 3. OpenWebMath — `math_science.jsonl` | |
| - **Dataset:** `open-web-math/open-web-math` | |
| - **What it is:** 14.7B tokens of mathematical content extracted from the web, including LaTeX, proofs, and problem sets | |
| - **Why it's used:** Improves numerical reasoning and scientific language understanding | |
| ### 4. Wikipedia — `multilingual.jsonl` | |
| - **Dataset:** `wikimedia/wikipedia` (20231101 dumps) | |
| - **Languages:** English, Spanish, French, German, Chinese, Japanese, Portuguese, Arabic, Russian, Hindi, Italian, Korean, Dutch, Polish, Swedish, Turkish, Vietnamese, Indonesian, Ukrainian, Persian | |
| - **Why it's used:** Clean, factual, encyclopedic text across 20 languages | |
| ### 5. OpenHermes 2.5 — `synthetic.jsonl` | |
| - **Dataset:** `teknium/OpenHermes-2.5` | |
| - **What it is:** ~1M high-quality instruction-following pairs formatted as `### Instruction` / `### Response` conversations | |
| - **Why it's used:** Teaches the model to follow instructions and produce structured, helpful responses | |
| --- | |
| ## What Happens After Download | |
| Once all files are ready, continue with the standard SAGE pipeline: | |
| ### Train the tokenizer | |
| ```bash | |
| python -m tokenizer.train_tokenizer \ | |
| --input data/raw/general_web.jsonl \ | |
| data/raw/code.jsonl \ | |
| data/raw/math_science.jsonl \ | |
| data/raw/multilingual.jsonl \ | |
| data/raw/synthetic.jsonl \ | |
| --model-prefix tokenizer/tokenizer \ | |
| --vocab-size 32000 | |
| ``` | |
| ### Build parquet shards | |
| ```bash | |
| python -m data.pipeline \ | |
| --tokenizer-model tokenizer/tokenizer.model \ | |
| --output-dir data/processed \ | |
| --shard-size 128 | |
| ``` | |
| ### Start training | |
| ```bash | |
| python -m train.trainer \ | |
| --model-config configs/model/1b.yaml \ | |
| --schedule-config configs/train/schedule.yaml \ | |
| --train-shards data/processed/shard-00000.parquet \ | |
| --validation-shards data/processed/shard-00001.parquet \ | |
| --output-dir runs/sage-1b | |
| ``` | |
| --- | |
| ## Troubleshooting | |
| **Download stalls or disconnects** | |
| Run with `--resume` to pick up exactly where you left off. The writer appends to existing files and counts already-written tokens before continuing. | |
| **A specific language or dataset fails** | |
| The downloader catches errors per-source and logs a warning, then moves on. The other files are unaffected. Re-run with `--only <filename> --resume` to retry just that file. | |
| **Running out of disk space mid-download** | |
| Use `--scale 0.5` to target 2.5B tokens total (~10–12 GB) instead of the full 5B. The model will be slightly less capable but the pipeline will still work end to end. | |
| **Slow download speed** | |
| All datasets are streamed — data is downloaded and written record by record, so you never need to load the entire dataset at once. If speed is consistently low, try running overnight or on a cloud VM closer to Hugging Face's CDN. | |
| --- | |
| ## Full Script | |
| ```python | |
| """ | |
| SAGE — 5 Billion Token Dataset Downloader | |
| ========================================== | |
| Downloads ~5B tokens from free Hugging Face datasets and saves them | |
| as JSONL files in your data/raw/ directory, ready for the SAGE pipeline. | |
| Token budget breakdown: | |
| general_web.jsonl → 2.5B tokens (FineWeb) | |
| code.jsonl → 1.0B tokens (The Stack v2 - Python, JS, Rust, Go, C++) | |
| math_science.jsonl → 0.5B tokens (OpenWebMath) | |
| multilingual.jsonl → 0.5B tokens (Wikipedia 20+ languages) | |
| synthetic.jsonl → 0.5B tokens (OpenHermes instruction data) | |
| ───────────────────────────────────── | |
| TOTAL → ~5.0B tokens | |
| Usage: | |
| pip install datasets huggingface_hub tqdm | |
| python debug/download_5b_tokens.py --output-dir data/raw | |
| python debug/download_5b_tokens.py --output-dir data/raw --resume | |
| """ | |
| import argparse | |
| import json | |
| import sys | |
| import time | |
| from pathlib import Path | |
| missing = [] | |
| try: | |
| from datasets import load_dataset | |
| except ImportError: | |
| missing.append("datasets") | |
| try: | |
| from tqdm import tqdm | |
| except ImportError: | |
| missing.append("tqdm") | |
| if missing: | |
| print(f"[ERROR] Missing packages: {', '.join(missing)}") | |
| print(f" Run: pip install {' '.join(missing)}") | |
| sys.exit(1) | |
| def estimate_tokens(text: str) -> int: | |
| return max(1, len(text) // 4) | |
| def human_tokens(n: int) -> str: | |
| if n >= 1_000_000_000: | |
| return f"{n/1_000_000_000:.2f}B" | |
| if n >= 1_000_000: | |
| return f"{n/1_000_000:.1f}M" | |
| return f"{n:,}" | |
| def human_bytes(n: int) -> str: | |
| for unit in ["B", "KB", "MB", "GB"]: | |
| if n < 1024: | |
| return f"{n:.1f} {unit}" | |
| n /= 1024 | |
| return f"{n:.1f} TB" | |
| class JSONLWriter: | |
| def __init__(self, path: Path, target_tokens: int, resume: bool = False): | |
| self.path = path | |
| self.target_tokens = target_tokens | |
| self.tokens_written = 0 | |
| self.records_written = 0 | |
| if resume and path.exists(): | |
| print(f" [resume] Counting existing tokens in {path.name}...") | |
| with open(path, "r", encoding="utf-8") as f: | |
| for line in f: | |
| try: | |
| rec = json.loads(line) | |
| self.tokens_written += estimate_tokens(rec.get("text", "")) | |
| self.records_written += 1 | |
| except json.JSONDecodeError: | |
| pass | |
| print(f" [resume] Already have {human_tokens(self.tokens_written)} / {human_tokens(target_tokens)}") | |
| self._file = open(path, "a", encoding="utf-8", buffering=1024 * 1024) | |
| else: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| self._file = open(path, "w", encoding="utf-8", buffering=1024 * 1024) | |
| @property | |
| def done(self) -> bool: | |
| return self.tokens_written >= self.target_tokens | |
| def write(self, record: dict) -> int: | |
| text = record.get("text", "") | |
| if not text or len(text.strip()) < 50: | |
| return 0 | |
| toks = estimate_tokens(text) | |
| self._file.write(json.dumps(record, ensure_ascii=False) + "\n") | |
| self.tokens_written += toks | |
| self.records_written += 1 | |
| return toks | |
| def close(self): | |
| self._file.flush() | |
| self._file.close() | |
| def __enter__(self): return self | |
| def __exit__(self, *_): self.close() | |
| def download_general_web(writer): | |
| print("\n[1/5] general_web.jsonl — FineWeb") | |
| bar = tqdm(total=writer.target_tokens, initial=writer.tokens_written, | |
| unit="tok", unit_scale=True, desc=" web tokens") | |
| ds = load_dataset("HuggingFaceFW/fineweb", name="sample-10BT", | |
| split="train", streaming=True) | |
| for sample in ds: | |
| if writer.done: break | |
| bar.update(writer.write({"text": sample["text"], "source": "fineweb", | |
| "url": sample.get("url", ""), "language": "en"})) | |
| bar.close() | |
| print(f" ✓ {human_tokens(writer.tokens_written)} tokens | {writer.records_written:,} records") | |
| def download_code(writer): | |
| print("\n[2/5] code.jsonl — The Stack v2") | |
| LANGUAGES = [("python","Python"),("javascript","JavaScript"),("typescript","TypeScript"), | |
| ("rust","Rust"),("go","Go"),("cpp","C++"),("java","Java"), | |
| ("bash","Bash"),("sql","SQL"),("html","HTML")] | |
| bar = tqdm(total=writer.target_tokens, initial=writer.tokens_written, | |
| unit="tok", unit_scale=True, desc=" code tokens") | |
| tokens_per_lang = writer.target_tokens // len(LANGUAGES) | |
| for lang_id, lang_name in LANGUAGES: | |
| if writer.done: break | |
| lang_tokens = 0 | |
| print(f" → {lang_name}...") | |
| try: | |
| ds = load_dataset("bigcode/the-stack-v2-train-smol-ids", | |
| data_dir=f"data/{lang_id}", split="train", | |
| streaming=True, trust_remote_code=True) | |
| for sample in ds: | |
| if writer.done or lang_tokens >= tokens_per_lang: break | |
| content = sample.get("content", "") or sample.get("text", "") | |
| if not content: continue | |
| t = writer.write({"text": content, "source": "the_stack_v2", | |
| "language": lang_id}) | |
| bar.update(t); lang_tokens += t | |
| except Exception as e: | |
| print(f" [warn] {lang_name} failed ({e}), skipping.") | |
| bar.close() | |
| print(f" ✓ {human_tokens(writer.tokens_written)} tokens | {writer.records_written:,} records") | |
| def download_math(writer): | |
| print("\n[3/5] math_science.jsonl — OpenWebMath") | |
| bar = tqdm(total=writer.target_tokens, initial=writer.tokens_written, | |
| unit="tok", unit_scale=True, desc=" math tokens") | |
| ds = load_dataset("open-web-math/open-web-math", split="train", streaming=True) | |
| for sample in ds: | |
| if writer.done: break | |
| bar.update(writer.write({"text": sample["text"], "source": "open_web_math", | |
| "url": sample.get("url", "")})) | |
| bar.close() | |
| print(f" ✓ {human_tokens(writer.tokens_written)} tokens | {writer.records_written:,} records") | |
| def download_multilingual(writer): | |
| print("\n[4/5] multilingual.jsonl — Wikipedia (20 languages)") | |
| LANGUAGES = [("en","English"),("es","Spanish"),("fr","French"),("de","German"), | |
| ("zh","Chinese"),("ja","Japanese"),("pt","Portuguese"),("ar","Arabic"), | |
| ("ru","Russian"),("hi","Hindi"),("it","Italian"),("ko","Korean"), | |
| ("nl","Dutch"),("pl","Polish"),("sv","Swedish"),("tr","Turkish"), | |
| ("vi","Vietnamese"),("id","Indonesian"),("uk","Ukrainian"),("fa","Persian")] | |
| bar = tqdm(total=writer.target_tokens, initial=writer.tokens_written, | |
| unit="tok", unit_scale=True, desc=" multilingual tokens") | |
| tokens_per_lang = writer.target_tokens // len(LANGUAGES) | |
| for lang_code, lang_name in LANGUAGES: | |
| if writer.done: break | |
| lang_tokens = 0 | |
| try: | |
| ds = load_dataset("wikimedia/wikipedia", f"20231101.{lang_code}", | |
| split="train", streaming=True, trust_remote_code=True) | |
| for sample in ds: | |
| if writer.done or lang_tokens >= tokens_per_lang: break | |
| text = sample.get("text", "") | |
| if not text: continue | |
| t = writer.write({"text": text, "source": "wikipedia", | |
| "language": lang_code, "title": sample.get("title","")}) | |
| bar.update(t); lang_tokens += t | |
| except Exception as e: | |
| print(f"\n [warn] Wikipedia {lang_name} failed: {e}") | |
| bar.close() | |
| print(f" ✓ {human_tokens(writer.tokens_written)} tokens | {writer.records_written:,} records") | |
| def download_synthetic(writer): | |
| print("\n[5/5] synthetic.jsonl — OpenHermes 2.5") | |
| bar = tqdm(total=writer.target_tokens, initial=writer.tokens_written, | |
| unit="tok", unit_scale=True, desc=" synthetic tokens") | |
| ds = load_dataset("teknium/OpenHermes-2.5", split="train", streaming=True) | |
| rounds = 0 | |
| while not writer.done and rounds < 10: | |
| for sample in ds: | |
| if writer.done: break | |
| convs = sample.get("conversations", []) | |
| parts = [] | |
| for turn in convs: | |
| role, value = turn.get("from",""), turn.get("value","") | |
| if role == "human": parts.append(f"### Instruction\n{value}") | |
| elif role == "gpt": parts.append(f"### Response\n{value}") | |
| text = "\n\n".join(parts) or sample.get("text","") | |
| if not text: continue | |
| bar.update(writer.write({"text": text, "source": "openhermes_2.5", | |
| "task": "instruction_following"})) | |
| rounds += 1 | |
| bar.close() | |
| print(f" ✓ {human_tokens(writer.tokens_written)} tokens | {writer.records_written:,} records") | |
| TARGETS = { | |
| "general_web.jsonl": 2_500_000_000, | |
| "code.jsonl": 1_000_000_000, | |
| "math_science.jsonl": 500_000_000, | |
| "multilingual.jsonl": 500_000_000, | |
| "synthetic.jsonl": 500_000_000, | |
| } | |
| DOWNLOADERS = { | |
| "general_web.jsonl": download_general_web, | |
| "code.jsonl": download_code, | |
| "math_science.jsonl": download_math, | |
| "multilingual.jsonl": download_multilingual, | |
| "synthetic.jsonl": download_synthetic, | |
| } | |
| def main(): | |
| parser = argparse.ArgumentParser(description="Download ~5B tokens for SAGE training.") | |
| parser.add_argument("--output-dir", default="data/raw") | |
| parser.add_argument("--resume", action="store_true") | |
| parser.add_argument("--only", nargs="+", choices=list(TARGETS.keys())) | |
| parser.add_argument("--scale", type=float, default=1.0) | |
| args = parser.parse_args() | |
| out_dir = Path(args.output_dir) | |
| out_dir.mkdir(parents=True, exist_ok=True) | |
| files_to_run = args.only or list(TARGETS.keys()) | |
| total_target = sum(int(TARGETS[f] * args.scale) for f in files_to_run) | |
| print("=" * 60) | |
| print(" SAGE — 5 Billion Token Downloader") | |
| print("=" * 60) | |
| print(f" Output dir : {out_dir.resolve()}") | |
| print(f" Resume : {args.resume}") | |
| print(f" Scale : {args.scale}x") | |
| print(f" Target : {human_tokens(total_target)} tokens") | |
| print(f" Est. disk : ~{total_target // 40_000_000} GB") | |
| print("=" * 60) | |
| grand_start = time.time() | |
| grand_tokens = 0 | |
| for filename in files_to_run: | |
| target = int(TARGETS[filename] * args.scale) | |
| with JSONLWriter(out_dir / filename, target, resume=args.resume) as writer: | |
| if writer.done: | |
| print(f"\n[skip] {filename} already complete ({human_tokens(writer.tokens_written)} tokens)") | |
| grand_tokens += writer.tokens_written | |
| continue | |
| t0 = time.time() | |
| DOWNLOADERS[filename](writer) | |
| elapsed = time.time() - t0 | |
| grand_tokens += writer.tokens_written | |
| size = (out_dir / filename).stat().st_size | |
| print(f" Time: {elapsed/60:.1f} min | Size: {human_bytes(size)}") | |
| elapsed_total = time.time() - grand_start | |
| print("\n" + "=" * 60) | |
| print(f" DONE — {human_tokens(grand_tokens)} tokens downloaded") | |
| print(f" Total time: {elapsed_total/3600:.2f} hours") | |
| print(f" Files: {out_dir.resolve()}/") | |
| print("=" * 60) | |
| if __name__ == "__main__": | |
| main() | |
| ``` |