116 lines
6.7 KiB
Python
Executable File
116 lines
6.7 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
|
|
|
|
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()
|