"""FineWeb-Edu -> GPT-2 BPE uint16 shards (train.bin / val.bin).
  python scripts/prepare_data.py --tokens 3e9 --out data/fineweb [--dataset HuggingFaceFW/fineweb-edu --name sample-10BT]"""
import os, sys, argparse
import numpy as np

p = argparse.ArgumentParser()
p.add_argument("--dataset", default="HuggingFaceFW/fineweb-edu")
p.add_argument("--name", default="sample-10BT")
p.add_argument("--tokens", type=float, default=3e9)
p.add_argument("--val_tokens", type=float, default=5e6)
p.add_argument("--out", default="data/fineweb")
a = p.parse_args()
os.makedirs(a.out, exist_ok=True)
import tiktoken
from datasets import load_dataset
enc = tiktoken.get_encoding("gpt2")
ds = load_dataset(a.dataset, name=a.name, split="train", streaming=True)
train_path, val_path = os.path.join(a.out, "train.bin"), os.path.join(a.out, "val.bin")
# write to .partial and rename on completion, so an interrupted prep never leaves a
# half-written train.bin that run_ladder.sh's existence check would mistake for done
ftr, fva = open(train_path + ".partial", "wb"), open(val_path + ".partial", "wb")
n_tr = n_va = 0
target_tr, target_va = int(a.tokens), int(a.val_tokens)
buf = []
for i, ex in enumerate(ds):
    ids = enc.encode_ordinary(ex["text"]) + [enc.eot_token]
    arr = np.array(ids, dtype=np.uint16)
    if n_va < target_va and i % 100 == 0:          # every 100th doc to val until filled
        fva.write(arr.tobytes()); n_va += len(arr)
    else:
        ftr.write(arr.tobytes()); n_tr += len(arr)
    if n_tr % 50_000_000 < len(arr):
        print(f"train {n_tr/1e6:.0f}M  val {n_va/1e6:.1f}M tokens", flush=True)
    if n_tr >= target_tr and n_va >= target_va:
        break
ftr.close(); fva.close()
os.rename(train_path + ".partial", train_path); os.rename(val_path + ".partial", val_path)
print(f"DONE train {n_tr/1e6:.1f}M val {n_va/1e6:.1f}M -> {a.out}", flush=True)
sys.stdout.flush(); sys.stderr.flush()
os._exit(0)  # hard-exit: HF streaming leaves non-daemon prefetch threads after `break` that hang interpreter shutdown