"""Phase 2.6 — tokenize + pack cleaned corpus → data/cpt_packed/ (HF Arrow)

Concatenates all data/cpt_clean/*.jsonl, tokenizes with the Qwen3 tokenizer,
inserts EOS between docs, and packs into fixed 2048-token blocks (ragged tail
dropped). Output columns: input_ids (+ labels==input_ids added by the collator
at train time). corpus_pack is the final corpus artifact train_cpt.py loads.

Usage:  python corpus_pack.py
"""
import json

from datasets import Dataset
from transformers import AutoTokenizer

from corpus_common import CLEAN, PACKED

MODEL = "Qwen/Qwen3-1.7B-Base"
BLOCK = 2048


def iter_texts():
    for fp in sorted(CLEAN.glob("*.jsonl")):
        with open(fp, encoding="utf-8") as f:
            for ln in f:
                ln = ln.strip()
                if ln:
                    yield json.loads(ln)["text"]


def main():
    tok = AutoTokenizer.from_pretrained(MODEL)
    eos = tok.eos_token_id
    if eos is None:
        raise SystemExit("tokenizer has no eos_token_id")

    buf = []
    blocks = []
    n_docs = 0
    for text in iter_texts():
        buf.extend(tok(text, add_special_tokens=False)["input_ids"])
        buf.append(eos)
        n_docs += 1
        while len(buf) >= BLOCK:
            blocks.append(buf[:BLOCK])
            buf = buf[BLOCK:]
        if n_docs % 10000 == 0:
            print(f"[pack] {n_docs} docs -> {len(blocks)} blocks "
                  f"({len(blocks)*BLOCK/1e6:.0f}M tok)")

    ds = Dataset.from_dict({"input_ids": blocks})
    ds.save_to_disk(str(PACKED))
    total_tok = len(blocks) * BLOCK
    print(f"[pack] DONE {n_docs} docs -> {len(blocks)} blocks = {total_tok/1e6:.0f}M tokens")
    print(f"[pack] saved {PACKED}  (target ~300M ± 50M)")

    # sanity: decode a random block, must read as Russian
    import random
    sample = tok.decode(random.choice(blocks))
    print(f"\n--- random packed block (verify Russian) ---\n{sample[:400]}")


if __name__ == "__main__":
    main()