init

llm/corpus_pack.py

@@ -0,0 +1,63 @@
+"""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()