Clean tokenizer swap migration

2026-06-18 10:10:57 +00:00
commit 5d1e9c4bc3
37 changed files with 9574 additions and 0 deletions
@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+import argparse, gzip, json, random, time
+from pathlib import Path
+import numpy as np
+from transformers import AutoTokenizer
+
+SOURCES = {
+    "english_web": "data/cpt_docmix_parquet_sources_8192_20260613/documents/english_web.jsonl.gz",
+    "english_edu": "data/cpt_docmix_parquet_sources_8192_20260613/documents/english_edu.jsonl.gz",
+    "chinese_clean": "data/cpt_docmix_cci3_science_fixed_8192_20260614/documents/chinese_clean.jsonl.gz",
+    "code": "data/cpt_docmix_parquet_sources_8192_20260613/documents/code.jsonl.gz",
+    "math": "data/cpt_docmix_available_sources_8192_20260613/documents/math.jsonl.gz",
+    "science": "data/cpt_docmix_cci3_science_fixed_8192_20260614/documents/science.jsonl.gz",
+    "qa_as_text": "data/cpt_docmix_available_sources_8192_20260613/documents/qa_as_text.jsonl.gz",
+}
+BUDGETS = {
+    "english_web": 250_000_000,
+    "english_edu": 200_000_000,
+    "chinese_clean": 250_000_000,
+    "code": 150_000_000,
+    "math": 100_000_000,
+    "science": 30_000_000,
+    "qa_as_text": 20_000_000,
+}
+
+def open_text(path):
+    return gzip.open(path, "rt", encoding="utf-8") if path.suffix == ".gz" else path.open("r", encoding="utf-8")
+
+def flush(out_dir, split, shard_idx, arrays):
+    if not arrays: return None
+    arr = np.stack(arrays, axis=0)
+    path = out_dir / f"{split}_{shard_idx:05d}.npy"
+    np.save(path, arr)
+    return {"path": path.name, "blocks": int(arr.shape[0]), "tokens": int(arr.size)}
+
+def maybe_write_block(block, category, args, rng, eval_quota_blocks, counts, train_arrays, eval_arrays, shards, shard_idx):
+    # Stratified eval: each category contributes proportional eval blocks, seed fixed.
+    if counts["eval_blocks_by_category"].get(category, 0) < eval_quota_blocks.get(category, 0) and rng.random() < args.eval_rate:
+        eval_arrays.append(block)
+        counts["eval_blocks"] += 1
+        counts["eval_tokens"] += args.seq_len
+        counts["eval_blocks_by_category"][category] = counts["eval_blocks_by_category"].get(category, 0) + 1
+        if len(eval_arrays) >= args.shard_blocks:
+            rec = flush(args.out_dir_path, "eval", shard_idx["eval"], eval_arrays)
+            shards["eval"].append(rec); shard_idx["eval"] += 1; eval_arrays.clear()
+    else:
+        train_arrays.append(block)
+        counts["train_blocks"] += 1
+        counts["train_tokens"] += args.seq_len
+        counts["train_blocks_by_category"][category] = counts["train_blocks_by_category"].get(category, 0) + 1
+        if len(train_arrays) >= args.shard_blocks:
+            rec = flush(args.out_dir_path, "train", shard_idx["train"], train_arrays)
+            shards["train"].append(rec); shard_idx["train"] += 1; train_arrays.clear()
+
+def main():
+    ap=argparse.ArgumentParser()
+    ap.add_argument("--base-dir", default="/ssd/yi/Tokenizer_Swap")
+    ap.add_argument("--tokenizer", default="model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
+    ap.add_argument("--out-dir", default="data/cpt_packed_1b_seq8192_seed42_stratified_20260614")
+    ap.add_argument("--seq-len", type=int, default=8192)
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--eval-blocks", type=int, default=1024)
+    ap.add_argument("--eval-rate", type=float, default=0.02)
+    ap.add_argument("--shard-blocks", type=int, default=2048)
+    ap.add_argument("--log-every-docs", type=int, default=10000)
+    args=ap.parse_args()
+    base=Path(args.base_dir); out_dir=base/args.out_dir; out_dir.mkdir(parents=True, exist_ok=True)
+    args.out_dir_path=out_dir
+    tok=AutoTokenizer.from_pretrained(base/args.tokenizer, trust_remote_code=True)
+    eos=tok.eos_token_id
+    total_budget=sum(BUDGETS.values())
+    eval_quota={k: round(args.eval_blocks*v/total_budget) for k,v in BUDGETS.items()}
+    diff=args.eval_blocks-sum(eval_quota.values())
+    if diff:
+        eval_quota["english_web"] += diff
+    rng=random.Random(args.seed)
+    counts={"docs_seen":0,"train_blocks":0,"eval_blocks":0,"train_tokens":0,"eval_tokens":0,"source_docs":{},"source_tokens":{},"train_blocks_by_category":{},"eval_blocks_by_category":{},"start_time":time.time()}
+    shards={"train":[],"eval":[]}; shard_idx={"train":0,"eval":0}
+    train_arrays=[]; eval_arrays=[]
+    for category,path_s in SOURCES.items():
+        path=base/path_s
+        target=BUDGETS[category]
+        cat_tokens=0; cat_docs=0; buffer=[]
+        with open_text(path) as f:
+            for line in f:
+                if cat_tokens >= target: break
+                if not line.strip(): continue
+                row=json.loads(line); text=row.get("text") or ""
+                if not text: continue
+                ids=tok.encode(text, add_special_tokens=False)
+                if not ids: continue
+                ids.append(eos)
+                if cat_tokens + len(ids) > target and cat_tokens > 0:
+                    # Keep category budgets tight; do not substantially overshoot.
+                    break
+                buffer.extend(ids); cat_tokens += len(ids); cat_docs += 1; counts["docs_seen"] += 1
+                counts["source_docs"][category]=cat_docs; counts["source_tokens"][category]=cat_tokens
+                while len(buffer) >= args.seq_len:
+                    block=np.asarray(buffer[:args.seq_len], dtype=np.uint32); del buffer[:args.seq_len]
+                    maybe_write_block(block, category, args, rng, eval_quota, counts, train_arrays, eval_arrays, shards, shard_idx)
+                if counts["docs_seen"] % args.log_every_docs == 0:
+                    rec={k:counts[k] for k in ["docs_seen","train_blocks","eval_blocks","train_tokens","eval_tokens"]}; rec["category"]=category; rec["category_tokens"]=cat_tokens; rec["elapsed_sec"]=time.time()-counts["start_time"]
+                    print(json.dumps(rec, ensure_ascii=False), flush=True)
+        while len(buffer) >= args.seq_len:
+            block=np.asarray(buffer[:args.seq_len], dtype=np.uint32); del buffer[:args.seq_len]
+            maybe_write_block(block, category, args, rng, eval_quota, counts, train_arrays, eval_arrays, shards, shard_idx)
+        counts.setdefault("leftover_tokens_by_category", {})[category]=len(buffer)
+    rec=flush(out_dir,"train",shard_idx["train"],train_arrays)
+    if rec: shards["train"].append(rec)
+    rec=flush(out_dir,"eval",shard_idx["eval"],eval_arrays)
+    if rec: shards["eval"].append(rec)
+    manifest={**counts,"tokenizer":str(base/args.tokenizer),"seq_len":args.seq_len,"seed":args.seed,"budgets":BUDGETS,"sources":SOURCES,"eval_quota_blocks":eval_quota,"train_shards":shards["train"],"eval_shards":shards["eval"],"elapsed_sec":time.time()-counts["start_time"]}
+    (out_dir/"manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
+    print(out_dir/"manifest.json")
+if __name__=="__main__": main()