File size: 2,859 Bytes
4397e12 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | """Fetch raw pretraining shards into $TA_DATA/raw. Idempotent; re-run to resume.
Sources (all permissive / openly licensed):
english: HuggingFaceFW/fineweb-edu sample/10BT (ODC-By) ~0.7B tok/file
math: HuggingFaceTB/finemath finemath-4plus (ODC-By) ~0.15B tok/file
code: common-pile/stackv2_edu_filtered (openly licensed code only)
"""
import argparse, os
from huggingface_hub import hf_hub_download
DATA = os.environ.get("TA_DATA", "data")
SOURCES = {
"english": ("HuggingFaceFW/fineweb-edu", [f"sample/10BT/{i:03d}_00000.parquet" for i in range(14)]),
"math": ("HuggingFaceTB/finemath", [f"finemath-4plus/train-{i:05d}-of-00064.parquet" for i in range(64)]),
"code": ("common-pile/stackv2_edu_filtered", None),
}
# stack-edu shards are sorted by language (checked 2026-10-06 by reading each shard's first rows):
# 0-17 Markdown, 18-19 C, 20-27 C++, 28-33 C#, 34-37 Go, 38-49 Java, 50-64 JavaScript, 65-72 PHP,
# 73-84 Python, 85 Ruby, 86 Rust, 87-88 Shell, 89-90 Swift, 91-94 TypeScript.
CODE_SHARDS = {"markdown": [0, 1], "python": [73, 74, 75, 76], "shell": [87, 88], "javascript": [50],
"typescript": [91], "go": [34], "rust": [86], "c": [18]}
# Small supervised sets rendered into the synthetic sources (scripts/gen_synth.py):
# OpenMathInstruct-2 (CC-BY-4.0, solutions written by Llama-3.1-405B) and GSM8K (MIT) -> synth_reasoning
# SQuAD v2 and HotpotQA (CC-BY-SA-4.0) -> synth_grounded (copy-the-answer reading comprehension)
EXTRA = [("nvidia/OpenMathInstruct-2", f"data/train-{i:05d}-of-00032.parquet") for i in range(2)] + [
("openai/gsm8k", "main/train-00000-of-00001.parquet"),
("rajpurkar/squad_v2", "squad_v2/train-00000-of-00001.parquet"),
("hotpotqa/hotpot_qa", "distractor/train-00000-of-00002.parquet"),
]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--english", type=int, default=4)
ap.add_argument("--math", type=int, default=14)
a = ap.parse_args()
for lang, shards in CODE_SHARDS.items():
for i in shards:
p = hf_hub_download(SOURCES["code"][0], f"stack-edu-{i:04d}.json.gz", repo_type="dataset",
local_dir=f"{DATA}/raw/code")
print("code", lang, p, os.path.getsize(p) // 2**20, "MiB", flush=True)
for repo, f in EXTRA:
p = hf_hub_download(repo, f, repo_type="dataset", local_dir=f"{DATA}/raw/extra/{repo.split('/')[1]}")
print("extra", p, os.path.getsize(p) // 2**20, "MiB", flush=True)
for name, n in [("math", a.math), ("english", a.english)]:
repo, files = SOURCES[name]
for f in files[:n]:
p = hf_hub_download(repo, f, repo_type="dataset", local_dir=f"{DATA}/raw/{name}")
print(name, p, os.path.getsize(p) // 2**20, "MiB", flush=True)
if __name__ == "__main__":
main()
|