"""Training data for the paper: FineWeb-Edu (sample-10BT), Gemma tokenizer (262,144 tokens), uint32 shards.
2.05B training tokens + 10M validation tokens, documents separated by <eos>. Also writes the token counts used
for the unigram prior (counted on the training tokens only).

    python prep_data.py --out data [--train 2.05e9] [--val 1e7]
"""
import argparse, pathlib, time
import numpy as np
from datasets import load_dataset
from transformers import AutoTokenizer

SHARD = 100_000_000


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="data")
    ap.add_argument("--train", type=float, default=2.05e9)
    ap.add_argument("--val", type=float, default=1e7)
    a = ap.parse_args()
    out = pathlib.Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    tok = AutoTokenizer.from_pretrained("google/gemma-3-270m")
    eos = tok.eos_token_id
    need_val, need_train = int(a.val), int(a.train)
    ds = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT", split="train", streaming=True)
    counts = np.zeros(262144, np.int64)
    buf = np.empty(SHARD + 1_000_000, np.uint32)
    n_buf, shard, done_val, total, t0, batch = 0, 0, False, 0, time.time(), []

    def write(arr, name):
        np.save(out / name, arr)
        print(time.strftime("%H:%M:%S"), f"wrote {name} ({len(arr) / 1e6:.0f}M tokens)", flush=True)

    for ex in ds:
        batch.append(ex["text"])
        if len(batch) < 512:
            continue
        for ids in tok(batch, add_special_tokens=False)["input_ids"]:
            ids.append(eos)
            buf[n_buf:n_buf + len(ids)] = ids
            n_buf += len(ids)
        batch = []
        if not done_val and n_buf >= need_val:
            write(buf[:need_val].copy(), "val.npy")
            rest = buf[need_val:n_buf].copy()
            buf[:len(rest)] = rest
            n_buf, done_val = len(rest), True
        elif done_val and n_buf >= SHARD:
            take = min(SHARD, need_train - total)
            arr = buf[:take].copy()
            counts += np.bincount(arr, minlength=262144)
            write(arr, f"train_{shard:03d}.npy")
            shard, total = shard + 1, total + take
            rest = buf[SHARD:n_buf].copy()
            buf[:len(rest)] = rest
            n_buf = len(rest)
            if total >= need_train:
                break
            print(f"   {total / 1e9:.2f}B / {need_train / 1e9:.2f}B  {total / (time.time() - t0) / 1e6:.2f}M tok/s", flush=True)
    if total < need_train and n_buf:
        arr = buf[:min(n_buf, need_train - total)].copy()
        counts += np.bincount(arr, minlength=262144)
        write(arr, f"train_{shard:03d}.npy")
        total += len(arr)
    np.save(out / "unigram_counts.npy", counts)
    print(time.strftime("%H:%M:%S"), f"done: {total / 1e9:.3f}B train tokens in {time.time() - t0:.0f} s", flush=True)


if __name__ == "__main__":
    main()