import sys, os, time
import numpy as np
repo, config, text_col, target_b, out = sys.argv[1], sys.argv[2], sys.argv[3], float(sys.argv[4]), sys.argv[5]
target = target_b * 1e9
from transformers import AutoTokenizer
from datasets import load_dataset
tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-14B", trust_remote_code=True)
EOS = tok.eos_token_id if tok.eos_token_id is not None else 151643
os.makedirs(os.path.dirname(out), exist_ok=True)
print("[START] %s [%s] -> %s target=%.1fB eos=%d" % (repo, config, out, target_b, EOS), flush=True)
if config and config != "-":
    ds = load_dataset(repo, config, split="train", streaming=True)
else:
    ds = load_dataset(repo, split="train", streaming=True)
fout = open(out, "wb")
written = 0; last_log = 0; t0 = time.time(); texts = []; BATCH = 2000
def flush():
    global written, texts
    if not texts: return
    enc = tok(texts, add_special_tokens=False)["input_ids"]
    buf = []
    for ids in enc:
        buf.extend(ids); buf.append(EOS)
    np.asarray(buf, dtype=np.uint32).tofile(fout)
    written += len(buf); texts = []
for ex in ds:
    t = ex.get(text_col)
    if not t: continue
    texts.append(t)
    if len(texts) >= BATCH:
        flush()
        if written - last_log >= 5e8:
            dt = max(1, time.time()-t0)
            print("  %.2fB tok (%.0fs, %.0f tok/s)" % (written/1e9, dt, written/dt), flush=True)
            last_log = written
        if written >= target:
            break
flush(); fout.close()
print("[DONE] %s -> %s  %.3fB tokens (%.0fs)" % (repo, out, written/1e9, time.time()-t0), flush=True)