64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
"""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()
|