"""Tokenize FineWeb-Edu (sample-10BT) with the GPT-2 tokenizer into flat uint16 files.

Each document is written as [EOT] + tokens. Train comes from shards 000-006
(~5B tokens), val from the first --val_tokens of shard 013.
"""

import argparse
from multiprocessing import Pool
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq
import tiktoken
from huggingface_hub import hf_hub_download

REPO = "HuggingFaceFW/fineweb-edu"
ENC = None


def encode(texts):
    global ENC
    if ENC is None:
        ENC = tiktoken.get_encoding("gpt2")
    out = []
    for ids in ENC.encode_ordinary_batch(texts):
        out.append(ENC.eot_token)
        out.extend(ids)
    return np.asarray(out, dtype=np.uint16)


def text_batches(path, batch=1024):
    for rb in pq.ParquetFile(path).iter_batches(batch_size=batch, columns=["text"]):
        yield rb.column("text").to_pylist()


def write_split(shards, dst, max_tokens, workers):
    n = 0
    with open(dst, "wb") as f, Pool(workers) as pool:
        for shard in shards:
            local = hf_hub_download(REPO, f"sample/10BT/{shard}_00000.parquet", repo_type="dataset")
            for arr in pool.imap(encode, text_batches(local), chunksize=4):
                if max_tokens and n + len(arr) > max_tokens:
                    arr = arr[: max_tokens - n]
                f.write(arr.tobytes())
                n += len(arr)
                if max_tokens and n >= max_tokens:
                    return n
            print(f"{dst.name}: {shard} done, {n / 1e9:.3f}B tokens", flush=True)
    return n


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--out_dir", default="/data/fineweb-edu-gpt2")
    p.add_argument("--train_shards", default="000,001,002,003,004,005,006")
    p.add_argument("--val_shard", default="013")
    p.add_argument("--val_tokens", type=int, default=20_000_000)
    p.add_argument("--workers", type=int, default=32)
    args = p.parse_args()
    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    nv = write_split([args.val_shard], out / "val.bin", args.val_tokens, args.workers)
    print(f"val: {nv} tokens", flush=True)
    nt = write_split(args.train_shards.split(","), out / "train.bin", 0, args.workers)
    print(f"train: {nt} tokens", flush=True)
    (out / "DONE").write_text(f"train={nt}\nval={nv}\n")


if __name__ == "__main__":
    main()