96 lines
5.9 KiB
Python
Executable File
96 lines
5.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import argparse, gzip, json, random, time
|
|
from pathlib import Path
|
|
import numpy as np
|
|
from transformers import AutoTokenizer
|
|
|
|
BUDGETS_1B = {
|
|
"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 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("--docmix-dir", default="data/cpt_docmix_5b_sources_8192_20260614")
|
|
ap.add_argument("--out-dir", default="data/cpt_packed_5b_seq8192_seed42_stratified_20260614")
|
|
ap.add_argument("--scale", type=float, default=5.0)
|
|
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=2048)
|
|
ap.add_argument("--eval-rate", type=float, default=0.01)
|
|
ap.add_argument("--shard-blocks", type=int, default=2048)
|
|
ap.add_argument("--log-every-docs", type=int, default=20000)
|
|
args=ap.parse_args()
|
|
base=Path(args.base_dir); docmix=base/args.docmix_dir; out_dir=base/args.out_dir; out_dir.mkdir(parents=True, exist_ok=True)
|
|
sources={cat: docmix/"documents"/f"{cat}.jsonl.gz" for cat in BUDGETS_1B}
|
|
missing=[str(p) for p in sources.values() if not p.exists()]
|
|
if missing: raise FileNotFoundError(missing)
|
|
budgets={k:int(v*args.scale) for k,v in BUDGETS_1B.items()}
|
|
tok=AutoTokenizer.from_pretrained(base/args.tokenizer, trust_remote_code=True); eos=tok.eos_token_id
|
|
rng=random.Random(args.seed)
|
|
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
|
|
stats={"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":{},"leftover_tokens_by_category":{},"start_time":time.time()}
|
|
shards={"train":[],"eval":[]}; shard_idx={"train":0,"eval":0}; train_arrays=[]; eval_arrays=[]
|
|
def add_block(block, category):
|
|
if stats["eval_blocks_by_category"].get(category,0) < eval_quota.get(category,0) and rng.random() < args.eval_rate:
|
|
eval_arrays.append(block); stats["eval_blocks"] += 1; stats["eval_tokens"] += args.seq_len
|
|
stats["eval_blocks_by_category"][category]=stats["eval_blocks_by_category"].get(category,0)+1
|
|
if len(eval_arrays) >= args.shard_blocks:
|
|
rec=flush(out_dir,"eval",shard_idx["eval"],eval_arrays); shards["eval"].append(rec); shard_idx["eval"] += 1; eval_arrays.clear()
|
|
else:
|
|
train_arrays.append(block); stats["train_blocks"] += 1; stats["train_tokens"] += args.seq_len
|
|
stats["train_blocks_by_category"][category]=stats["train_blocks_by_category"].get(category,0)+1
|
|
if len(train_arrays) >= args.shard_blocks:
|
|
rec=flush(out_dir,"train",shard_idx["train"],train_arrays); shards["train"].append(rec); shard_idx["train"] += 1; train_arrays.clear()
|
|
for category,path in sources.items():
|
|
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:
|
|
break
|
|
buffer.extend(ids); cat_tokens += len(ids); cat_docs += 1; stats["docs_seen"] += 1
|
|
stats["source_docs"][category]=cat_docs; stats["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]; add_block(block, category)
|
|
if stats["docs_seen"] % args.log_every_docs == 0:
|
|
rec={k:stats[k] for k in ["docs_seen","train_blocks","eval_blocks","train_tokens","eval_tokens"]}; rec.update({"category":category,"category_tokens":cat_tokens,"elapsed_sec":time.time()-stats["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]; add_block(block, category)
|
|
stats["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={**stats,"tokenizer":str(base/args.tokenizer),"seq_len":args.seq_len,"seed":args.seed,"scale":args.scale,"budgets":budgets,"sources":{k:str(v) for k,v in sources.items()},"eval_quota_blocks":eval_quota,"train_shards":shards["train"],"eval_shards":shards["eval"],"elapsed_sec":time.time()-stats["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()
|