yi_lu/Tokenizer_Swap
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Clean tokenizer swap migration

Browse Source
This commit is contained in:
Codex
2026-06-18 10:10:57 +00:00
commit 5d1e9c4bc3
37 changed files with 9574 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
.gitignore vendored Normal file
View File

@@ -0,0 +1,34 @@
# Runtime outputs and experiment artifacts inside the four workflow folders
model_building/generated_models/
model_training/checkpoints/
evaluation_reporting/outputs/
evaluation_reporting/reports/
dataset_building/generated/
*/logs/
logs/
# Large model/data artifacts
*.safetensors
*.bin
*.pt
*.pth
*.ckpt
*.npy
*.npz
*.arrow
*.parquet
*.jsonl.gz
*.pkl
*.pickle
# Caches
__pycache__/
*.py[cod]
.cache/
.hf_cache/
wandb/
# Local/editor noise
.DS_Store
.vscode/
.idea/

209
README.md Normal file
View File

@@ -0,0 +1,209 @@
# Tokenizer Swap
This is the cleaned migration of the tokenizer-swap experiments. The repo is organized by workflow boundary, not by experiment date or temporary artifact type.
It keeps only the final useful recipes:
- SFT 1M
- CPT 1B
- CPT 5B
- CPT 5B + SFT 1M
- tokenizer swap v2 algorithm
- latest public heldout 2K validation/evaluation set
Trained model weights, optimizer states, generated packed data, and experiment logs are intentionally excluded.
## Repository Structure
```text
dataset_building/ Build SFT/CPT/validation datasets and keep final manifests
model_building/ Build the tokenizer-swapped base model with the v2 algorithm
model_training/ Train SFT/CPT models from existing model and data artifacts
evaluation_reporting/ Run heldout evaluation, merge shards, and generate summaries
```
Each workflow folder has its own README with stage-specific inputs, outputs, scripts, and artifact policy.
The four folders are independent stages:
1. `dataset_building/` produces data artifacts.
2. `model_building/` produces the remapped base checkpoint.
3. `model_training/` consumes a checkpoint plus data and produces trained checkpoints.
4. `evaluation_reporting/` consumes checkpoints plus heldout data and produces reports.
Runtime outputs stay inside these folders but are ignored by git:
```text
dataset_building/generated/
model_building/generated_models/
model_training/checkpoints/
evaluation_reporting/outputs/
evaluation_reporting/reports/
```
## 1. Dataset Building
This folder contains the builders for:
- SFT 1M chat mixture and DSV4-tokenized chat JSONL
- CPT document mixtures
- CPT packed training/eval blocks
- public heldout 2K validation/evaluation set
The only heldout dataset kept in the repo is:
```text
dataset_building/heldout_public_mcq_2k_20260607/heldout_public_mcq_2k.jsonl
```
The in-domain heldout and earlier ratio-imbalanced 2K heldout are not included.
### SFT Data
The final SFT recipe is the no-upsample, no-BBH v4 chat mix derived from the 1M instruction mixture. Build metadata is kept in:
```text
dataset_building/metadata/sft_v4_mix_build_stats.json
dataset_building/metadata/sft_v4_tokenization_build_stats.json
```
The final training scripts expect generated tokenized SFT data under:
```text
dataset_building/generated/dsv4_chat_tokenized_v4_noupsample_nobbh_921k/
train_dsv4_chat_tokenized.jsonl.gz
validation_dsv4_chat_tokenized.jsonl.gz
```
### CPT Data
The final CPT data uses stratified packing with sequence length 8192 and seed 42.
Target source proportions:
| Source bucket | Ratio |
|---|---:|
| English web | 25% |
| English education | 20% |
| Chinese clean | 25% |
| Code | 15% |
| Math | 10% |
| Science | 3% |
| QA as text | 2% |
Final token counts:
| Bucket | CPT 1B source tokens | CPT 5B source tokens |
|---|---:|---:|
| English web | 249,999,374 | 1,249,998,964 |
| English education | 199,998,420 | 999,999,328 |
| Chinese clean | 249,999,618 | 1,249,999,598 |
| Code | 149,999,958 | 749,998,673 |
| Math | 99,987,860 | 499,995,208 |
| Science | 29,999,428 | 149,998,447 |
| QA as text | 19,999,497 | 99,999,341 |
Packed output sizes:
| Dataset | Train tokens | Eval tokens | Train blocks | Eval blocks | Seq len |
|---|---:|---:|---:|---:|---:|
| CPT 1B | 991,567,872 | 8,388,608 | 121,041 | 1,024 | 8192 |
| CPT 5B | 4,983,177,216 | 16,777,216 | 608,298 | 2,048 | 8192 |
Source families:
- FineWeb for English web
- FineWeb-Edu for English education
- BAAI/CCI3-HQ cleaned Chinese data
- StarCoderData/code parquet sources
- OpenWebMath-derived documents
- science reasoning datasets including MedMCQA, ProofWriter, ScienceQA, MedQA, SciQ, QASC, and OpenBookQA
- QA rendered as plain text from the final instruction mixture plus recovery fallback sources
Final manifests are kept in `dataset_building/metadata/`.
## 2. Model Building
`model_building/build_qwen3_dsv4_remap_checkpoint_v2.py` builds the tokenizer-swapped base checkpoint.
Run:
```bash
ROOT=/ssd/yi/Tokenizer_Swap bash model_building/run_remap_v2.sh
```
The algorithm:
- load the source Qwen model/tokenizer and target DSV4 tokenizer
- resize/rebuild input embedding and LM-head rows to the DSV4 vocab size
- initialize each target token row by priority:
- exact same token surface in Qwen vocab
- functional special-token mapping, for example DSV BOS to Qwen `<|im_start|>` and DSV EOS to Qwen EOS
- byte-level decode followed by Qwen tokenization, averaging old rows
- raw token decomposition fallback, averaging old rows
- global embedding/head mean fallback
- save the remapped checkpoint and `tokenizer_remap_v2_report.json`
The model-building step does not build datasets and does not train.
## 3. Model Training
Training scripts consume existing model/data paths. They do not perform tokenizer remapping or dataset construction.
Final entrypoints:
| Experiment | Script |
|---|---|
| SFT 1M on remapped base | `model_training/run_sft1m_remap_v2_5epoch.sh` |
| SFT 1M plus v4 no-upsample continuation | `model_training/run_sft1m_remap_v2_then_v4_noupsample_5epoch_bsz16.sh` |
| CPT 1B | `model_training/run_cpt1b_seed42_train_eval.sh` |
| CPT 5B | `model_training/run_cpt5b_seed42_train_eval.sh` |
| CPT 5B + SFT 1M | `model_training/run_cpt5b_then_sft1m_5epoch.sh` |
Most paths are configurable with environment variables:
```bash
MODEL=... DATA=... TRAIN=... EVAL=... OUT=... NPROC=8 bash model_training/run_cpt5b_then_sft1m_5epoch.sh
```
Default outputs go to `model_training/checkpoints/`, which is ignored by git.
## 4. Evaluation And Report Generation
The public heldout 2K evaluation entrypoint is:
```bash
ROOT=/ssd/yi/Tokenizer_Swap \
MODEL=/path/to/checkpoint \
LABEL=my_model \
bash evaluation_reporting/run_public_heldout_eval_8gpu.sh
```
This runs sharded evaluation, merges per-shard outputs, and writes summaries under `evaluation_reporting/outputs/`.
Main public heldout 2K results from the final sweep:
| Model | MCQ acc avg-norm | MCQ acc sum | PPL | NLL/token |
|---|---:|---:|---:|---:|
| Native Qwen3-0.6B tokenizer baseline | 0.2960 | 0.2510 | 61.08 | 3.5399 |
| Remap v2, no training | 0.2940 | 0.2410 | 313.62 | 4.6920 |
| Remap v2 + CPT 1B | 0.3005 | 0.2540 | 71.90 | 3.6580 |
| Remap v2 + CPT 5B | 0.3020 | 0.2615 | 66.95 | 3.5806 |
| Remap v2 + SFT 1M | 0.3105 | 0.2590 | 114.75 | 3.9400 |
| Remap v2 + SFT 1M + v4 continuation | 0.3165 | 0.2595 | 117.33 | 3.9755 |
| Remap v2 + CPT 5B + SFT 1M | 0.3280 | 0.2740 | 88.76 | 3.7899 |
CPT mainly repairs language-modeling quality after tokenizer replacement. SFT improves heldout MCQ/task behavior but can raise perplexity because the objective focuses on assistant-answer tokens rather than generic next-token modeling over heldout text.
## Artifact Policy
Do not commit:
- trained checkpoints or partial checkpoints
- optimizer states
- generated packed CPT data
- generated SFT `.jsonl.gz` data
- `.safetensors`, `.bin`, `.pt`, `.npy`, `.parquet`
- experiment logs and evaluation output directories
The original repo's experiment logs were cleared during migration.

98
dataset_building/README.md Normal file
View File

@@ -0,0 +1,98 @@
# Dataset Building
This folder owns all dataset construction for the migration: SFT data, CPT data, and the public heldout validation/evaluation set.
It does not build model weights and does not run training.
## Inputs
The builders expect raw or intermediate source data outside the git repo, such as FineWeb/FineWeb-Edu shards, BAAI/CCI3-HQ Chinese data, StarCoder/code data, OpenWebMath-derived data, science reasoning caches, and instruction QA mixes.
Tokenizer paths are configurable with CLI flags. Generated data should be written under:
```text
dataset_building/generated/
```
That directory is ignored by git.
## Kept Dataset
Only the latest public heldout 2K set is versioned here:
```text
heldout_public_mcq_2k_20260607/heldout_public_mcq_2k.jsonl
heldout_public_mcq_2k_20260607/heldout_public_mcq_2k_stats.json
```
Older in-domain heldout and earlier ratio-imbalanced 2K heldout datasets are intentionally not included.
## SFT Builders
Main scripts:
```text
build_training_and_test_mix_v3.py
build_dsv4_chat_tokenized_messages_jsonl.py
build_dsv4_chat_tokenized_custom.py
```
Expected final generated layout:
```text
dataset_building/generated/dsv4_chat_tokenized_v4_noupsample_nobbh_921k/
train_dsv4_chat_tokenized.jsonl.gz
validation_dsv4_chat_tokenized.jsonl.gz
```
Build metadata from the final SFT recipe is kept in:
```text
metadata/sft_v4_mix_build_stats.json
metadata/sft_v4_tokenization_build_stats.json
```
## CPT Builders
Main scripts:
```text
build_cpt_docmix_1b.py
build_cpt_packed_stratified.py
build_cpt_packed_5b_stratified.py
build_cci3_chinese_docmix_fix.py
build_math_docmix_fix.py
build_science_docmix_fix.py
```
The final CPT recipes use stratified packing with sequence length 8192 and seed 42.
Source proportions:
| Source bucket | Ratio |
|---|---:|
| English web | 25% |
| English education | 20% |
| Chinese clean | 25% |
| Code | 15% |
| Math | 10% |
| Science | 3% |
| QA as text | 2% |
Final manifests are kept in:
```text
metadata/cpt_packed_1b_seed42_stratified_manifest.json
metadata/cpt_packed_5b_seed42_stratified_manifest.json
metadata/cpt_docmix_5b_manifest.json
metadata/cpt_docmix_5b_stats.json
```
## Output Contract
Downstream training expects either:
- tokenized SFT `.jsonl.gz` files for `train_dsv4_tokenized_full_sft.py`
- packed CPT arrays/manifests for `train_cpt_packed_full.py`
Do not commit generated `.jsonl.gz`, `.npy`, `.parquet`, or other large intermediate files.

173
dataset_building/build_cci3_chinese_docmix_fix.py Executable file
View File

@@ -0,0 +1,173 @@
#!/usr/bin/env python3
import argparse
import gzip
import json
import os
import re
import time
from collections import Counter
from pathlib import Path
from huggingface_hub import HfApi, hf_hub_download
from transformers import AutoTokenizer
def clean_text(text):
if text is None:
return ""
text = str(text).replace("\x00", " ")
text = re.sub(r"[ \t\r\f\v]+", " ", text)
text = re.sub(r"\n{4,}", "\n\n\n", text)
return text.strip()
def iter_jsonl(path):
with open(path, "r", encoding="utf-8", errors="replace") as f:
for line in f:
line = line.strip()
if not line:
continue
try:
yield json.loads(line)
except json.JSONDecodeError:
continue
def download_with_retry(filename, args):
last = None
for attempt in range(1, args.retries + 1):
try:
return hf_hub_download(
repo_id=args.repo,
repo_type="dataset",
filename=filename,
endpoint=args.endpoint,
token=os.environ.get("HF_TOKEN"),
local_dir=args.raw_dir,
)
except Exception as exc:
last = exc
print(json.dumps({"event": "download_retry", "file": filename, "attempt": attempt, "error": repr(exc)[:800]}, ensure_ascii=False), flush=True)
time.sleep(min(120, 5 * attempt))
raise RuntimeError(f"download failed for {filename}: {last!r}")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--repo", default="BAAI/CCI3-HQ")
ap.add_argument("--endpoint", default=os.environ.get("HF_ENDPOINT", "https://hf-mirror.com"))
ap.add_argument("--tokenizer", default="model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
ap.add_argument("--raw-dir", default="data/raw_jsonl/cci3_hq_probe")
ap.add_argument("--out-dir", default="data/cpt_docmix_5b_sources_8192_20260614")
ap.add_argument("--target-tokens", type=int, default=1_250_000_000)
ap.add_argument("--min-tokens", type=int, default=128)
ap.add_argument("--max-doc-tokens", type=int, default=32768)
ap.add_argument("--log-every", type=int, default=5000)
ap.add_argument("--retries", type=int, default=16)
ap.add_argument("--start-index", type=int, default=0)
ap.add_argument("--keep-raw", action="store_true")
args = ap.parse_args()
base = Path.cwd()
raw_dir = base / args.raw_dir
out_dir = base / args.out_dir
doc_dir = out_dir / "documents"
doc_dir.mkdir(parents=True, exist_ok=True)
raw_dir.mkdir(parents=True, exist_ok=True)
token = os.environ.get("HF_TOKEN")
if not token:
raise SystemExit("HF_TOKEN is required for gated CCI3-HQ")
api = HfApi(endpoint=args.endpoint, token=token)
files = [f for f in api.list_repo_files(args.repo, repo_type="dataset") if f.startswith("data/part_") and f.endswith(".jsonl")]
files = sorted(files)
if args.start_index:
files = [f for f in files if int(Path(f).stem.split("_")[-1]) >= args.start_index]
tok = AutoTokenizer.from_pretrained(base / args.tokenizer, trust_remote_code=True)
tmp_out = doc_dir / "chinese_clean.jsonl.gz.tmp"
final_out = doc_dir / "chinese_clean.jsonl.gz"
if tmp_out.exists():
tmp_out.unlink()
stats = {
"target_tokens": args.target_tokens,
"tokens": 0,
"docs": 0,
"rows_seen": 0,
"files_done": [],
"tokens_by_file": Counter(),
"rejected": Counter(),
"started_at": time.time(),
"repo": args.repo,
"endpoint": args.endpoint,
}
with gzip.open(tmp_out, "wt", encoding="utf-8") as w:
for filename in files:
if stats["tokens"] >= args.target_tokens:
break
local_path = Path(download_with_retry(filename, args))
file_tokens = 0
file_docs = 0
for row in iter_jsonl(local_path):
stats["rows_seen"] += 1
if stats["tokens"] >= args.target_tokens:
break
text = clean_text(row.get("text") or row.get("content"))
if not text:
stats["rejected"]["empty"] += 1
continue
ntok = len(tok.encode(text, add_special_tokens=False))
if ntok < args.min_tokens:
stats["rejected"]["too_short"] += 1
continue
if ntok > args.max_doc_tokens:
stats["rejected"]["too_long"] += 1
continue
idx = stats["docs"]
rec = {
"id": f"chinese_clean_cci3_hq_{idx:09d}",
"category": "chinese_clean",
"source": f"{args.repo}:{filename}",
"text": text,
"token_count": ntok,
"metadata": {"cci3_id": row.get("id"), "cci3_score": row.get("score")},
}
w.write(json.dumps(rec, ensure_ascii=False) + "\n")
stats["docs"] += 1
stats["tokens"] += ntok
file_docs += 1
file_tokens += ntok
if stats["docs"] % args.log_every == 0:
print(json.dumps({"event": "progress", "docs": stats["docs"], "tokens": stats["tokens"], "target": args.target_tokens, "file": filename, "elapsed_sec": time.time() - stats["started_at"]}, ensure_ascii=False), flush=True)
stats["files_done"].append({"file": filename, "docs": file_docs, "tokens": file_tokens, "path": str(local_path)})
stats["tokens_by_file"][filename] += file_tokens
print(json.dumps({"event": "file_done", "file": filename, "docs": file_docs, "tokens": file_tokens, "total_tokens": stats["tokens"]}, ensure_ascii=False), flush=True)
if not args.keep_raw:
try:
local_path.unlink()
except Exception:
pass
if stats["tokens"] < args.target_tokens:
raise SystemExit(f"only collected {stats['tokens']} / {args.target_tokens} tokens")
backup = final_out.with_suffix(".jsonl.gz.failed_empty_20260614")
if final_out.exists() and final_out.stat().st_size < 1024:
final_out.replace(backup)
elif final_out.exists():
final_out.replace(final_out.with_suffix(".jsonl.gz.backup_20260614"))
tmp_out.replace(final_out)
stats["elapsed_sec"] = time.time() - stats["started_at"]
stats["tokens_by_file"] = dict(stats["tokens_by_file"])
stats["rejected"] = dict(stats["rejected"])
(out_dir / "chinese_clean_5b_fix_stats.json").write_text(json.dumps(stats, ensure_ascii=False, indent=2), encoding="utf-8")
(out_dir / ".chinese_clean_5b_ready").write_text(json.dumps({"tokens": stats["tokens"], "docs": stats["docs"], "elapsed_sec": stats["elapsed_sec"]}, ensure_ascii=False), encoding="utf-8")
print(json.dumps({"event": "done", "tokens": stats["tokens"], "docs": stats["docs"], "elapsed_sec": stats["elapsed_sec"], "output": str(final_out)}, ensure_ascii=False), flush=True)
if __name__ == "__main__":
main()

489
dataset_building/build_cpt_docmix_1b.py Normal file
View File

@@ -0,0 +1,489 @@
#!/usr/bin/env python3
import argparse
import gzip
import json
import os
import re
import time
from collections import Counter, defaultdict
from pathlib import Path
from datasets import load_dataset
from transformers import AutoTokenizer
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,
}
STREAM_SOURCES = {
"english_web": [
{"kind": "hf", "name": "HuggingFaceFW/fineweb", "config": "CC-MAIN-2025-26", "split": "train", "max_rows": 0},
{"kind": "hf", "name": "HuggingFaceFW/fineweb", "config": "CC-MAIN-2025-21", "split": "train", "max_rows": 0},
],
"english_edu": [
{"kind": "hf", "name": "HuggingFaceFW/fineweb-edu", "config": None, "split": "train", "max_rows": 0},
],
"chinese_clean": [
{"kind": "hf", "name": "BAAI/CCI3-HQ", "config": None, "split": "train", "max_rows": 0},
{"kind": "hf", "name": "Skywork/SkyPile-150B", "config": None, "split": "train", "max_rows": 0},
],
"code": [
{"kind": "hf", "name": "bigcode/starcoderdata", "config": None, "split": "train", "max_rows": 0},
{"kind": "hf", "name": "codeparrot/github-code", "config": None, "split": "train", "max_rows": 0},
],
"math": [
{"kind": "hf", "name": "open-web-math/open-web-math", "config": None, "split": "train", "max_rows": 0},
{"kind": "hf", "name": "GAIR/MathPile", "config": None, "split": "train", "max_rows": 0},
],
}
LOCAL_GLOBS = {
"science": [
"data/offline_text_only_reasoning_sources_20260611/science_reasoning__*.jsonl",
"data/offline_text_only_reasoning_sources_20260611/logic__*.jsonl",
],
"qa_as_text": [
"data/training_mix_v4_train1m_test2p8k_noupsample_nobbh_20260611/train_1m.jsonl",
"data/open_recovery_sft_mix_alt_sources_1m_parquet_20260607/normalized.jsonl",
],
}
LOCAL_PARQUET_GLOBS = {
"english_web": [
"data/raw_parquets/fineweb_2025/*.parquet",
],
"english_edu": [
"data/raw_parquets/fineweb_edu/*.parquet",
],
"code": [
"data/raw_parquets/starcoder/python/*.parquet",
"data/raw_parquets/starcoder/javascript/*.parquet",
"data/raw_parquets/starcoder/typescript/*.parquet",
"data/raw_parquets/starcoder/java/*.parquet",
"data/raw_parquets/starcoder/cpp/*.parquet",
"data/raw_parquets/starcoder/go/*.parquet",
"data/raw_parquets/starcoder/rust/*.parquet",
"data/raw_parquets/starcoder/shell/*.parquet",
],
}
def clean_text(text):
if text is None:
return ""
text = str(text).replace("\x00", " ")
text = re.sub(r"[ \t\r\f\v]+", " ", text)
text = re.sub(r"\n{4,}", "\n\n\n", text)
return text.strip()
def first_user_assistant(messages):
user = None
assistant = None
for msg in messages or []:
role = msg.get("role")
content = clean_text(msg.get("content"))
if not content:
continue
if role == "user" and user is None:
user = content
elif role == "assistant" and user is not None:
assistant = content
break
if user and assistant:
return user, assistant
return None, None
def qa_text(user, assistant):
if not user or not assistant:
return ""
user_label = "" if user.lstrip().lower().startswith(("question:", "problem:", "context:", "support:", "fact1:")) else "Question:\n"
assistant_label = "" if assistant.lstrip().lower().startswith(("answer:", "solution:", "final answer:", "explanation:")) else "Answer:\n"
return f"{user_label}{user}\n\n{assistant_label}{assistant}"
def extract_text(row, category, source_name):
if category == "science":
user, assistant = first_user_assistant(row.get("messages"))
if user and assistant:
return qa_text(user, assistant)
if category == "qa_as_text" or row.get("messages"):
user, assistant = first_user_assistant(row.get("messages"))
if user and assistant:
return qa_text(user, assistant)
if category == "code":
content = clean_text(row.get("content") or row.get("text") or row.get("code"))
if not content:
return ""
path = clean_text(row.get("path") or row.get("max_stars_repo_path") or row.get("repo_name") or "")
lang = clean_text(row.get("lang") or row.get("language") or "")
if path or lang:
attrs = []
if path:
attrs.append(f'path="{path[:300]}"')
if lang:
attrs.append(f'language="{lang[:80]}"')
return f"<file {' '.join(attrs)}>\n{content}\n</file>"
return content
for key in ("text", "content", "markdown", "raw_content"):
text = clean_text(row.get(key))
if text:
return text
return ""
def token_count(tok, text):
return len(tok.encode(text, add_special_tokens=False))
def open_gz_writer(path):
path.parent.mkdir(parents=True, exist_ok=True)
return gzip.open(path, "at", encoding="utf-8")
def write_doc(writer, doc_id, category, source, text, ntok, metadata=None):
writer.write(
json.dumps(
{
"id": doc_id,
"category": category,
"source": source,
"text": text,
"token_count": ntok,
"metadata": metadata or {},
},
ensure_ascii=False,
)
+ "\n"
)
def should_keep(row, category, text, ntok, args):
if ntok < args.min_tokens:
return False, "too_short"
if ntok > args.max_doc_tokens:
return False, "too_long"
if category == "english_web":
lang = row.get("language")
score = row.get("language_score")
if lang and str(lang).lower() != "en":
return False, "non_en"
if score is not None:
try:
if float(score) < args.min_fineweb_language_score:
return False, "low_language_score"
except Exception:
pass
return True, ""
def iter_local_jsonl(paths):
for path in paths:
with open(path, "r", encoding="utf-8") as f:
for line in f:
if line.strip():
yield path, json.loads(line)
def flush_pack(writer, category, source, parts, ntok, stats, args):
if not parts:
return [], 0
if ntok < args.min_tokens:
stats["rejected"][category]["pack_too_short"] += 1
return [], 0
idx = stats["docs_by_category"][category]
text = "\n\n---\n\n".join(parts)
write_doc(
writer,
f"{category}_packed_{idx:09d}",
category,
source,
text,
ntok,
{"local_path": source, "packed_items": len(parts)},
)
stats["docs_by_category"][category] += 1
stats["tokens_by_category"][category] += ntok
stats["tokens_by_source"][source] += ntok
if stats["docs_by_category"][category] % args.log_every == 0:
print_progress(stats, category)
return [], 0
def collect_local(base, tok, category, out_dir, args, stats):
paths = []
for pattern in LOCAL_GLOBS.get(category, []):
paths.extend(str(p) for p in sorted((base).glob(pattern)))
if not paths:
stats["sources"].append({"category": category, "kind": "local", "error": "no local files"})
return
out_path = out_dir / "documents" / f"{category}.jsonl.gz"
with open_gz_writer(out_path) as writer:
pack_parts = []
pack_tokens = 0
pack_source = ""
for path, row in iter_local_jsonl(paths):
if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
break
text = extract_text(row, category, path)
if not text:
stats["rejected"][category]["empty"] += 1
continue
ntok = token_count(tok, text)
if category == "science" and ntok < args.min_tokens:
if ntok > args.max_doc_tokens:
stats["rejected"][category]["too_long"] += 1
continue
if pack_source and pack_source != path:
pack_parts, pack_tokens = flush_pack(writer, category, pack_source, pack_parts, pack_tokens, stats, args)
pack_source = path
if pack_tokens + ntok > args.science_pack_tokens and pack_parts:
pack_parts, pack_tokens = flush_pack(writer, category, pack_source, pack_parts, pack_tokens, stats, args)
pack_parts.append(text)
pack_tokens += ntok
continue
keep, reason = should_keep(row, category, text, ntok, args)
if not keep:
stats["rejected"][category][reason] += 1
continue
if category == "science" and pack_parts:
pack_parts, pack_tokens = flush_pack(writer, category, pack_source, pack_parts, pack_tokens, stats, args)
idx = stats["docs_by_category"][category]
write_doc(writer, f"{category}_local_{idx:09d}", category, path, text, ntok, {"local_path": path})
stats["docs_by_category"][category] += 1
stats["tokens_by_category"][category] += ntok
stats["tokens_by_source"][path] += ntok
if stats["docs_by_category"][category] % args.log_every == 0:
print_progress(stats, category)
if category == "science" and pack_parts and stats["tokens_by_category"][category] < BUDGETS[category] * args.scale:
flush_pack(writer, category, pack_source, pack_parts, pack_tokens, stats, args)
def collect_local_parquet(base, tok, category, out_dir, args, stats):
try:
import pyarrow.parquet as pq
except Exception as exc:
stats["sources"].append({"category": category, "kind": "local_parquet", "error": repr(exc)})
return
paths = []
for pattern in LOCAL_PARQUET_GLOBS.get(category, []):
paths.extend(sorted(base.glob(pattern)))
if not paths:
stats["sources"].append({"category": category, "kind": "local_parquet", "error": "no parquet files"})
return
out_path = out_dir / "documents" / f"{category}.jsonl.gz"
with open_gz_writer(out_path) as writer:
for path in paths:
if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
break
source_label = str(path.relative_to(base))
source_rec = {"category": category, "kind": "local_parquet", "source": source_label, "rows_seen": 0, "docs_written": 0, "tokens": 0, "error": ""}
try:
pf = pq.ParquetFile(path)
for batch in pf.iter_batches(batch_size=args.parquet_batch_size):
rows = batch.to_pylist()
for row in rows:
source_rec["rows_seen"] += 1
if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
break
text = extract_text(row, category, source_label)
if not text:
stats["rejected"][category]["empty"] += 1
continue
ntok = token_count(tok, text)
keep, reason = should_keep(row, category, text, ntok, args)
if not keep:
stats["rejected"][category][reason] += 1
continue
idx = stats["docs_by_category"][category]
write_doc(
writer,
f"{category}_{safe_source(source_label)}_{idx:09d}",
category,
source_label,
text,
ntok,
{k: row.get(k) for k in ("url", "dump", "date", "language", "language_score", "path", "lang", "license") if k in row},
)
stats["docs_by_category"][category] += 1
stats["tokens_by_category"][category] += ntok
stats["tokens_by_source"][source_label] += ntok
source_rec["docs_written"] += 1
source_rec["tokens"] += ntok
if stats["docs_by_category"][category] % args.log_every == 0:
print_progress(stats, category)
if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
break
except Exception as exc:
source_rec["error"] = repr(exc)
print(json.dumps({"event": "local_parquet_error", **source_rec}, ensure_ascii=False), flush=True)
stats["sources"].append(source_rec)
def collect_stream(tok, category, out_dir, args, stats):
out_path = out_dir / "documents" / f"{category}.jsonl.gz"
with open_gz_writer(out_path) as writer:
for spec in STREAM_SOURCES.get(category, []):
if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
break
source_label = spec["name"] if not spec.get("config") else f"{spec['name']}:{spec['config']}"
source_rec = {"category": category, "kind": "hf_stream", "source": source_label, "rows_seen": 0, "docs_written": 0, "tokens": 0, "error": ""}
try:
ds = load_dataset(
spec["name"],
spec.get("config"),
split=spec.get("split", "train"),
streaming=True,
trust_remote_code=True,
)
for row in ds:
source_rec["rows_seen"] += 1
if spec.get("max_rows") and source_rec["rows_seen"] > spec["max_rows"]:
break
if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
break
text = extract_text(row, category, source_label)
if not text:
stats["rejected"][category]["empty"] += 1
continue
ntok = token_count(tok, text)
keep, reason = should_keep(row, category, text, ntok, args)
if not keep:
stats["rejected"][category][reason] += 1
continue
idx = stats["docs_by_category"][category]
write_doc(
writer,
f"{category}_{safe_source(source_label)}_{idx:09d}",
category,
source_label,
text,
ntok,
{k: row.get(k) for k in ("url", "dump", "date", "language", "language_score", "path", "lang", "license") if k in row},
)
stats["docs_by_category"][category] += 1
stats["tokens_by_category"][category] += ntok
stats["tokens_by_source"][source_label] += ntok
source_rec["docs_written"] += 1
source_rec["tokens"] += ntok
if stats["docs_by_category"][category] % args.log_every == 0:
print_progress(stats, category)
except Exception as exc:
source_rec["error"] = repr(exc)
print(json.dumps({"event": "source_error", **source_rec}, ensure_ascii=False), flush=True)
stats["sources"].append(source_rec)
def safe_source(text):
return re.sub(r"[^A-Za-z0-9._-]+", "_", text)[:120]
def print_progress(stats, category):
print(
json.dumps(
{
"event": "progress",
"category": category,
"docs": stats["docs_by_category"][category],
"tokens": stats["tokens_by_category"][category],
"target": int(BUDGETS[category] * stats["scale"]),
"elapsed_sec": time.time() - stats["start_time"],
},
ensure_ascii=False,
),
flush=True,
)
def dump_stats(out_dir, stats):
serializable = {
"scale": stats["scale"],
"budgets": {k: int(v * stats["scale"]) for k, v in BUDGETS.items()},
"tokens_by_category": dict(stats["tokens_by_category"]),
"docs_by_category": dict(stats["docs_by_category"]),
"tokens_by_source": dict(stats["tokens_by_source"].most_common()),
"rejected": {k: dict(v) for k, v in stats["rejected"].items()},
"sources": stats["sources"],
"elapsed_sec": time.time() - stats["start_time"],
}
(out_dir / "stats.json").write_text(json.dumps(serializable, ensure_ascii=False, indent=2), encoding="utf-8")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--base-dir", default="/ssd/yi/Tokenizer_Swap")
parser.add_argument("--tokenizer", default="model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
parser.add_argument("--out-dir", default="data/cpt_docmix_1b_8192_20260613")
parser.add_argument("--scale", type=float, default=1.0, help="1.0 = 1B token budget; 0.001 = 1M token smoke")
parser.add_argument("--categories", default="english_web,english_edu,chinese_clean,code,math,science,qa_as_text")
parser.add_argument("--min-tokens", type=int, default=128)
parser.add_argument("--max-doc-tokens", type=int, default=32768)
parser.add_argument("--min-fineweb-language-score", type=float, default=0.65)
parser.add_argument("--log-every", type=int, default=1000)
parser.add_argument("--parquet-batch-size", type=int, default=1000)
parser.add_argument("--science-pack-tokens", type=int, default=2048)
args = parser.parse_args()
base = Path(args.base_dir)
out_dir = base / args.out_dir
out_dir.mkdir(parents=True, exist_ok=True)
(out_dir / "manifest.json").write_text(
json.dumps(
{
"tokenizer": str(base / args.tokenizer),
"budgets": {k: int(v * args.scale) for k, v in BUDGETS.items()},
"seq_len_for_later_packing": 8192,
"stream_sources": STREAM_SOURCES,
"local_globs": LOCAL_GLOBS,
"local_parquet_globs": LOCAL_PARQUET_GLOBS,
"min_tokens": args.min_tokens,
"max_doc_tokens": args.max_doc_tokens,
},
ensure_ascii=False,
indent=2,
),
encoding="utf-8",
)
tok = AutoTokenizer.from_pretrained(base / args.tokenizer, trust_remote_code=True)
stats = {
"scale": args.scale,
"start_time": time.time(),
"tokens_by_category": Counter(),
"docs_by_category": Counter(),
"tokens_by_source": Counter(),
"rejected": defaultdict(Counter),
"sources": [],
}
for category in [x.strip() for x in args.categories.split(",") if x.strip()]:
if category in LOCAL_PARQUET_GLOBS:
collect_local_parquet(base, tok, category, out_dir, args, stats)
if category in STREAM_SOURCES:
collect_stream(tok, category, out_dir, args, stats)
if category in LOCAL_GLOBS:
collect_local(base, tok, category, out_dir, args, stats)
dump_stats(out_dir, stats)
dump_stats(out_dir, stats)
print(out_dir / "stats.json")
if __name__ == "__main__":
main()

95
dataset_building/build_cpt_packed_5b_stratified.py Executable file
View File

@@ -0,0 +1,95 @@
#!/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()

115
dataset_building/build_cpt_packed_stratified.py Executable file
View File

@@ -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()

166
dataset_building/build_dsv4_chat_tokenized_custom.py Normal file
View File

@@ -0,0 +1,166 @@
#!/usr/bin/env python3
import argparse
import gzip
import json
import sys
from pathlib import Path
from transformers import AutoTokenizer
def import_encoding(encoding_dir: Path):
sys.path.insert(0, str(encoding_dir))
import encoding_dsv4 # type: ignore
return encoding_dsv4
def read_json(path):
return json.loads(Path(path).read_text(encoding="utf-8"))
def open_writer(path: Path):
path.parent.mkdir(parents=True, exist_ok=True)
if path.suffix == ".gz":
return gzip.open(path, "wt", encoding="utf-8")
return path.open("w", encoding="utf-8")
def quantiles(xs):
if not xs:
return {}
xs = sorted(xs)
return {
"p50": xs[int((len(xs) - 1) * 0.50)],
"p90": xs[int((len(xs) - 1) * 0.90)],
"p95": xs[int((len(xs) - 1) * 0.95)],
"p99": xs[int((len(xs) - 1) * 0.99)],
"max": xs[-1],
}
def build_split(name, src_path, out_path, tok, enc, cutoff_len):
rows = read_json(src_path)
stats = {
"split": name,
"source": str(src_path),
"output": str(out_path),
"n": len(rows),
"cutoff_len": cutoff_len,
"truncated": 0,
"prompt_tokens": [],
"response_tokens": [],
"total_tokens": [],
"eos_in_labels": 0,
"prefix_mismatch": 0,
}
with open_writer(out_path) as f:
for idx, row in enumerate(rows):
instruction = (row.get("instruction") or "").strip()
output = row.get("output") or ""
messages_prompt = [{"role": "user", "content": instruction}]
messages_full = [
{"role": "user", "content": instruction},
{"role": "assistant", "content": output},
]
prompt_text = enc.encode_messages(messages_prompt, thinking_mode="chat")
full_text = enc.encode_messages(messages_full, thinking_mode="chat")
if not full_text.startswith(prompt_text):
stats["prefix_mismatch"] += 1
prompt_ids = tok(prompt_text, add_special_tokens=False).input_ids
full_ids = tok(full_text, add_special_tokens=False).input_ids
response_ids = full_ids[len(prompt_ids) :]
labels = [-100] * len(prompt_ids) + response_ids
truncated = False
if len(full_ids) > cutoff_len:
truncated = True
full_ids = full_ids[:cutoff_len]
labels = labels[:cutoff_len]
stats["truncated"] += 1
eos_in_labels = tok.eos_token_id in [x for x in labels if x != -100]
stats["eos_in_labels"] += int(eos_in_labels)
stats["prompt_tokens"].append(len(prompt_ids))
stats["response_tokens"].append(len(response_ids))
stats["total_tokens"].append(len(prompt_ids) + len(response_ids))
out = {
"id": row.get("id", f"{name}_{idx:07d}"),
"split": name,
"source_task": row.get("task_type") or row.get("task") or row.get("category"),
"source": row.get("source"),
"thinking_mode": "chat",
"messages": messages_full,
"prompt_text": prompt_text,
"full_text": full_text,
"prompt_tokens": len(prompt_ids),
"response_tokens": len(response_ids),
"total_tokens": len(prompt_ids) + len(response_ids),
"truncated": truncated,
"eos_in_labels": eos_in_labels,
"input_ids": full_ids,
"labels": labels,
}
f.write(json.dumps(out, ensure_ascii=False) + "\n")
for key in ["prompt_tokens", "response_tokens", "total_tokens"]:
stats[key] = quantiles(stats[key])
stats["truncated_rate"] = stats["truncated"] / max(1, stats["n"])
stats["eos_label_rate"] = stats["eos_in_labels"] / max(1, stats["n"])
return stats
def main():
p = argparse.ArgumentParser()
p.add_argument("--base-dir", default="/ssd/yi/Tokenizer_Swap")
p.add_argument("--tokenizer", default="model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
p.add_argument("--encoding-dir", default="external/deepseek_v4_encoding")
p.add_argument("--data-dir", required=True)
p.add_argument("--out-dir", required=True)
p.add_argument("--train-file", required=True)
p.add_argument("--validation-file", default="fixed_validation.json")
p.add_argument("--case-file", default="fixed_case.json")
p.add_argument("--cutoff-len", type=int, default=2048)
p.add_argument("--gzip", action="store_true")
args = p.parse_args()
base = Path(args.base_dir)
data_dir = base / args.data_dir
out_dir = base / args.out_dir
enc = import_encoding(base / args.encoding_dir)
tok = AutoTokenizer.from_pretrained(base / args.tokenizer, trust_remote_code=True)
suffix = ".jsonl.gz" if args.gzip else ".jsonl"
split_files = {
"train": args.train_file,
"validation": args.validation_file,
"case": args.case_file,
}
all_stats = {
"tokenizer": str(base / args.tokenizer),
"encoding_dir": str(base / args.encoding_dir),
"data_dir": str(data_dir),
"cutoff_len": args.cutoff_len,
"eos_token": tok.eos_token,
"eos_token_id": tok.eos_token_id,
"splits": {},
}
for split, filename in split_files.items():
stats = build_split(
split,
data_dir / filename,
out_dir / f"{split}_dsv4_chat_tokenized{suffix}",
tok,
enc,
args.cutoff_len,
)
all_stats["splits"][split] = stats
print(json.dumps(stats, ensure_ascii=False), flush=True)
(out_dir / "build_stats.json").write_text(json.dumps(all_stats, ensure_ascii=False, indent=2), encoding="utf-8")
print(out_dir / "build_stats.json")
if __name__ == "__main__":
main()

208
dataset_building/build_dsv4_chat_tokenized_messages_jsonl.py Normal file
View File

@@ -0,0 +1,208 @@
#!/usr/bin/env python3
import argparse
import gzip
import json
import sys
from collections import Counter
from pathlib import Path
from transformers import AutoTokenizer
def import_encoding(encoding_dir: Path):
sys.path.insert(0, str(encoding_dir))
import encoding_dsv4 # type: ignore
return encoding_dsv4
def open_reader(path: Path):
if path.suffix == ".gz":
return gzip.open(path, "rt", encoding="utf-8")
return path.open("r", encoding="utf-8")
def open_writer(path: Path):
path.parent.mkdir(parents=True, exist_ok=True)
if path.suffix == ".gz":
return gzip.open(path, "wt", encoding="utf-8")
return path.open("w", encoding="utf-8")
def quantiles(xs):
if not xs:
return {}
xs = sorted(xs)
return {
"p50": xs[int((len(xs) - 1) * 0.50)],
"p90": xs[int((len(xs) - 1) * 0.90)],
"p95": xs[int((len(xs) - 1) * 0.95)],
"p99": xs[int((len(xs) - 1) * 0.99)],
"max": xs[-1],
}
def first_user_assistant(messages):
user = None
assistant = None
for msg in messages or []:
role = msg.get("role")
content = (msg.get("content") or "").strip()
if not content:
continue
if role == "user" and user is None:
user = content
elif role == "assistant" and user is not None:
assistant = content
break
if not user or not assistant:
return None
return [
{"role": "user", "content": user},
{"role": "assistant", "content": assistant},
]
def tokenize_split(name, src_path, out_path, tok, enc, cutoff_len, max_rows):
stats = {
"split": name,
"source": str(src_path),
"output": str(out_path),
"cutoff_len": cutoff_len,
"rows_seen": 0,
"rows_written": 0,
"skipped_no_messages": 0,
"truncated": 0,
"eos_in_labels": 0,
"prefix_mismatch": 0,
"capability_counts": Counter(),
"source_counts": Counter(),
"prompt_tokens": [],
"response_tokens": [],
"total_tokens": [],
}
with open_reader(src_path) as src, open_writer(out_path) as dst:
for idx, line in enumerate(src):
if max_rows and stats["rows_written"] >= max_rows:
break
line = line.strip()
if not line:
continue
stats["rows_seen"] += 1
row = json.loads(line)
messages_full = first_user_assistant(row.get("messages"))
if not messages_full:
stats["skipped_no_messages"] += 1
continue
messages_prompt = [messages_full[0]]
prompt_text = enc.encode_messages(messages_prompt, thinking_mode="chat")
full_text = enc.encode_messages(messages_full, thinking_mode="chat")
if not full_text.startswith(prompt_text):
stats["prefix_mismatch"] += 1
prompt_ids = tok(prompt_text, add_special_tokens=False).input_ids
full_ids = tok(full_text, add_special_tokens=False).input_ids
response_ids = full_ids[len(prompt_ids) :]
labels = [-100] * len(prompt_ids) + response_ids
truncated = False
if len(full_ids) > cutoff_len:
truncated = True
full_ids = full_ids[:cutoff_len]
labels = labels[:cutoff_len]
stats["truncated"] += 1
eos_in_labels = tok.eos_token_id in [x for x in labels if x != -100]
stats["eos_in_labels"] += int(eos_in_labels)
stats["prompt_tokens"].append(len(prompt_ids))
stats["response_tokens"].append(len(response_ids))
stats["total_tokens"].append(len(prompt_ids) + len(response_ids))
stats["capability_counts"][row.get("capability") or "unknown"] += 1
stats["source_counts"][row.get("source_id") or row.get("source") or "unknown"] += 1
out = {
"id": row.get("id", f"{name}_{idx:08d}"),
"split": name,
"capability": row.get("capability"),
"source": row.get("source_id") or row.get("source"),
"thinking_mode": "chat",
"messages": messages_full,
"prompt_text": prompt_text,
"full_text": full_text,
"prompt_tokens": len(prompt_ids),
"response_tokens": len(response_ids),
"total_tokens": len(prompt_ids) + len(response_ids),
"truncated": truncated,
"eos_in_labels": eos_in_labels,
"input_ids": full_ids,
"labels": labels,
}
dst.write(json.dumps(out, ensure_ascii=False) + "\n")
stats["rows_written"] += 1
if stats["rows_written"] % 50000 == 0:
print(json.dumps({"split": name, "rows_written": stats["rows_written"]}, ensure_ascii=False), flush=True)
for key in ["prompt_tokens", "response_tokens", "total_tokens"]:
stats[key] = quantiles(stats[key])
stats["truncated_rate"] = stats["truncated"] / max(1, stats["rows_written"])
stats["eos_label_rate"] = stats["eos_in_labels"] / max(1, stats["rows_written"])
stats["capability_counts"] = dict(stats["capability_counts"].most_common())
stats["source_counts_top50"] = dict(stats["source_counts"].most_common(50))
stats.pop("source_counts")
return stats
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--base-dir", default="/ssd/yi/Tokenizer_Swap")
parser.add_argument("--tokenizer", default="model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
parser.add_argument("--encoding-dir", default="external/deepseek_v4_encoding")
parser.add_argument("--train-jsonl", required=True)
parser.add_argument("--eval-jsonl", required=True)
parser.add_argument("--out-dir", required=True)
parser.add_argument("--cutoff-len", type=int, default=2048)
parser.add_argument("--max-train-rows", type=int, default=0)
parser.add_argument("--max-eval-rows", type=int, default=0)
parser.add_argument("--gzip", action="store_true")
args = parser.parse_args()
base = Path(args.base_dir)
out_dir = base / args.out_dir
enc = import_encoding(base / args.encoding_dir)
tok = AutoTokenizer.from_pretrained(base / args.tokenizer, trust_remote_code=True)
suffix = ".jsonl.gz" if args.gzip else ".jsonl"
all_stats = {
"tokenizer": str(base / args.tokenizer),
"encoding_dir": str(base / args.encoding_dir),
"cutoff_len": args.cutoff_len,
"eos_token": tok.eos_token,
"eos_token_id": tok.eos_token_id,
"splits": {},
}
all_stats["splits"]["train"] = tokenize_split(
"train",
Path(args.train_jsonl),
out_dir / f"train_dsv4_chat_tokenized{suffix}",
tok,
enc,
args.cutoff_len,
args.max_train_rows,
)
all_stats["splits"]["validation"] = tokenize_split(
"validation",
Path(args.eval_jsonl),
out_dir / f"validation_dsv4_chat_tokenized{suffix}",
tok,
enc,
args.cutoff_len,
args.max_eval_rows,
)
(out_dir / "build_stats.json").write_text(json.dumps(all_stats, ensure_ascii=False, indent=2), encoding="utf-8")
print(out_dir / "build_stats.json", flush=True)
if __name__ == "__main__":
main()

377
dataset_building/build_heldout_public_mcq_benchmark.py Executable file
View File

@@ -0,0 +1,377 @@
#!/usr/bin/env python3
import argparse
import json
import os
import random
from pathlib import Path
from datasets import get_dataset_config_names, load_dataset
MMLU_DATASET = "cais/mmlu"
GPQA_DATASET = "Idavidrein/gpqa"
CEVAL_DATASET = "ceval/ceval-exam"
CMMLU_DATASET = "haonan-li/cmmlu"
MMLU_CONFIGS = [
"abstract_algebra",
"anatomy",
"astronomy",
"business_ethics",
"clinical_knowledge",
"college_biology",
"college_chemistry",
"college_computer_science",
"college_mathematics",
"college_physics",
"computer_security",
"conceptual_physics",
"econometrics",
"electrical_engineering",
"elementary_mathematics",
"formal_logic",
"global_facts",
"high_school_biology",
"high_school_chemistry",
"high_school_computer_science",
"high_school_government_and_politics",
"high_school_macroeconomics",
"high_school_mathematics",
"high_school_physics",
"high_school_statistics",
"international_law",
"jurisprudence",
"machine_learning",
"management",
"marketing",
"medical_genetics",
"moral_disputes",
"nutrition",
"philosophy",
"professional_law",
"professional_medicine",
"professional_psychology",
"public_relations",
"security_studies",
"sociology",
"us_foreign_policy",
"world_religions",
]
CEVAL_CONFIGS = [
"computer_network",
"operating_system",
"computer_architecture",
"college_programming",
"college_physics",
"college_chemistry",
"advanced_mathematics",
"probability_and_statistics",
"discrete_mathematics",
"electrical_engineer",
"metrology_engineer",
"high_school_mathematics",
"high_school_physics",
"high_school_chemistry",
"high_school_biology",
"legal_professional",
"business_administration",
"marxism",
"mao_zedong_thought",
"education_science",
"teacher_qualification",
"modern_chinese_history",
"chinese_language_and_literature",
"logic",
]
CMMLU_CONFIGS = [
"agronomy",
"anatomy",
"ancient_chinese",
"arts",
"astronomy",
"business_ethics",
"chinese_civil_service_exam",
"chinese_driving_rule",
"chinese_food_culture",
"chinese_foreign_policy",
"chinese_history",
"college_actuarial_science",
"college_education",
"college_engineering_hydrology",
"college_law",
"college_mathematics",
"college_medical_statistics",
"college_medicine",
"computer_science",
"conceptual_physics",
"econometrics",
"education",
"electrical_engineering",
"elementary_chinese",
"elementary_commonsense",
"elementary_information_and_technology",
"elementary_mathematics",
"ethnology",
"food_science",
"genetics",
"global_facts",
"high_school_biology",
"high_school_chemistry",
"high_school_geography",
"high_school_mathematics",
"high_school_physics",
"human_sexuality",
"international_law",
"journalism",
"jurisprudence",
"legal_and_moral_basis",
"logical",
"machine_learning",
"management",
"marketing",
"marxist_theory",
"modern_chinese",
"nutrition",
"philosophy",
"professional_accounting",
"professional_law",
"professional_medicine",
"professional_psychology",
"public_relations",
"security_study",
"sociology",
"sports_science",
"traditional_chinese_medicine",
"virology",
"world_history",
"world_religions",
]
def norm(text):
return " ".join(str(text or "").replace("\x00", " ").split()).strip()
def stable_sample(rng, rows, n):
if len(rows) <= n:
return list(rows)
return rng.sample(rows, n)
def add_item(items, prefix, idx, source, subset, question, choices, answer_idx, metadata=None):
question = norm(question)
choices = [norm(x) for x in choices]
if not question or len(choices) < 2:
return
if answer_idx < 0 or answer_idx >= len(choices):
return
answer_text = choices[answer_idx]
if not answer_text:
return
mcq_prompt = (
"Answer the following multiple-choice question. Choose the single best option.\n\n"
f"Question:\n{question}\n\nAnswer:"
)
ppl_text = f"Question: {question}\nCorrect answer: {answer_text}"
items.append(
{
"id": f"{prefix}_{idx:05d}",
"category": prefix,
"source": source,
"subset": subset,
"ppl_text": ppl_text,
"mcq_prompt": mcq_prompt,
"choices": choices,
"answer_idx": int(answer_idx),
"answer_text": answer_text,
"metadata": metadata or {},
}
)
def build_mmlu(rng, quota):
configs = list(MMLU_CONFIGS)
rng.shuffle(configs)
rows = []
per_config = max(20, quota // 24)
for cfg in configs:
print(f"[info] load mmlu {cfg}", flush=True)
try:
ds = load_dataset(MMLU_DATASET, cfg, split="test")
except Exception as exc:
print(f"[warn] skip MMLU {cfg}: {type(exc).__name__}: {str(exc)[:160]}")
continue
sampled = stable_sample(rng, list(ds), min(per_config, len(ds)))
for r in sampled:
rows.append((cfg, r))
if len(rows) >= quota * 2:
break
rng.shuffle(rows)
items = []
for i, (cfg, r) in enumerate(rows[:quota]):
add_item(
items,
"mmlu",
i,
MMLU_DATASET,
cfg,
r.get("question"),
r.get("choices", []),
r.get("answer"),
{"subject": r.get("subject", cfg)},
)
return items
def build_gpqa(rng, quota, config="gpqa_main"):
ds = load_dataset(GPQA_DATASET, config, split="train", token=True)
rows = stable_sample(rng, list(ds), min(quota, len(ds)))
items = []
for i, r in enumerate(rows):
choices = [
r.get("Correct Answer"),
r.get("Incorrect Answer 1"),
r.get("Incorrect Answer 2"),
r.get("Incorrect Answer 3"),
]
order = list(range(4))
rng.shuffle(order)
shuffled = [choices[j] for j in order]
answer_idx = order.index(0)
add_item(
items,
"gpqa",
i,
GPQA_DATASET,
config,
r.get("Question"),
shuffled,
answer_idx,
{
"high_level_domain": r.get("High-level domain"),
"subdomain": r.get("Subdomain"),
"record_id": r.get("Record ID"),
},
)
return items
def letter_answer_idx(ans):
if ans is None:
return -1
s = str(ans).strip()
if s in {"0", "1", "2", "3"}:
return int(s)
if s:
ch = s[0].upper()
if ch in "ABCD":
return ord(ch) - ord("A")
return -1
def build_abcd_dataset(rng, dataset_name, prefix, quota):
if prefix == "ceval":
configs = list(CEVAL_CONFIGS)
elif prefix == "cmmlu":
configs = list(CMMLU_CONFIGS)
else:
configs = get_dataset_config_names(dataset_name)
rng.shuffle(configs)
rows = []
per_config = max(24, quota // 18)
for cfg in configs:
print(f"[info] load {prefix} {cfg}", flush=True)
split = "test"
try:
ds = load_dataset(dataset_name, cfg, split=split)
except Exception:
try:
split = "val"
ds = load_dataset(dataset_name, cfg, split=split)
except Exception as exc:
print(f"[warn] skip {prefix} {cfg}: {type(exc).__name__}: {str(exc)[:160]}")
continue
sampled = stable_sample(rng, list(ds), min(per_config, len(ds)))
for r in sampled:
rows.append((cfg, split, r))
if len(rows) >= quota * 2:
break
rng.shuffle(rows)
items = []
for i, (cfg, split, r) in enumerate(rows[:quota]):
choices = [r.get("A"), r.get("B"), r.get("C"), r.get("D")]
add_item(
items,
prefix,
i,
dataset_name,
cfg,
r.get("question"),
choices,
letter_answer_idx(r.get("answer")),
{"split": split},
)
return items
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--out-dir", required=True)
ap.add_argument("--seed", type=int, default=20260607)
ap.add_argument("--mmlu", type=int, default=700)
ap.add_argument("--gpqa", type=int, default=400)
ap.add_argument("--ceval", type=int, default=450)
ap.add_argument("--cmmlu", type=int, default=450)
args = ap.parse_args()
rng = random.Random(args.seed)
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
print("[info] HF_ENDPOINT", os.environ.get("HF_ENDPOINT", ""))
builders = [
("mmlu", lambda: build_mmlu(rng, args.mmlu)),
("gpqa", lambda: build_gpqa(rng, args.gpqa)),
("ceval", lambda: build_abcd_dataset(rng, CEVAL_DATASET, "ceval", args.ceval)),
("cmmlu", lambda: build_abcd_dataset(rng, CMMLU_DATASET, "cmmlu", args.cmmlu)),
]
all_items = []
errors = {}
for name, fn in builders:
if getattr(args, name) <= 0:
print(f"[info] skip {name}: quota=0")
continue
try:
items = fn()
print(f"[info] built {name}: {len(items)}")
all_items.extend(items)
except Exception as exc:
errors[name] = f"{type(exc).__name__}: {str(exc)[:500]}"
print(f"[error] {name}: {errors[name]}")
rng.shuffle(all_items)
bench_path = out_dir / "heldout_public_mcq_2k.jsonl"
with bench_path.open("w", encoding="utf-8") as f:
for item in all_items:
f.write(json.dumps(item, ensure_ascii=False) + "\n")
counts = {}
for item in all_items:
counts[item["category"]] = counts.get(item["category"], 0) + 1
stats = {
"seed": args.seed,
"requested": {"mmlu": args.mmlu, "gpqa": args.gpqa, "ceval": args.ceval, "cmmlu": args.cmmlu},
"total_items": len(all_items),
"counts": counts,
"errors": errors,
"output": str(bench_path),
}
with (out_dir / "heldout_public_mcq_2k_stats.json").open("w", encoding="utf-8") as f:
json.dump(stats, f, ensure_ascii=False, indent=2)
print(json.dumps(stats, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()

264
dataset_building/build_math_docmix_fix.py Executable file
View File

@@ -0,0 +1,264 @@
#!/usr/bin/env python3
import argparse
import gzip
import json
import os
import re
import time
from collections import Counter
from pathlib import Path
from huggingface_hub import HfApi, hf_hub_download
from transformers import AutoTokenizer
MATHPILE_PRIORITY = [
"train/textbooks/textbooks_markdown.jsonl.gz",
"train/textbooks/synthetic_textbooks_markdown.jsonl.gz",
"train/wikipedia/wikipedia_en_mathematics_nopic_2023-08_v0.2.jsonl.gz",
"train/proofwiki/ProofWiki_definitions.jsonl.gz",
"train/proofwiki/ProofWiki_theorem_proofs.jsonl.gz",
"train/stackexchange/math.stackexchange.com.jsonl.gz",
"train/stackexchange/mathoverflow.net.jsonl.gz",
"train/stackexchange/physics.stackexchange.com.jsonl.gz",
"train/arXiv/math_arXiv_v0.2_chunk_1.jsonl.gz",
"train/arXiv/math_arXiv_v0.2_chunk_2.jsonl.gz",
"train/arXiv/math_arXiv_v0.2_chunk_3.jsonl.gz",
"train/arXiv/math_arXiv_v0.2_chunk_4.jsonl.gz",
"train/commoncrawl/C4_math_docs_chunk_0.jsonl.gz",
"train/commoncrawl/CC_math_docs_chunk_0.jsonl.gz",
]
def clean_text(text):
if text is None:
return ""
text = str(text).replace("\x00", " ")
text = re.sub(r"[ \t\r\f\v]+", " ", text)
text = re.sub(r"\n{4,}", "\n\n\n", text)
return text.strip()
def extract_text(row):
for key in ("text", "content", "markdown", "raw_content", "document"):
if isinstance(row, dict):
text = clean_text(row.get(key))
if text:
return text
return ""
def safe_source(text):
return re.sub(r"[^A-Za-z0-9._-]+", "_", text)[:120]
def download_with_retry(repo, filename, args):
last = None
for attempt in range(1, args.retries + 1):
try:
return hf_hub_download(
repo_id=repo,
repo_type="dataset",
filename=filename,
endpoint=args.endpoint,
token=os.environ.get("HF_TOKEN"),
local_dir=args.raw_dir,
)
except Exception as exc:
last = exc
print(json.dumps({"event": "download_retry", "repo": repo, "file": filename, "attempt": attempt, "error": repr(exc)[:800]}, ensure_ascii=False), flush=True)
time.sleep(min(120, 5 * attempt))
raise RuntimeError(f"download failed for {repo}:{filename}: {last!r}")
def copy_existing(existing, writer, target):
tokens = 0
docs = 0
if not existing.exists() or existing.stat().st_size < 1024:
return tokens, docs
with gzip.open(existing, "rt", encoding="utf-8", errors="replace") as f:
for line in f:
if not line.strip() or tokens >= target:
continue
try:
rec = json.loads(line)
except json.JSONDecodeError:
continue
ntok = int(rec.get("token_count") or 0)
if ntok <= 0:
continue
writer.write(json.dumps(rec, ensure_ascii=False) + "\n")
tokens += ntok
docs += 1
return tokens, docs
def iter_jsonl_gz(path):
opener = gzip.open if str(path).endswith(".gz") else open
with opener(path, "rt", encoding="utf-8", errors="replace") as f:
for line in f:
if not line.strip():
continue
try:
yield json.loads(line)
except json.JSONDecodeError:
continue
def add_jsonl_source(path, source_label, tok, writer, stats, args):
for row in iter_jsonl_gz(path):
if stats["tokens"] >= args.target_tokens:
break
text = extract_text(row)
if not text:
stats["rejected"]["empty"] += 1
continue
ntok = len(tok.encode(text, add_special_tokens=False))
if ntok < args.min_tokens:
stats["rejected"]["too_short"] += 1
continue
if ntok > args.max_doc_tokens:
stats["rejected"]["too_long"] += 1
continue
idx = stats["docs"]
writer.write(json.dumps({
"id": f"math_fix_{idx:09d}",
"category": "math",
"source": source_label,
"text": text,
"token_count": ntok,
"metadata": {k: row.get(k) for k in ("id", "url", "source", "title") if isinstance(row, dict) and k in row},
}, ensure_ascii=False) + "\n")
stats["docs"] += 1
stats["tokens"] += ntok
stats["tokens_by_source"][source_label] += ntok
if stats["docs"] % args.log_every == 0:
print(json.dumps({"event": "progress", "docs": stats["docs"], "tokens": stats["tokens"], "target": args.target_tokens, "source": source_label, "elapsed_sec": time.time() - stats["started_at"]}, ensure_ascii=False), flush=True)
def add_parquet_source(path, source_label, tok, writer, stats, args):
import pyarrow.parquet as pq
pf = pq.ParquetFile(path)
for batch in pf.iter_batches(batch_size=args.parquet_batch_size):
if stats["tokens"] >= args.target_tokens:
break
for row in batch.to_pylist():
if stats["tokens"] >= args.target_tokens:
break
text = extract_text(row)
if not text:
stats["rejected"]["empty"] += 1
continue
ntok = len(tok.encode(text, add_special_tokens=False))
if ntok < args.min_tokens:
stats["rejected"]["too_short"] += 1
continue
if ntok > args.max_doc_tokens:
stats["rejected"]["too_long"] += 1
continue
idx = stats["docs"]
writer.write(json.dumps({
"id": f"math_fix_{idx:09d}",
"category": "math",
"source": source_label,
"text": text,
"token_count": ntok,
"metadata": {k: row.get(k) for k in ("url", "date", "language", "language_score") if k in row},
}, ensure_ascii=False) + "\n")
stats["docs"] += 1
stats["tokens"] += ntok
stats["tokens_by_source"][source_label] += ntok
if stats["docs"] % args.log_every == 0:
print(json.dumps({"event": "progress", "docs": stats["docs"], "tokens": stats["tokens"], "target": args.target_tokens, "source": source_label, "elapsed_sec": time.time() - stats["started_at"]}, ensure_ascii=False), flush=True)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--endpoint", default=os.environ.get("HF_ENDPOINT", "https://hf-mirror.com"))
ap.add_argument("--tokenizer", default="model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
ap.add_argument("--raw-dir", default="data/raw_math_fix_20260614")
ap.add_argument("--out-dir", default="data/cpt_docmix_5b_sources_8192_20260614")
ap.add_argument("--target-tokens", type=int, default=500_000_000)
ap.add_argument("--min-tokens", type=int, default=128)
ap.add_argument("--max-doc-tokens", type=int, default=32768)
ap.add_argument("--log-every", type=int, default=5000)
ap.add_argument("--retries", type=int, default=16)
ap.add_argument("--parquet-batch-size", type=int, default=1000)
ap.add_argument("--keep-raw", action="store_true")
ap.add_argument("--skip-mathpile", action="store_true")
ap.add_argument("--skip-openwebmath-prefix-count", type=int, default=0)
args = ap.parse_args()
base = Path.cwd()
out_dir = base / args.out_dir
doc_dir = out_dir / "documents"
doc_dir.mkdir(parents=True, exist_ok=True)
raw_dir = base / args.raw_dir
raw_dir.mkdir(parents=True, exist_ok=True)
args.raw_dir = str(raw_dir)
tok = AutoTokenizer.from_pretrained(base / args.tokenizer, trust_remote_code=True)
final_out = doc_dir / "math.jsonl.gz"
tmp_out = doc_dir / "math.jsonl.gz.tmp"
if tmp_out.exists():
tmp_out.unlink()
stats = {"target_tokens": args.target_tokens, "tokens": 0, "docs": 0, "tokens_by_source": Counter(), "rejected": Counter(), "sources": [], "started_at": time.time()}
api = HfApi(endpoint=args.endpoint, token=os.environ.get("HF_TOKEN"))
with gzip.open(tmp_out, "wt", encoding="utf-8") as writer:
existing_tokens, existing_docs = copy_existing(final_out, writer, args.target_tokens)
stats["tokens"] += existing_tokens
stats["docs"] += existing_docs
stats["tokens_by_source"]["existing_math_docmix"] += existing_tokens
print(json.dumps({"event": "copied_existing", "docs": existing_docs, "tokens": existing_tokens}, ensure_ascii=False), flush=True)
for filename in ([] if args.skip_mathpile else MATHPILE_PRIORITY):
if stats["tokens"] >= args.target_tokens:
break
try:
local = Path(download_with_retry("GAIR/MathPile", filename, args))
before = stats["tokens"]
add_jsonl_source(local, f"GAIR/MathPile:{filename}", tok, writer, stats, args)
stats["sources"].append({"repo": "GAIR/MathPile", "file": filename, "tokens": stats["tokens"] - before})
print(json.dumps({"event": "source_done", "source": f"GAIR/MathPile:{filename}", "tokens_added": stats["tokens"] - before, "total_tokens": stats["tokens"]}, ensure_ascii=False), flush=True)
if not args.keep_raw:
try: local.unlink()
except Exception: pass
except Exception as exc:
stats["sources"].append({"repo": "GAIR/MathPile", "file": filename, "error": repr(exc)[:1000]})
print(json.dumps({"event": "source_error", "source": f"GAIR/MathPile:{filename}", "error": repr(exc)[:1000]}, ensure_ascii=False), flush=True)
if stats["tokens"] < args.target_tokens:
files = [f for f in api.list_repo_files("open-web-math/open-web-math", repo_type="dataset") if f.startswith("data/") and f.endswith(".parquet")]
for filename in sorted(files)[args.skip_openwebmath_prefix_count:]:
if stats["tokens"] >= args.target_tokens:
break
try:
local = Path(download_with_retry("open-web-math/open-web-math", filename, args))
before = stats["tokens"]
add_parquet_source(local, f"open-web-math/open-web-math:{filename}", tok, writer, stats, args)
stats["sources"].append({"repo": "open-web-math/open-web-math", "file": filename, "tokens": stats["tokens"] - before})
print(json.dumps({"event": "source_done", "source": f"open-web-math/open-web-math:{filename}", "tokens_added": stats["tokens"] - before, "total_tokens": stats["tokens"]}, ensure_ascii=False), flush=True)
if not args.keep_raw:
try: local.unlink()
except Exception: pass
except Exception as exc:
stats["sources"].append({"repo": "open-web-math/open-web-math", "file": filename, "error": repr(exc)[:1000]})
print(json.dumps({"event": "source_error", "source": f"open-web-math/open-web-math:{filename}", "error": repr(exc)[:1000]}, ensure_ascii=False), flush=True)
if stats["tokens"] < args.target_tokens:
raise SystemExit(f"only collected {stats['tokens']} / {args.target_tokens} tokens")
if final_out.exists():
final_out.replace(final_out.with_suffix(".jsonl.gz.underfilled_20260614"))
tmp_out.replace(final_out)
stats["elapsed_sec"] = time.time() - stats["started_at"]
stats["tokens_by_source"] = dict(stats["tokens_by_source"])
stats["rejected"] = dict(stats["rejected"])
(out_dir / "math_5b_fix_stats.json").write_text(json.dumps(stats, ensure_ascii=False, indent=2), encoding="utf-8")
(out_dir / ".math_5b_ready").write_text(json.dumps({"tokens": stats["tokens"], "docs": stats["docs"], "elapsed_sec": stats["elapsed_sec"]}, ensure_ascii=False), encoding="utf-8")
print(json.dumps({"event": "done", "tokens": stats["tokens"], "docs": stats["docs"], "elapsed_sec": stats["elapsed_sec"], "output": str(final_out)}, ensure_ascii=False), flush=True)
if __name__ == "__main__":
main()

208
dataset_building/build_science_docmix_fix.py Executable file
View File

@@ -0,0 +1,208 @@
#!/usr/bin/env python3
import argparse
import gzip
import json
import re
import time
from collections import Counter
from pathlib import Path
from transformers import AutoTokenizer
SCIENCE_KEYWORDS = re.compile(
r"\b(physics|chemistry|biology|biological|medical|medicine|clinical|anatomy|physiology|genetics|ecology|astronomy|geology|neuroscience|experiment|hypothesis|laboratory|scientific|research|disease|protein|cell|molecule|atom|energy|force|gravity|electric|magnetic|quantum|planet|star|organism|evolution|climate|weather|ecosystem|bacteria|virus|vaccine|enzyme|hormone|blood|brain|heart|lung|kidney|cancer|therapy|diagnosis)\b",
re.I,
)
def clean_text(text):
if text is None:
return ""
text = str(text).replace("\x00", " ")
text = re.sub(r"[ \t\r\f\v]+", " ", text)
text = re.sub(r"\n{4,}", "\n\n\n", text)
return text.strip()
def first_user_assistant(messages):
user = None
assistant = None
for msg in messages or []:
role = msg.get("role")
content = clean_text(msg.get("content"))
if not content:
continue
if role == "user" and user is None:
user = content
elif role == "assistant" and user is not None:
assistant = content
break
return user, assistant
def qa_text(user, assistant):
if not user or not assistant:
return ""
return f"Question:\n{user}\n\nAnswer:\n{assistant}"
def extract_row_text(row):
if isinstance(row, dict) and row.get("messages"):
u, a = first_user_assistant(row.get("messages"))
if u and a:
return qa_text(u, a)
for key in ("text", "content", "markdown", "raw_content"):
if isinstance(row, dict):
text = clean_text(row.get(key))
if text:
return text
return ""
def write_record(writer, stats, source, text, ntok, metadata=None):
idx = stats["docs"]
writer.write(json.dumps({
"id": f"science_fix_{idx:09d}",
"category": "science",
"source": source,
"text": text,
"token_count": ntok,
"metadata": metadata or {},
}, ensure_ascii=False) + "\n")
stats["docs"] += 1
stats["tokens"] += ntok
stats["tokens_by_source"][source] += ntok
def copy_existing(existing, writer, target):
tokens = docs = 0
if not existing.exists() or existing.stat().st_size < 1024:
return tokens, docs
with gzip.open(existing, "rt", encoding="utf-8", errors="replace") as f:
for line in f:
if not line.strip() or tokens >= target:
continue
try:
rec = json.loads(line)
except json.JSONDecodeError:
continue
ntok = int(rec.get("token_count") or 0)
if ntok <= 0:
continue
writer.write(json.dumps(rec, ensure_ascii=False) + "\n")
tokens += ntok
docs += 1
return tokens, docs
def add_jsonl(path, source, tok, writer, stats, args, capability_filter=None):
with open(path, "r", encoding="utf-8", errors="replace") as f:
for line in f:
if stats["tokens"] >= args.target_tokens:
break
if not line.strip():
continue
try:
row = json.loads(line)
except json.JSONDecodeError:
continue
if capability_filter and row.get("capability") not in capability_filter:
continue
text = extract_row_text(row)
if not text:
stats["rejected"]["empty"] += 1
continue
ntok = len(tok.encode(text, add_special_tokens=False))
if ntok < args.min_tokens:
stats["rejected"]["too_short"] += 1
continue
if ntok > args.max_doc_tokens:
stats["rejected"]["too_long"] += 1
continue
write_record(writer, stats, source, text, ntok, {k: row.get(k) for k in ("id", "capability", "source_id", "split") if k in row})
if stats["docs"] % args.log_every == 0:
print(json.dumps({"event": "progress", "docs": stats["docs"], "tokens": stats["tokens"], "target": args.target_tokens, "source": source, "elapsed_sec": time.time() - stats["started_at"]}, ensure_ascii=False), flush=True)
def add_fineweb_edu_parquets(paths, tok, writer, stats, args):
import pyarrow.parquet as pq
for path in paths:
if stats["tokens"] >= args.target_tokens:
break
source = f"fineweb_edu_science_supplement:{path.name}"
before = stats["tokens"]
pf = pq.ParquetFile(path)
for batch in pf.iter_batches(batch_size=args.parquet_batch_size):
if stats["tokens"] >= args.target_tokens:
break
for row in batch.to_pylist():
if stats["tokens"] >= args.target_tokens:
break
text = extract_row_text(row)
if not text or not SCIENCE_KEYWORDS.search(text[:12000]):
stats["rejected"]["not_science_like"] += 1
continue
ntok = len(tok.encode(text, add_special_tokens=False))
if ntok < args.min_tokens:
stats["rejected"]["too_short"] += 1
continue
if ntok > args.max_doc_tokens:
stats["rejected"]["too_long"] += 1
continue
write_record(writer, stats, source, text, ntok, {k: row.get(k) for k in ("url", "dump", "language", "language_score") if k in row})
if stats["docs"] % args.log_every == 0:
print(json.dumps({"event": "progress", "docs": stats["docs"], "tokens": stats["tokens"], "target": args.target_tokens, "source": source, "elapsed_sec": time.time() - stats["started_at"]}, ensure_ascii=False), flush=True)
print(json.dumps({"event": "source_done", "source": source, "tokens_added": stats["tokens"] - before, "total_tokens": stats["tokens"]}, ensure_ascii=False), flush=True)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--tokenizer", default="model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
ap.add_argument("--out-dir", default="data/cpt_docmix_5b_sources_8192_20260614")
ap.add_argument("--target-tokens", type=int, default=150_000_000)
ap.add_argument("--min-tokens", type=int, default=128)
ap.add_argument("--max-doc-tokens", type=int, default=32768)
ap.add_argument("--log-every", type=int, default=5000)
ap.add_argument("--parquet-batch-size", type=int, default=1000)
args = ap.parse_args()
base = Path.cwd()
out_dir = base / args.out_dir
doc_dir = out_dir / "documents"
final_out = doc_dir / "science.jsonl.gz"
tmp_out = doc_dir / "science.jsonl.gz.tmp"
if tmp_out.exists():
tmp_out.unlink()
tok = AutoTokenizer.from_pretrained(base / args.tokenizer, trust_remote_code=True)
stats = {"target_tokens": args.target_tokens, "tokens": 0, "docs": 0, "tokens_by_source": Counter(), "rejected": Counter(), "started_at": time.time(), "sources": []}
with gzip.open(tmp_out, "wt", encoding="utf-8") as writer:
t, d = copy_existing(final_out, writer, args.target_tokens)
stats["tokens"] += t; stats["docs"] += d; stats["tokens_by_source"]["existing_science_docmix"] += t
print(json.dumps({"event": "copied_existing", "docs": d, "tokens": t}, ensure_ascii=False), flush=True)
mix = base / "data/training_mix_v4_train1m_test2p8k_noupsample_nobbh_20260611/train_1m.jsonl"
if mix.exists() and stats["tokens"] < args.target_tokens:
before = stats["tokens"]
add_jsonl(mix, "training_mix_v4_science_logic", tok, writer, stats, args, {"science_reasoning", "logic"})
stats["sources"].append({"source": str(mix), "tokens": stats["tokens"] - before})
paths = sorted((base / "data/raw_parquets/fineweb_edu").glob("*.parquet"))
if paths and stats["tokens"] < args.target_tokens:
add_fineweb_edu_parquets(paths, tok, writer, stats, args)
if stats["tokens"] < args.target_tokens:
raise SystemExit(f"only collected {stats['tokens']} / {args.target_tokens} tokens")
if final_out.exists():
final_out.replace(final_out.with_suffix(".jsonl.gz.underfilled_20260614"))
tmp_out.replace(final_out)
stats["elapsed_sec"] = time.time() - stats["started_at"]
stats["tokens_by_source"] = dict(stats["tokens_by_source"])
stats["rejected"] = dict(stats["rejected"])
(out_dir / "science_5b_fix_stats.json").write_text(json.dumps(stats, ensure_ascii=False, indent=2), encoding="utf-8")
(out_dir / ".science_5b_ready").write_text(json.dumps({"tokens": stats["tokens"], "docs": stats["docs"], "elapsed_sec": stats["elapsed_sec"]}, ensure_ascii=False), encoding="utf-8")
print(json.dumps({"event": "done", "tokens": stats["tokens"], "docs": stats["docs"], "elapsed_sec": stats["elapsed_sec"], "output": str(final_out)}, ensure_ascii=False), flush=True)
if __name__ == "__main__":
main()

745
dataset_building/build_training_and_test_mix_v3.py Normal file
View File

@@ -0,0 +1,745 @@
#!/usr/bin/env python3
import argparse
import copy
import csv
import hashlib
import json
import os
import random
import re
from collections import Counter, defaultdict
from pathlib import Path
from tokenizers import Tokenizer
from transformers import AutoTokenizer
GROUPS = [
"chinese_exam",
"chinese_dialogue",
"code",
"math",
"logic",
"science_reasoning",
"english_dialogue",
]
TRAIN_RATIOS = {
"chinese_exam": 0,
"chinese_dialogue": 35,
"code": 30,
"math": 20,
"logic": 5,
"science_reasoning": 5,
"english_dialogue": 5,
}
LOCAL_TRAIN_SOURCES = [
("data/open_recovery_sft_mix_alt_sources_1m_parquet_20260607/normalized.jsonl", "normalized"),
("data/open_recovery_sft_mix_100k_4_3_2_1_20260604/normalized.jsonl", "normalized"),
("data/modelscope_alt_sources_20260607/CodeAlpaca-20k.jsonl", "jsonl:code:AI-ModelScope/CodeAlpaca-20k"),
("data/modelscope_alt_sources_20260607/alpaca-gpt4-data-zh_train.csv", "csv:chinese_dialogue:AI-ModelScope/alpaca-gpt4-data-zh"),
]
HF_TRAIN_SOURCES = {
"chinese_exam": [],
"chinese_dialogue": [
{"name": "BelleGroup/train_0.5M_CN", "splits": ["train"]},
{"name": "m-a-p/COIG-CQIA", "configs": ["coig_pc", "zhihu", "wikihow"], "splits": ["train"]},
],
"code": [
{"name": "ise-uiuc/Magicoder-OSS-Instruct-75K", "splits": ["train"]},
{"name": "bigcode/self-oss-instruct-sc2-exec-filter-50k", "splits": ["train"]},
{"name": "nvidia/OpenCodeInstruct", "splits": ["train"]},
{"name": "HuggingFaceTB/smoltalk", "configs": ["self-oss-instruct", "apigen-80k"], "splits": ["train"]},
],
"math": [
{"name": "TIGER-Lab/MathInstruct", "splits": ["train"]},
{"name": "nvidia/OpenMathInstruct-2", "splits": ["train"]},
{"name": "HuggingFaceTB/smoltalk", "configs": ["numina-cot-100k", "metamathqa-50k"], "splits": ["train"]},
],
"logic": [
{"name": "metaeval/reclor", "splits": ["train", "validation"]},
{"name": "tasksource/bigbench", "configs": [
"causal_judgment",
"date_understanding",
"disambiguation_qa",
"logical_args",
"logical_deduction",
"logical_fallacy_detection",
"social_iqa",
"strategyqa",
"temporal_sequences",
], "splits": ["train"]},
{"name": "tau/commonsense_qa", "splits": ["train", "validation"]},
],
"science_reasoning": [
{"name": "allenai/ai2_arc", "configs": ["ARC-Challenge"], "splits": ["train", "validation"]},
{"name": "allenai/ai2_arc", "configs": ["ARC-Easy"], "splits": ["train", "validation"]},
{"name": "allenai/qasc", "splits": ["train", "validation"]},
{"name": "allenai/openbookqa", "configs": ["main", "additional"], "splits": ["train", "validation"]},
{"name": "sciq", "splits": ["train", "validation"]},
{"name": "qiaojin/PubMedQA", "configs": ["pqa_labeled"], "splits": ["train"]},
],
"english_dialogue": [
{"name": "HuggingFaceTB/smoltalk", "configs": ["all"], "splits": ["train"]},
{"name": "allenai/tulu-3-sft-mixture", "splits": ["train"]},
{"name": "HuggingFaceH4/ultrachat_200k", "splits": ["train_sft", "train"]},
],
}
HF_TEST_SOURCES = {
"chinese_exam": [
{"name": "ceval/ceval-exam", "configs": "all", "splits": ["val", "dev"]},
],
"chinese_dialogue": [
{"name": "BelleGroup/train_0.5M_CN", "splits": ["train"]},
{"name": "m-a-p/COIG-CQIA", "configs": ["zhihu", "wikihow", "coig_pc"], "splits": ["train"]},
],
"code": [
{"name": "openai/openai_humaneval", "splits": ["test"]},
{"name": "google-research-datasets/mbpp", "splits": ["test", "validation", "train"]},
],
"math": [
{"name": "gsm8k", "configs": ["main"], "splits": ["test"]},
{"name": "hendrycks/competition_math", "splits": ["test"]},
],
"logic": [
{"name": "metaeval/reclor", "splits": ["test", "validation"]},
{"name": "lighteval/bbh", "configs": ["logical_deduction_three_objects", "logical_deduction_five_objects", "logical_deduction_seven_objects"], "splits": ["test", "train"]},
{"name": "cais/mmlu", "configs": ["formal_logic", "logical_fallacies"], "splits": ["test", "validation", "dev"]},
],
"science_reasoning": [
{"name": "Idavidrein/gpqa", "splits": ["train"]},
{"name": "allenai/ai2_arc", "configs": ["ARC-Challenge"], "splits": ["test", "validation"]},
{"name": "sciq", "splits": ["test", "validation"]},
],
"english_dialogue": [
{"name": "HuggingFaceH4/ultrachat_200k", "splits": ["test_sft", "test", "train_sft"]},
{"name": "HuggingFaceTB/smoltalk", "configs": ["all"], "splits": ["test", "train"]},
],
}
def clean_text(x):
if x is None:
return ""
if not isinstance(x, str):
x = str(x)
x = x.replace("\r\n", "\n").replace("\r", "\n")
return re.sub(r"\n{4,}", "\n\n\n", x).strip()
def normalize_text_for_hash(text):
text = clean_text(text).lower()
return re.sub(r"\s+", " ", text)
def pair_hash(user, assistant):
payload = normalize_text_for_hash(user) + "\n\n---\n\n" + normalize_text_for_hash(assistant)
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
def prompt_hash(user):
return hashlib.sha256(normalize_text_for_hash(user).encode("utf-8")).hexdigest()
def load_tokenizer(path):
try:
return AutoTokenizer.from_pretrained(path, trust_remote_code=True)
except Exception:
return Tokenizer.from_file(str(Path(path) / "tokenizer.json"))
def token_len(tok, text):
if isinstance(tok, Tokenizer):
return len(tok.encode(text, add_special_tokens=False).ids)
return len(tok(text, add_special_tokens=False)["input_ids"])
def message_content(m):
if isinstance(m, dict):
return clean_text(m.get("content") or m.get("value") or m.get("text") or "")
return clean_text(m)
def role_of(m, default):
if isinstance(m, dict):
role = (m.get("role") or m.get("from") or m.get("speaker") or default).lower()
if role in {"human", "user"}:
return "user"
if role in {"gpt", "assistant", "model"}:
return "assistant"
if role == "system":
return "system"
return default
def normalize_messages(row):
if not isinstance(row, dict):
return None
# Code generation schemas: HumanEval / MBPP-like.
if clean_text(row.get("prompt")) and clean_text(row.get("canonical_solution")):
prompt = clean_text(row.get("prompt"))
tests = clean_text(row.get("test"))
user = "Complete the following Python function.\n\n" + prompt
if tests:
user += "\n\nThe solution should pass these tests:\n" + tests
return [{"role": "user", "content": user}, {"role": "assistant", "content": clean_text(row.get("canonical_solution"))}]
if clean_text(row.get("text")) and clean_text(row.get("code")):
user = clean_text(row.get("text"))
tests = row.get("test_list")
if isinstance(tests, list) and tests:
user += "\n\nTests:\n" + "\n".join(str(x) for x in tests)
return [{"role": "user", "content": user}, {"role": "assistant", "content": clean_text(row.get("code"))}]
# BBH-like schemas.
if clean_text(row.get("input")) and isinstance(row.get("choices"), list) and row.get("target_idx") is not None:
choices_list = [clean_text(x) for x in row.get("choices") if clean_text(x)]
try:
ans_idx = int(row.get("target_idx"))
except Exception:
ans_idx = None
prefix = clean_text(row.get("task_prefix"))
prompt = (prefix + "\n\n" if prefix else "") + clean_text(row.get("input"))
if choices_list:
prompt += "\n" + "\n".join(f"{chr(ord('A') + i)}. {x}" for i, x in enumerate(choices_list))
if ans_idx is not None and 0 <= ans_idx < len(choices_list):
assistant = f"{chr(ord('A') + ans_idx)}. {choices_list[ans_idx]}"
else:
assistant = clean_text(row.get("target_idx"))
if assistant:
return [{"role": "user", "content": prompt}, {"role": "assistant", "content": assistant}]
if clean_text(row.get("input")) and clean_text(row.get("target")):
return [{"role": "user", "content": clean_text(row.get("input"))}, {"role": "assistant", "content": clean_text(row.get("target"))}]
# BIG-bench/tasksource schemas.
if clean_text(row.get("inputs")) and isinstance(row.get("targets"), list):
prompt = clean_text(row.get("inputs"))
targets = [clean_text(x) for x in row.get("targets") if clean_text(x)]
choices = [clean_text(x) for x in row.get("multiple_choice_targets") or [] if clean_text(x)]
scores = row.get("multiple_choice_scores") or []
if choices:
prompt += "\n" + "\n".join(f"{chr(ord('A') + i)}. {x}" for i, x in enumerate(choices))
answer = targets[0] if targets else ""
if choices and scores:
try:
best_idx = max(range(len(scores)), key=lambda i: scores[i])
except Exception:
best_idx = None
if best_idx is not None and 0 <= best_idx < len(choices):
answer = f"{chr(ord('A') + best_idx)}. {choices[best_idx]}"
if answer:
return [{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}]
# ReClor/LogiQA-like schemas.
if clean_text(row.get("context")) and clean_text(row.get("question")) and isinstance(row.get("answers"), list):
answers = [clean_text(x) for x in row.get("answers") if clean_text(x)]
label = row.get("label")
try:
label_idx = int(label)
except Exception:
label_idx = None
prompt = clean_text(row.get("context")) + "\n\nQuestion: " + clean_text(row.get("question"))
prompt += "\n" + "\n".join(f"{chr(ord('A') + i)}. {x}" for i, x in enumerate(answers))
if label_idx is not None and 0 <= label_idx < len(answers):
assistant = f"{chr(ord('A') + label_idx)}. {answers[label_idx]}"
else:
assistant = clean_text(label)
if assistant:
return [{"role": "user", "content": prompt}, {"role": "assistant", "content": assistant}]
# GPQA-like schemas.
gpqa_question = clean_text(row.get("Question") or row.get("question"))
gpqa_correct = clean_text(row.get("Correct Answer") or row.get("correct_answer"))
if gpqa_question and gpqa_correct and any(clean_text(row.get(f"Incorrect Answer {i}")) for i in [1, 2, 3]):
choices = [gpqa_correct] + [clean_text(row.get(f"Incorrect Answer {i}")) for i in [1, 2, 3] if clean_text(row.get(f"Incorrect Answer {i}"))]
prompt = gpqa_question + "\nChoices:\n" + "\n".join(f"- {x}" for x in choices)
return [{"role": "user", "content": prompt}, {"role": "assistant", "content": gpqa_correct}]
# PubMedQA-like schemas.
if clean_text(row.get("question")) and isinstance(row.get("context"), dict) and clean_text(row.get("long_answer")):
contexts = row.get("context", {}).get("contexts") or []
context_text = "\n".join(clean_text(x) for x in contexts if clean_text(x))
prompt = clean_text(row.get("question"))
if context_text:
prompt = context_text + "\n\nQuestion: " + prompt
answer = clean_text(row.get("final_decision"))
long_answer = clean_text(row.get("long_answer"))
assistant = answer + "\n\n" + long_answer if answer else long_answer
return [{"role": "user", "content": prompt}, {"role": "assistant", "content": assistant}]
# MedMCQA-like schemas.
if clean_text(row.get("question")) and all(clean_text(row.get(k)) for k in ["opa", "opb", "opc", "opd"]) and row.get("cop") is not None:
labels = ["A", "B", "C", "D"]
values = [clean_text(row.get(k)) for k in ["opa", "opb", "opc", "opd"]]
try:
ans_idx = int(row.get("cop"))
except Exception:
ans_idx = None
prompt = clean_text(row.get("question")) + "\n" + "\n".join(f"{label}. {value}" for label, value in zip(labels, values))
if ans_idx is not None and 0 <= ans_idx < len(values):
assistant = f"{labels[ans_idx]}. {values[ans_idx]}"
else:
assistant = clean_text(row.get("cop"))
exp = clean_text(row.get("exp"))
if exp:
assistant += "\n\n" + exp
return [{"role": "user", "content": prompt}, {"role": "assistant", "content": assistant}]
# Common MCQ schemas: CEval/CMMLU, ARC, SciQ.
question = clean_text(row.get("question") or row.get("question_stem") or row.get("formatted_question"))
if question:
choices = []
answer_text = ""
explanation = clean_text(row.get("explanation") or row.get("support") or "")
for label in ["A", "B", "C", "D", "E"]:
if clean_text(row.get(label)):
choices.append((label, clean_text(row.get(label))))
if choices and clean_text(row.get("answer")):
ans = clean_text(row.get("answer")).strip()
answer_value = dict(choices).get(ans, ans)
prompt = question + "\n" + "\n".join(f"{label}. {text}" for label, text in choices)
answer_text = f"{ans}. {answer_value}"
if explanation:
answer_text += "\n\n" + explanation
return [{"role": "user", "content": prompt}, {"role": "assistant", "content": answer_text}]
# MMLU/ScienceQA-style: choices is a list and answer is an integer index.
if isinstance(row.get("choices"), list) and row.get("answer") is not None:
choices_list = [clean_text(x) for x in row.get("choices") if clean_text(x)]
try:
ans_idx = int(row.get("answer"))
except Exception:
ans_idx = None
prompt_parts = [question]
hint = clean_text(row.get("hint"))
lecture = clean_text(row.get("lecture"))
solution = clean_text(row.get("solution"))
if hint:
prompt_parts.append("Hint: " + hint)
prompt_parts.append("\n".join(f"{chr(ord('A') + i)}. {x}" for i, x in enumerate(choices_list)))
prompt = "\n\n".join(x for x in prompt_parts if x)
if ans_idx is not None and 0 <= ans_idx < len(choices_list):
assistant = f"{chr(ord('A') + ans_idx)}. {choices_list[ans_idx]}"
else:
assistant = clean_text(row.get("answer"))
extra = "\n\n".join(x for x in [lecture, solution] if x)
if assistant and extra:
assistant += "\n\n" + extra
if assistant:
return [{"role": "user", "content": prompt}, {"role": "assistant", "content": assistant}]
arc_choices = row.get("choices")
if isinstance(arc_choices, dict) and arc_choices.get("text") and row.get("answerKey"):
labels = arc_choices.get("label") or [chr(ord("A") + i) for i in range(len(arc_choices["text"]))]
choices = list(zip(labels, arc_choices["text"]))
ans = clean_text(row.get("answerKey"))
answer_value = dict(choices).get(ans, ans)
prompt = question + "\n" + "\n".join(f"{label}. {text}" for label, text in choices)
facts = [clean_text(row.get(k)) for k in ["combinedfact", "fact1", "fact2"] if clean_text(row.get(k))]
assistant = f"{ans}. {answer_value}"
if facts:
assistant += "\n\n" + "\n".join(facts)
return [{"role": "user", "content": prompt}, {"role": "assistant", "content": assistant}]
if row.get("correct_answer"):
distractors = [clean_text(row.get(k)) for k in ["distractor1", "distractor2", "distractor3"] if clean_text(row.get(k))]
prompt = question
if distractors:
all_choices = distractors + [clean_text(row.get("correct_answer"))]
prompt += "\nChoices:\n" + "\n".join(f"- {x}" for x in all_choices)
answer_text = clean_text(row.get("correct_answer"))
if explanation:
answer_text += "\n\n" + explanation
return [{"role": "user", "content": prompt}, {"role": "assistant", "content": answer_text}]
for key in ("messages", "conversations", "conversation"):
val = row.get(key)
if isinstance(val, list) and len(val) >= 2:
out = []
for i, m in enumerate(val):
content = message_content(m)
role = role_of(m, "user" if i % 2 == 0 else "assistant")
if content:
out.append({"role": role, "content": content})
if any(m["role"] == "user" for m in out) and any(m["role"] == "assistant" for m in out):
return out
prompt = clean_text(row.get("instruction") or row.get("prompt") or row.get("question") or row.get("problem") or row.get("query") or row.get("input"))
if clean_text(row.get("instruction")) and clean_text(row.get("input")):
prompt = clean_text(row.get("instruction")) + "\n\n" + clean_text(row.get("input"))
answer = clean_text(row.get("response") or row.get("output") or row.get("answer") or row.get("solution") or row.get("completion") or row.get("target") or row.get("label"))
if not answer and isinstance(row.get("choices"), list) and row.get("answer") is not None:
try:
answer = str(row["choices"][int(row["answer"])])
except Exception:
answer = clean_text(row.get("answer"))
if prompt and answer:
return [{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}]
return None
def first_user_assistant(messages):
system_parts = []
user = None
assistant = None
for m in messages:
if m.get("role") == "system" and m.get("content"):
system_parts.append(m["content"])
elif m.get("role") == "user" and user is None:
user = m.get("content", "")
elif m.get("role") == "assistant" and user is not None:
assistant = m.get("content", "")
break
if not user or not assistant:
return None, None
if system_parts:
user = "System context:\n" + "\n\n".join(system_parts) + "\n\nUser request:\n" + user
return clean_text(user), clean_text(assistant)
def infer_local_group(row):
task = row.get("task_type")
sid = str(row.get("source_id") or row.get("source") or "")
if task == "chinese":
return "chinese_exam" if ":exam" in sid or "ceval" in sid.lower() or "cmmlu" in sid.lower() else "chinese_dialogue"
if task == "code":
return "code"
if task == "math":
return "math"
if task == "dialogue":
return "english_dialogue"
return None
def valid_item(group, user, assistant, tok, args, preset_lengths=None):
if preset_lengths and args.trust_metadata_lengths:
ptok, atok = preset_lengths
else:
ptok = token_len(tok, user)
atok = token_len(tok, assistant)
if ptok < args.min_prompt_tokens or atok < args.min_answer_tokens:
return False, "too_short", ptok, atok
if ptok > args.max_prompt_tokens or atok > args.max_answer_tokens:
return False, "too_long", ptok, atok
if group == "code" and "```" in assistant and assistant.count("```") % 2 != 0:
return False, "bad_code_fence", ptok, atok
return True, "ok", ptok, atok
def make_item(group, source_id, user, assistant, ptok, atok, split, metadata):
return {
"capability": group,
"source": source_id.split(":")[0],
"source_id": source_id,
"split": split,
"messages": [{"role": "user", "content": user}, {"role": "assistant", "content": assistant}],
"hashes": {"pair_sha256": pair_hash(user, assistant), "prompt_sha256": prompt_hash(user)},
"metadata": {"prompt_tokens": ptok, "answer_tokens": atok, **metadata},
}
def add_row(buckets, seen_pair, excluded_pairs, excluded_prompts, group, source_id, split, messages, tok, args, stats, metadata=None, preset_lengths=None):
user, assistant = first_user_assistant(messages)
if not user or not assistant:
stats["no_first_pair"] += 1
return False
ph = prompt_hash(user)
h = pair_hash(user, assistant)
if h in excluded_pairs or ph in excluded_prompts:
stats["heldout_excluded"] += 1
return False
if h in seen_pair:
stats["duplicate"] += 1
return False
ok, reason, ptok, atok = valid_item(group, user, assistant, tok, args, preset_lengths)
if not ok:
stats[reason] += 1
return False
seen_pair.add(h)
buckets[group].append(make_item(group, source_id, user, assistant, ptok, atok, split, metadata or {}))
stats["accepted"] += 1
return True
def read_jsonl(path):
with Path(path).open("r", encoding="utf-8") as f:
for line in f:
if line.strip():
yield json.loads(line)
def load_local_train(buckets, seen_pair, excluded_pairs, excluded_prompts, tok, args):
stats = {}
for path, mode in LOCAL_TRAIN_SOURCES:
path = Path(path)
if not path.exists():
continue
s = Counter()
if mode == "normalized":
for row in read_jsonl(path):
group = infer_local_group(row)
if not group:
s["unknown_group"] += 1
continue
meta = row.get("metadata") or {}
preset = None
if meta.get("prompt_tokens") is not None and meta.get("answer_tokens") is not None:
preset = (int(meta["prompt_tokens"]), int(meta["answer_tokens"]))
add_row(buckets, seen_pair, excluded_pairs, excluded_prompts, group, row.get("source_id") or row.get("source") or str(path), row.get("split", "local"), row.get("messages") or [], tok, args, s, {"loader": "local_normalized", "path": str(path)}, preset)
elif mode.startswith("jsonl:"):
_, group, source_id = mode.split(":", 2)
for row in read_jsonl(path):
msg = normalize_messages(row)
if not msg:
s["no_messages"] += 1
continue
add_row(buckets, seen_pair, excluded_pairs, excluded_prompts, group, source_id, "local", msg, tok, args, s, {"loader": "local_jsonl", "path": str(path)})
elif mode.startswith("csv:"):
_, group, source_id = mode.split(":", 2)
with path.open("r", encoding="utf-8", newline="") as f:
for row in csv.DictReader(f):
msg = normalize_messages(row)
if not msg:
s["no_messages"] += 1
continue
add_row(buckets, seen_pair, excluded_pairs, excluded_prompts, group, source_id, "local", msg, tok, args, s, {"loader": "local_csv", "path": str(path)})
stats[str(path)] = dict(s)
return stats
def iter_configs(src):
configs = src.get("configs")
if not configs:
yield None
elif configs == "all":
from datasets import get_dataset_config_names
try:
for cfg in get_dataset_config_names(src["name"]):
yield cfg
except Exception:
yield None
else:
for cfg in configs:
yield cfg
def load_dataset_iter(src, config, split):
from datasets import load_dataset
if config:
return load_dataset(src["name"], config, split=split, streaming=True)
return load_dataset(src["name"], split=split, streaming=True)
def fill_from_hf(buckets, seen_pair, excluded_pairs, excluded_prompts, tok, quotas, args, split_sources, phase):
stats = {}
rng = random.Random(args.seed + (101 if phase == "test" else 202))
for group in GROUPS:
rng.shuffle(split_sources[group])
for src in split_sources[group]:
if len(buckets[group]) >= quotas[group]:
break
source_base = src["name"]
source_stats = Counter()
print(
f"[{phase}] group={group} source={source_base} "
f"have={len(buckets[group])} target={quotas[group]}",
flush=True,
)
for config in iter_configs(src):
if len(buckets[group]) >= quotas[group]:
break
for split in src.get("splits", ["train"]):
if len(buckets[group]) >= quotas[group]:
break
source_id = f"{source_base}:{config}" if config else source_base
print(f"[{phase}] start group={group} source={source_id} split={split}", flush=True)
try:
ds = load_dataset_iter(src, config, split)
except Exception as exc:
print(f"[{phase}] load_error group={group} source={source_id} split={split} error={exc!r}", flush=True)
source_stats[f"load_error:{split}:{config}"] += 1
source_stats[f"error:{repr(exc)[:160]}"] += 1
continue
skipped = 0
accepted_before = source_stats["accepted"]
try:
for row in ds:
source_stats["seen"] += 1
if skipped < int(src.get("sample_offset", 0)):
skipped += 1
continue
msg = normalize_messages(row)
if not msg:
source_stats["no_messages"] += 1
continue
add_row(buckets, seen_pair, excluded_pairs, excluded_prompts, group, source_id, split, msg, tok, args, source_stats, {"loader": f"hf_{phase}", "row_keys": sorted(row.keys())})
if len(buckets[group]) >= quotas[group]:
break
if source_stats["seen"] >= args.max_seen_per_source:
source_stats["max_seen_stop"] += 1
break
except Exception as exc:
print(f"[{phase}] iter_error group={group} source={source_id} split={split} error={exc!r}", flush=True)
source_stats[f"iter_error:{split}:{config}"] += 1
source_stats[f"error:{repr(exc)[:160]}"] += 1
print(
f"[{phase}] done group={group} source={source_id} split={split} "
f"accepted_delta={source_stats['accepted'] - accepted_before} "
f"accepted_total={source_stats['accepted']} seen={source_stats['seen']} "
f"group_have={len(buckets[group])}",
flush=True,
)
if source_stats["accepted"] == accepted_before and args.stop_empty_source_early:
break
stats[f"{phase}:{group}:{source_base}"] = dict(source_stats)
return stats
def quotas_from_ratios(total, ratios):
denom = sum(ratios.values())
quotas = {g: total * ratios[g] // denom for g in GROUPS}
quotas["chinese_dialogue"] += total - sum(quotas.values())
return quotas
def write_jsonl(path, rows):
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("w", encoding="utf-8") as f:
for row in rows:
f.write(json.dumps(row, ensure_ascii=False) + "\n")
def write_stats(path, stats):
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(stats, ensure_ascii=False, indent=2), encoding="utf-8")
def upsample_bucket(rows, target, rng):
if len(rows) >= target:
return rows[:target], 0
if not rows:
return [], 0
out = list(rows)
added = 0
while len(out) < target:
item = copy.deepcopy(rng.choice(rows))
item.setdefault("metadata", {})["upsampled"] = True
item["metadata"]["upsample_index"] = added
out.append(item)
added += 1
rng.shuffle(out)
return out, added
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--out", required=True)
ap.add_argument("--tokenizer", default="/ssd/yi/Tokenizer_Swap/model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
ap.add_argument("--train-total", type=int, default=1000000)
ap.add_argument("--test-per-group", type=int, default=400)
ap.add_argument("--seed", type=int, default=20260611)
ap.add_argument("--hf-endpoint", default="https://hf-mirror.com")
ap.add_argument("--use-hf", action="store_true")
ap.add_argument("--allow-shortfall", action="store_true")
ap.add_argument("--upsample-train-shortfall", action="store_true")
ap.add_argument("--min-prompt-tokens", type=int, default=4)
ap.add_argument("--min-answer-tokens", type=int, default=8)
ap.add_argument("--max-prompt-tokens", type=int, default=1024)
ap.add_argument("--max-answer-tokens", type=int, default=1024)
ap.add_argument("--trust-metadata-lengths", action="store_true")
ap.add_argument("--max-seen-per-source", type=int, default=300000)
ap.add_argument("--stop-empty-source-early", action="store_true")
args = ap.parse_args()
os.environ.setdefault("HF_ENDPOINT", args.hf_endpoint)
out = Path(args.out)
tok = load_tokenizer(args.tokenizer)
rng = random.Random(args.seed)
test_quotas = {g: args.test_per_group for g in GROUPS}
test_buckets = defaultdict(list)
test_seen = set()
empty_excluded = set()
stats = {"args": vars(args), "train_ratios": TRAIN_RATIOS, "test_quotas": test_quotas, "loaders": {}}
if args.use_hf:
stats["loaders"]["test_hf"] = fill_from_hf(test_buckets, test_seen, set(), set(), tok, test_quotas, args, HF_TEST_SOURCES, "test")
test_rows = []
test_shortfall = {}
for g in GROUPS:
rng.shuffle(test_buckets[g])
take = min(test_quotas[g], len(test_buckets[g]))
test_rows.extend(test_buckets[g][:take])
if take < test_quotas[g]:
test_shortfall[g] = test_quotas[g] - take
test_pair_hashes = {r["hashes"]["pair_sha256"] for r in test_rows}
test_prompt_hashes = {r["hashes"]["prompt_sha256"] for r in test_rows}
train_quotas = quotas_from_ratios(args.train_total, TRAIN_RATIOS)
train_buckets = defaultdict(list)
train_seen = set()
stats["loaders"]["train_local"] = load_local_train(train_buckets, train_seen, test_pair_hashes, test_prompt_hashes, tok, args)
if args.use_hf:
stats["loaders"]["train_hf"] = fill_from_hf(train_buckets, train_seen, test_pair_hashes, test_prompt_hashes, tok, train_quotas, args, HF_TRAIN_SOURCES, "train")
train_rows = []
train_shortfall = {}
train_upsampled = {}
for g in GROUPS:
rng.shuffle(train_buckets[g])
if args.upsample_train_shortfall:
selected, added = upsample_bucket(train_buckets[g], train_quotas[g], rng)
take = len(selected)
train_rows.extend(selected)
if added:
train_upsampled[g] = added
else:
take = min(train_quotas[g], len(train_buckets[g]))
train_rows.extend(train_buckets[g][:take])
if take < train_quotas[g]:
train_shortfall[g] = train_quotas[g] - take
if (test_shortfall or train_shortfall) and not args.allow_shortfall:
stats["result"] = {
"failed": True,
"test_shortfall": test_shortfall,
"train_shortfall": train_shortfall,
"train_upsampled": train_upsampled,
"test_available": {g: len(test_buckets[g]) for g in GROUPS},
"train_available": {g: len(train_buckets[g]) for g in GROUPS},
"train_quotas": train_quotas,
}
write_stats(out / "build_stats.json", stats)
raise SystemExit(json.dumps(stats["result"], ensure_ascii=False, indent=2))
rng.shuffle(test_rows)
rng.shuffle(train_rows)
write_jsonl(out / "heldout_2p8k.jsonl", test_rows)
write_jsonl(out / "train_1m.jsonl", train_rows)
write_jsonl(out / "heldout_exclusion_hashes.jsonl", [{"pair_sha256": r["hashes"]["pair_sha256"], "prompt_sha256": r["hashes"]["prompt_sha256"], "capability": r["capability"], "source_id": r["source_id"]} for r in test_rows])
stats["result"] = {
"failed": False,
"test_total": len(test_rows),
"train_total": len(train_rows),
"test_counts": dict(Counter(r["capability"] for r in test_rows)),
"train_counts": dict(Counter(r["capability"] for r in train_rows)),
"test_shortfall": test_shortfall,
"train_shortfall": train_shortfall,
"train_upsampled": train_upsampled,
"train_quotas": train_quotas,
"train_source_top": dict(Counter(r["source_id"] for r in train_rows).most_common(80)),
"test_source_top": dict(Counter(r["source_id"] for r in test_rows).most_common(80)),
}
write_stats(out / "build_stats.json", stats)
print(json.dumps(stats["result"], ensure_ascii=False, indent=2))
print(out)
if __name__ == "__main__":
main()

2000
dataset_building/heldout_public_mcq_2k_20260607/heldout_public_mcq_2k.jsonl Normal file
View File

File diff suppressed because one or more lines are too long

17
dataset_building/heldout_public_mcq_2k_20260607/heldout_public_mcq_2k_stats.json Normal file
View File

@@ -0,0 +1,17 @@
{
"seed": 20260607,
"requested": {
"mmlu": 1000,
"gpqa": 400,
"ceval": 600,
"cmmlu": 0
},
"total_items": 2000,
"counts": {
"gpqa": 400,
"ceval": 600,
"mmlu": 1000
},
"errors": {},
"output": "data/heldout_public_mcq_2k_20260607/heldout_public_mcq_2k.jsonl"
}

118
dataset_building/metadata/cpt_docmix_5b_manifest.json Normal file
View File

@@ -0,0 +1,118 @@
{
"tokenizer": "/ssd/yi/tokenizer_swap_cepe/models/Qwen3-0.6B-DSV4-tokenizer-remap-v2",
"budgets": {
"english_web": 1250000000,
"english_edu": 1000000000,
"chinese_clean": 1250000000,
"code": 750000000,
"math": 500000000,
"science": 150000000,
"qa_as_text": 100000000
},
"seq_len_for_later_packing": 8192,
"stream_sources": {
"english_web": [
{
"kind": "hf",
"name": "HuggingFaceFW/fineweb",
"config": "CC-MAIN-2025-26",
"split": "train",
"max_rows": 0
},
{
"kind": "hf",
"name": "HuggingFaceFW/fineweb",
"config": "CC-MAIN-2025-21",
"split": "train",
"max_rows": 0
}
],
"english_edu": [
{
"kind": "hf",
"name": "HuggingFaceFW/fineweb-edu",
"config": null,
"split": "train",
"max_rows": 0
}
],
"chinese_clean": [
{
"kind": "hf",
"name": "BAAI/CCI3-HQ",
"config": null,
"split": "train",
"max_rows": 0
},
{
"kind": "hf",
"name": "Skywork/SkyPile-150B",
"config": null,
"split": "train",
"max_rows": 0
}
],
"code": [
{
"kind": "hf",
"name": "bigcode/starcoderdata",
"config": null,
"split": "train",
"max_rows": 0
},
{
"kind": "hf",
"name": "codeparrot/github-code",
"config": null,
"split": "train",
"max_rows": 0
}
],
"math": [
{
"kind": "hf",
"name": "open-web-math/open-web-math",
"config": null,
"split": "train",
"max_rows": 0
},
{
"kind": "hf",
"name": "GAIR/MathPile",
"config": null,
"split": "train",
"max_rows": 0
}
]
},
"local_globs": {
"science": [
"data/offline_text_only_reasoning_sources_20260611/science_reasoning__*.jsonl",
"data/offline_text_only_reasoning_sources_20260611/logic__*.jsonl"
],
"qa_as_text": [
"data/training_mix_v4_train1m_test2p8k_noupsample_nobbh_20260611/train_1m.jsonl",
"data/open_recovery_sft_mix_alt_sources_1m_parquet_20260607/normalized.jsonl"
]
},
"local_parquet_globs": {
"english_web": [
"data/raw_parquets/fineweb_2025/*.parquet"
],
"english_edu": [
"data/raw_parquets/fineweb_edu/*.parquet"
],
"code": [
"data/raw_parquets/starcoder/python/*.parquet",
"data/raw_parquets/starcoder/javascript/*.parquet",
"data/raw_parquets/starcoder/typescript/*.parquet",
"data/raw_parquets/starcoder/java/*.parquet",
"data/raw_parquets/starcoder/cpp/*.parquet",
"data/raw_parquets/starcoder/go/*.parquet",
"data/raw_parquets/starcoder/rust/*.parquet",
"data/raw_parquets/starcoder/shell/*.parquet"
]
},
"min_tokens": 128,
"max_doc_tokens": 32768
}

215
dataset_building/metadata/cpt_docmix_5b_stats.json Normal file
View File

@@ -0,0 +1,215 @@
{
"scale": 5.0,
"budgets": {
"english_web": 1250000000,
"english_edu": 1000000000,
"chinese_clean": 1250000000,
"code": 750000000,
"math": 500000000,
"science": 150000000,
"qa_as_text": 100000000
},
"tokens_by_category": {
"english_web": 1250000445,
"english_edu": 1000000462,
"code": 750001624,
"math": 96992014,
"science": 30016910,
"qa_as_text": 100000055
},
"docs_by_category": {
"english_web": 1571857,
"english_edu": 973081,
"code": 589463,
"math": 51576,
"science": 48872,
"qa_as_text": 244246
},
"tokens_by_source": {
"data/raw_parquets/fineweb_2025/CC-MAIN-2025-26_000_00000.parquet": 684777289,
"data/raw_parquets/fineweb_edu/CC-MAIN-2024-10_000_00000.parquet": 627654914,
"data/raw_parquets/fineweb_2025/CC-MAIN-2025-26_001_00000.parquet": 565223156,
"data/raw_parquets/fineweb_edu/CC-MAIN-2024-10_000_00001.parquet": 372345548,
"data/raw_parquets/starcoder/python/train-00000-of-00059.parquet": 238778818,
"data/raw_parquets/starcoder/python/train-00002-of-00059.parquet": 233459009,
"data/raw_parquets/starcoder/python/train-00001-of-00059.parquet": 233210510,
"/ssd/yi/tokenizer_swap_cepe/data/training_mix_v4_train1m_test2p8k_noupsample_nobbh_20260611/train_1m.jsonl": 100000055,
"open-web-math/open-web-math": 96992014,
"data/raw_parquets/starcoder/python/train-00003-of-00059.parquet": 44553287,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__openlifescienceai_medmcqa__train.jsonl": 10027949,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_proofwriter__train.jsonl": 6138410,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__derek-thomas_ScienceQA__train.jsonl": 3119670,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_proofwriter__test.jsonl": 1837710,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__GBaker_MedQA-USMLE-4-options__train.jsonl": 1572096,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__sciq__train.jsonl": 1182676,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__derek-thomas_ScienceQA__test.jsonl": 1033827,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__derek-thomas_ScienceQA__validation.jsonl": 1032383,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_proofwriter__validation.jsonl": 895384,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__allenai_qasc__train.jsonl": 688074,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__openlifescienceai_medmcqa__test.jsonl": 307416,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_folio__train.jsonl": 261306,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__openlifescienceai_medmcqa__validation.jsonl": 242852,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__allenai_openbookqa__main__train.jsonl": 221520,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__GBaker_MedQA-USMLE-4-options__test.jsonl": 203198,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_bigbench__temporal_sequences__train.jsonl": 141867,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__allenai_ai2_arc__ARC-Easy__test.jsonl": 139093,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__allenai_ai2_arc__ARC-Easy__train.jsonl": 130752,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_bigbench__logical_deduction__train.jsonl": 118299,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__sciq__test.jsonl": 102393,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__sciq__validation.jsonl": 99297,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__allenai_ai2_arc__ARC-Challenge__test.jsonl": 79458,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__allenai_qasc__validation.jsonl": 79086,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__allenai_ai2_arc__ARC-Challenge__train.jsonl": 74207,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_folio__validation.jsonl": 53607,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_bigbench__temporal_sequences__validation.jsonl": 35418,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__allenai_ai2_arc__ARC-Easy__validation.jsonl": 33149,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_bigbench__logical_deduction__validation.jsonl": 31008,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__allenai_openbookqa__main__validation.jsonl": 23331,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__allenai_openbookqa__main__test.jsonl": 22745,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__qiaojin_PubMedQA__pqa_labeled__train.jsonl": 22169,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/science_reasoning__allenai_ai2_arc__ARC-Challenge__validation.jsonl": 19963,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_bigbench__disambiguation_qa__train.jsonl": 15771,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_bigbench__date_understanding__train.jsonl": 14547,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_bigbench__disambiguation_qa__validation.jsonl": 3843,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_bigbench__date_understanding__validation.jsonl": 3629,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_bigbench__logical_args__validation.jsonl": 3136,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_bigbench__logical_args__train.jsonl": 2729,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_bigbench__analytic_entailment__train.jsonl": 2277,
"/ssd/yi/tokenizer_swap_cepe/data/offline_text_only_reasoning_sources_20260611/logic__tasksource_bigbench__analytic_entailment__validation.jsonl": 665
},
"rejected": {
"english_web": {
"too_short": 168039,
"too_long": 179
},
"english_edu": {
"too_short": 18375,
"too_long": 221
},
"code": {
"too_short": 106203,
"too_long": 816
},
"math": {
"too_short": 3669,
"too_long": 152
},
"science": {
"pack_too_short": 5591
},
"qa_as_text": {
"too_short": 108331
}
},
"sources": [
{
"category": "english_web",
"kind": "local_parquet",
"source": "data/raw_parquets/fineweb_2025/CC-MAIN-2025-26_000_00000.parquet",
"rows_seen": 961000,
"docs_written": 864351,
"tokens": 684777289,
"error": ""
},
{
"category": "english_web",
"kind": "local_parquet",
"source": "data/raw_parquets/fineweb_2025/CC-MAIN-2025-26_001_00000.parquet",
"rows_seen": 779076,
"docs_written": 707506,
"tokens": 565223156,
"error": ""
},
{
"category": "english_edu",
"kind": "local_parquet",
"source": "data/raw_parquets/fineweb_edu/CC-MAIN-2024-10_000_00000.parquet",
"rows_seen": 621759,
"docs_written": 610163,
"tokens": 627654914,
"error": ""
},
{
"category": "english_edu",
"kind": "local_parquet",
"source": "data/raw_parquets/fineweb_edu/CC-MAIN-2024-10_000_00001.parquet",
"rows_seen": 369919,
"docs_written": 362918,
"tokens": 372345548,
"error": ""
},
{
"category": "chinese_clean",
"kind": "hf_stream",
"source": "BAAI/CCI3-HQ",
"rows_seen": 0,
"docs_written": 0,
"tokens": 0,
"error": "DatasetNotFoundError(\"Dataset 'BAAI/CCI3-HQ' is a gated dataset on the Hub. You must be authenticated to access it.\")"
},
{
"category": "chinese_clean",
"kind": "hf_stream",
"source": "Skywork/SkyPile-150B",
"rows_seen": 0,
"docs_written": 0,
"tokens": 0,
"error": "RuntimeError('Cannot send a request, as the client has been closed.')"
},
{
"category": "code",
"kind": "local_parquet",
"source": "data/raw_parquets/starcoder/python/train-00000-of-00059.parquet",
"rows_seen": 218079,
"docs_written": 184958,
"tokens": 238778818,
"error": ""
},
{
"category": "code",
"kind": "local_parquet",
"source": "data/raw_parquets/starcoder/python/train-00001-of-00059.parquet",
"rows_seen": 218079,
"docs_written": 184495,
"tokens": 233210510,
"error": ""
},
{
"category": "code",
"kind": "local_parquet",
"source": "data/raw_parquets/starcoder/python/train-00002-of-00059.parquet",
"rows_seen": 218079,
"docs_written": 184331,
"tokens": 233459009,
"error": ""
},
{
"category": "code",
"kind": "local_parquet",
"source": "data/raw_parquets/starcoder/python/train-00003-of-00059.parquet",
"rows_seen": 42246,
"docs_written": 35679,
"tokens": 44553287,
"error": ""
},
{
"category": "math",
"kind": "hf_stream",
"source": "open-web-math/open-web-math",
"rows_seen": 55397,
"docs_written": 51576,
"tokens": 96992014,
"error": "ProxyError('503 Service Unavailable')"
},
{
"category": "math",
"kind": "hf_stream",
"source": "GAIR/MathPile",
"rows_seen": 0,
"docs_written": 0,
"tokens": 0,
"error": "DatasetNotFoundError(\"Dataset 'GAIR/MathPile' is a gated dataset on the Hub. You must be authenticated to access it.\")"
}
],
"elapsed_sec": 10098.068085193634
}

393
dataset_building/metadata/cpt_packed_1b_seed42_stratified_manifest.json Normal file
View File

@@ -0,0 +1,393 @@
{
"docs_seen": 960549,
"train_blocks": 121041,
"eval_blocks": 1024,
"train_tokens": 991567872,
"eval_tokens": 8388608,
"source_docs": {
"english_web": 318271,
"english_edu": 194585,
"chinese_clean": 181655,
"code": 116242,
"math": 53125,
"science": 47973,
"qa_as_text": 48698
},
"source_tokens": {
"english_web": 249999374,
"english_edu": 199998420,
"chinese_clean": 249999618,
"code": 149999958,
"math": 99987860,
"science": 29999428,
"qa_as_text": 19999497
},
"train_blocks_by_category": {
"english_web": 30261,
"english_edu": 24208,
"chinese_clean": 30261,
"code": 18156,
"math": 12103,
"science": 3631,
"qa_as_text": 2421
},
"eval_blocks_by_category": {
"english_web": 256,
"english_edu": 205,
"chinese_clean": 256,
"code": 154,
"math": 102,
"science": 31,
"qa_as_text": 20
},
"start_time": 1781429820.1093097,
"leftover_tokens_by_category": {
"english_web": 4110,
"english_edu": 7124,
"chinese_clean": 4354,
"code": 4438,
"math": 4500,
"science": 324,
"qa_as_text": 2825
},
"tokenizer": "/ssd/yi/tokenizer_swap_cepe/models/Qwen3-0.6B-DSV4-tokenizer-remap-v2",
"seq_len": 8192,
"seed": 42,
"budgets": {
"english_web": 250000000,
"english_edu": 200000000,
"chinese_clean": 250000000,
"code": 150000000,
"math": 100000000,
"science": 30000000,
"qa_as_text": 20000000
},
"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"
},
"eval_quota_blocks": {
"english_web": 256,
"english_edu": 205,
"chinese_clean": 256,
"code": 154,
"math": 102,
"science": 31,
"qa_as_text": 20
},
"train_shards": [
{
"path": "train_00000.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00001.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00002.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00003.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00004.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00005.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00006.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00007.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00008.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00009.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00010.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00011.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00012.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00013.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00014.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00015.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00016.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00017.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00018.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00019.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00020.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00021.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00022.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00023.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00024.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00025.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00026.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00027.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00028.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00029.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00030.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00031.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00032.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00033.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00034.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00035.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00036.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00037.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00038.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00039.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00040.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00041.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00042.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00043.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00044.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00045.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00046.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00047.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00048.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00049.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00050.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00051.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00052.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00053.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00054.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00055.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00056.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00057.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00058.npy",
"blocks": 2048,
"tokens": 16777216
},
{
"path": "train_00059.npy",
"blocks": 209,
"tokens": 1712128
}
],
"eval_shards": [
{
"path": "eval_00000.npy",
"blocks": 1024,
"tokens": 8388608
}
],
"elapsed_sec": 2387.9453027248383
}

1584
dataset_building/metadata/cpt_packed_5b_seed42_stratified_manifest.json Normal file
View File

File diff suppressed because it is too large Load Diff

309
dataset_building/metadata/sft_v4_mix_build_stats.json Normal file
View File

@@ -0,0 +1,309 @@
{
"args": {
"out": "data/training_mix_v4_train1m_test2p8k_noupsample_nobbh_20260611",
"tokenizer": "/ssd/yi/tokenizer_swap_cepe/models/Qwen3-0.6B-DSV4-tokenizer-remap-v2",
"train_total": 1000000,
"test_per_group": 400,
"seed": 20260611,
"hf_endpoint": "https://hf-mirror.com",
"use_hf": true,
"allow_shortfall": true,
"upsample_train_shortfall": false,
"min_prompt_tokens": 4,
"min_answer_tokens": 8,
"max_prompt_tokens": 1024,
"max_answer_tokens": 1024,
"trust_metadata_lengths": true,
"max_seen_per_source": 500000,
"stop_empty_source_early": false
},
"train_ratios": {
"chinese_exam": 0,
"chinese_dialogue": 35,
"code": 30,
"math": 20,
"logic": 5,
"science_reasoning": 5,
"english_dialogue": 5
},
"test_quotas": {
"chinese_exam": 400,
"chinese_dialogue": 400,
"code": 400,
"math": 400,
"logic": 400,
"science_reasoning": 400,
"english_dialogue": 400
},
"loaders": {
"test_hf": {
"test:chinese_exam:ceval/ceval-exam": {
"seen": 728,
"accepted": 400,
"too_short": 328
},
"test:chinese_dialogue:m-a-p/COIG-CQIA": {
"seen": 414,
"accepted": 400,
"too_long": 14
},
"test:code:openai/openai_humaneval": {
"seen": 164,
"accepted": 154,
"too_short": 8,
"too_long": 2
},
"test:code:google-research-datasets/mbpp": {
"seen": 246,
"accepted": 246
},
"test:math:gsm8k": {
"seen": 400,
"accepted": 400
},
"test:logic:cais/mmlu": {
"seen": 331,
"accepted": 164,
"too_short": 167
},
"test:logic:lighteval/bbh": {
"load_error:test:logical_deduction_three_objects": 1,
"error:ValueError(\"Bad split: test. Available splits: ['train']\")": 1,
"seen": 296,
"accepted": 236,
"too_short": 60
},
"test:science_reasoning:Idavidrein/gpqa": {
"seen": 448,
"too_short": 195,
"accepted": 252,
"too_long": 1
},
"test:science_reasoning:allenai/ai2_arc": {
"seen": 287,
"accepted": 148,
"too_short": 139
},
"test:english_dialogue:HuggingFaceH4/ultrachat_200k": {
"seen": 420,
"too_long": 20,
"accepted": 400
}
},
"train_local": {
"data/open_recovery_sft_mix_alt_sources_1m_parquet_20260607/normalized.jsonl": {
"accepted": 1051916,
"heldout_excluded": 411,
"duplicate": 30
},
"data/open_recovery_sft_mix_100k_4_3_2_1_20260604/normalized.jsonl": {
"duplicate": 88687,
"accepted": 9020,
"heldout_excluded": 401
},
"data/modelscope_alt_sources_20260607/CodeAlpaca-20k.jsonl": {
"duplicate": 18820,
"too_short": 1193,
"no_messages": 6,
"too_long": 3
},
"data/modelscope_alt_sources_20260607/alpaca-gpt4-data-zh_train.csv": {
"duplicate": 46476,
"too_short": 2342
}
},
"train_hf": {
"train:chinese_dialogue:m-a-p/COIG-CQIA": {
"seen": 10116,
"duplicate": 7211,
"too_short": 1138,
"too_long": 1354,
"no_messages": 1,
"accepted": 1,
"heldout_excluded": 411
},
"train:chinese_dialogue:BelleGroup/train_0.5M_CN": {
"seen": 330837,
"too_short": 34263,
"duplicate": 2188,
"accepted": 294334,
"too_long": 32,
"no_messages": 20
},
"train:code:ise-uiuc/Magicoder-OSS-Instruct-75K": {
"seen": 75197,
"duplicate": 69998,
"bad_code_fence": 34,
"too_long": 48,
"accepted": 5117
},
"train:code:bigcode/self-oss-instruct-sc2-exec-filter-50k": {
"seen": 50661,
"duplicate": 31578,
"accepted": 19076,
"too_long": 6,
"bad_code_fence": 1
},
"train:code:HuggingFaceTB/smoltalk": {
"seen": 131271,
"duplicate": 78743,
"too_long": 1498,
"bad_code_fence": 1,
"accepted": 51029
},
"train:code:nvidia/OpenCodeInstruct": {
"seen": 43455,
"duplicate": 6678,
"too_long": 39,
"accepted": 36738
},
"train:logic:tasksource/bigbench": {
"seen": 8291,
"too_short": 5767,
"accepted": 2338,
"heldout_excluded": 186
},
"train:logic:tau/commonsense_qa": {
"seen": 10962,
"too_short": 10961,
"accepted": 1
},
"train:logic:metaeval/reclor": {
"seen": 5138,
"accepted": 5117,
"too_short": 21
},
"train:science_reasoning:allenai/ai2_arc": {
"seen": 1418,
"too_short": 739,
"accepted": 679
},
"train:science_reasoning:allenai/qasc": {
"iter_error:train:None": 1,
"error:ProxyError('503 Service Unavailable')": 2,
"iter_error:validation:None": 1
},
"train:science_reasoning:allenai/openbookqa": {
"iter_error:train:main": 1,
"error:ProxyError('503 Service Unavailable')": 4,
"iter_error:validation:main": 1,
"iter_error:train:additional": 1,
"iter_error:validation:additional": 1
},
"train:science_reasoning:sciq": {
"seen": 12679,
"accepted": 11370,
"too_short": 1309
},
"train:science_reasoning:qiaojin/PubMedQA": {
"seen": 1000,
"accepted": 1000
}
}
},
"result": {
"failed": false,
"test_total": 2800,
"train_total": 921360,
"test_counts": {
"science_reasoning": 400,
"logic": 400,
"code": 400,
"chinese_exam": 400,
"math": 400,
"chinese_dialogue": 400,
"english_dialogue": 400
},
"train_counts": {
"chinese_dialogue": 350000,
"code": 300000,
"english_dialogue": 50000,
"math": 200000,
"science_reasoning": 13904,
"logic": 7456
},
"test_shortfall": {},
"train_shortfall": {
"logic": 42544,
"science_reasoning": 36096
},
"train_upsampled": {},
"train_quotas": {
"chinese_exam": 0,
"chinese_dialogue": 350000,
"code": 300000,
"math": 200000,
"logic": 50000,
"science_reasoning": 50000,
"english_dialogue": 50000
},
"train_source_top": {
"BelleGroup/train_0.5M_CN": 296334,
"nvidia/OpenMathInstruct-2": 89260,
"TIGER-Lab/MathInstruct": 85173,
"ise-uiuc/Magicoder-OSS-Instruct-75K": 75115,
"bigcode/self-oss-instruct-sc2-exec-filter-50k": 66632,
"HuggingFaceTB/smoltalk:apigen-80k": 51029,
"AI-ModelScope/alpaca-gpt4-data-zh:file": 46476,
"nvidia/OpenCodeInstruct": 43404,
"AI-ModelScope/smoltalk:self-oss-instruct-file": 25000,
"AI-ModelScope/smoltalk:apigen-80k-file": 20000,
"AI-ModelScope/CodeAlpaca-20k:file": 18820,
"AI-ModelScope/smoltalk:numina-cot-100k-file": 17080,
"HuggingFaceTB/smoltalk:all": 15260,
"HuggingFaceH4/ultrachat_200k": 15238,
"allenai/tulu-3-sft-mixture": 15198,
"sciq": 11370,
"AI-ModelScope/smoltalk:metamathqa-50k-file": 8487,
"metaeval/reclor": 5117,
"m-a-p/COIG-CQIA:zhihu": 4978,
"AI-ModelScope/smoltalk:openhermes-100k": 4304,
"m-a-p/COIG-CQIA:coig_pc": 1877,
"qiaojin/PubMedQA:pqa_labeled": 1000,
"allenai/ai2_arc:ARC-Easy": 855,
"tasksource/bigbench:logical_deduction": 812,
"tasksource/bigbench:temporal_sequences": 800,
"allenai/ai2_arc:ARC-Challenge": 679,
"tasksource/bigbench:social_iqa": 343,
"m-a-p/COIG-CQIA:wikihow": 335,
"tasksource/bigbench:date_understanding": 296,
"tasksource/bigbench:disambiguation_qa": 71,
"tasksource/bigbench:logical_args": 16,
"tau/commonsense_qa": 1
},
"test_source_top": {
"gsm8k:main": 400,
"m-a-p/COIG-CQIA:zhihu": 400,
"HuggingFaceH4/ultrachat_200k": 400,
"Idavidrein/gpqa": 252,
"google-research-datasets/mbpp": 246,
"lighteval/bbh:logical_deduction_three_objects": 236,
"openai/openai_humaneval": 154,
"allenai/ai2_arc:ARC-Challenge": 148,
"cais/mmlu:formal_logic": 98,
"cais/mmlu:logical_fallacies": 66,
"ceval/ceval-exam:civil_servant": 38,
"ceval/ceval-exam:college_economics": 36,
"ceval/ceval-exam:accountant": 31,
"ceval/ceval-exam:college_programming": 23,
"ceval/ceval-exam:fire_engineer": 22,
"ceval/ceval-exam:advanced_mathematics": 22,
"ceval/ceval-exam:high_school_chinese": 21,
"ceval/ceval-exam:high_school_biology": 19,
"ceval/ceval-exam:college_chemistry": 18,
"ceval/ceval-exam:college_physics": 18,
"ceval/ceval-exam:high_school_chemistry": 17,
"ceval/ceval-exam:environmental_impact_assessment_engineer": 16,
"ceval/ceval-exam:business_administration": 15,
"ceval/ceval-exam:discrete_mathematics": 14,
"ceval/ceval-exam:education_science": 13,
"ceval/ceval-exam:clinical_medicine": 13,
"ceval/ceval-exam:electrical_engineer": 12,
"ceval/ceval-exam:computer_network": 11,
"ceval/ceval-exam:art_studies": 11,
"ceval/ceval-exam:computer_architecture": 11,
"ceval/ceval-exam:chinese_language_and_literature": 10,
"ceval/ceval-exam:basic_medicine": 9
}
}
}

164
dataset_building/metadata/sft_v4_tokenization_build_stats.json Normal file
View File

@@ -0,0 +1,164 @@
{
"tokenizer": "/ssd/yi/tokenizer_swap_cepe/models/dsv4_chat_full_sft_remap_v2_alt1m_5epoch_bsz8_accum16_20260610",
"encoding_dir": "/ssd/yi/tokenizer_swap_cepe/external/deepseek_v4_encoding",
"cutoff_len": 2048,
"eos_token": "<｜end▁of▁sentence｜>",
"eos_token_id": 1,
"splits": {
"train": {
"split": "train",
"source": "/ssd/yi/tokenizer_swap_cepe/data/training_mix_v4_train1m_test2p8k_noupsample_nobbh_20260611/train_1m.jsonl",
"output": "/ssd/yi/tokenizer_swap_cepe/data/dsv4_chat_tokenized_v4_noupsample_nobbh_921k_20260611/train_dsv4_chat_tokenized.jsonl.gz",
"cutoff_len": 2048,
"rows_seen": 921360,
"rows_written": 921360,
"skipped_no_messages": 0,
"truncated": 0,
"eos_in_labels": 921360,
"prefix_mismatch": 0,
"capability_counts": {
"chinese_dialogue": 350000,
"code": 300000,
"math": 200000,
"english_dialogue": 50000,
"science_reasoning": 13904,
"logic": 7456
},
"prompt_tokens": {
"p50": 65,
"p90": 461,
"p95": 698,
"p99": 951,
"max": 1028
},
"response_tokens": {
"p50": 101,
"p90": 340,
"p95": 457,
"p99": 748,
"max": 1025
},
"total_tokens": {
"p50": 223,
"p90": 699,
"p95": 889,
"p99": 1165,
"max": 1957
},
"truncated_rate": 0.0,
"eos_label_rate": 1.0,
"source_counts_top50": {
"BelleGroup/train_0.5M_CN": 296334,
"nvidia/OpenMathInstruct-2": 89260,
"TIGER-Lab/MathInstruct": 85173,
"ise-uiuc/Magicoder-OSS-Instruct-75K": 75115,
"bigcode/self-oss-instruct-sc2-exec-filter-50k": 66632,
"HuggingFaceTB/smoltalk:apigen-80k": 51029,
"AI-ModelScope/alpaca-gpt4-data-zh:file": 46476,
"nvidia/OpenCodeInstruct": 43404,
"AI-ModelScope/smoltalk:self-oss-instruct-file": 25000,
"AI-ModelScope/smoltalk:apigen-80k-file": 20000,
"AI-ModelScope/CodeAlpaca-20k:file": 18820,
"AI-ModelScope/smoltalk:numina-cot-100k-file": 17080,
"HuggingFaceTB/smoltalk:all": 15260,
"HuggingFaceH4/ultrachat_200k": 15238,
"allenai/tulu-3-sft-mixture": 15198,
"sciq": 11370,
"AI-ModelScope/smoltalk:metamathqa-50k-file": 8487,
"metaeval/reclor": 5117,
"m-a-p/COIG-CQIA:zhihu": 4978,
"AI-ModelScope/smoltalk:openhermes-100k": 4304,
"m-a-p/COIG-CQIA:coig_pc": 1877,
"qiaojin/PubMedQA:pqa_labeled": 1000,
"allenai/ai2_arc:ARC-Easy": 855,
"tasksource/bigbench:logical_deduction": 812,
"tasksource/bigbench:temporal_sequences": 800,
"allenai/ai2_arc:ARC-Challenge": 679,
"tasksource/bigbench:social_iqa": 343,
"m-a-p/COIG-CQIA:wikihow": 335,
"tasksource/bigbench:date_understanding": 296,
"tasksource/bigbench:disambiguation_qa": 71,
"tasksource/bigbench:logical_args": 16,
"tau/commonsense_qa": 1
}
},
"validation": {
"split": "validation",
"source": "/ssd/yi/tokenizer_swap_cepe/data/training_mix_v4_train1m_test2p8k_noupsample_nobbh_20260611/heldout_2p8k.jsonl",
"output": "/ssd/yi/tokenizer_swap_cepe/data/dsv4_chat_tokenized_v4_noupsample_nobbh_921k_20260611/validation_dsv4_chat_tokenized.jsonl.gz",
"cutoff_len": 2048,
"rows_seen": 2800,
"rows_written": 2800,
"skipped_no_messages": 0,
"truncated": 0,
"eos_in_labels": 2800,
"prefix_mismatch": 0,
"capability_counts": {
"science_reasoning": 400,
"logic": 400,
"code": 400,
"chinese_exam": 400,
"math": 400,
"chinese_dialogue": 400,
"english_dialogue": 400
},
"prompt_tokens": {
"p50": 85,
"p90": 238,
"p95": 346,
"p99": 687,
"max": 1020
},
"response_tokens": {
"p50": 53,
"p90": 416,
"p95": 547,
"p99": 837,
"max": 1019
},
"total_tokens": {
"p50": 177,
"p90": 543,
"p95": 688,
"p99": 981,
"max": 1162
},
"truncated_rate": 0.0,
"eos_label_rate": 1.0,
"source_counts_top50": {
"gsm8k:main": 400,
"m-a-p/COIG-CQIA:zhihu": 400,
"HuggingFaceH4/ultrachat_200k": 400,
"Idavidrein/gpqa": 252,
"google-research-datasets/mbpp": 246,
"lighteval/bbh:logical_deduction_three_objects": 236,
"openai/openai_humaneval": 154,
"allenai/ai2_arc:ARC-Challenge": 148,
"cais/mmlu:formal_logic": 98,
"cais/mmlu:logical_fallacies": 66,
"ceval/ceval-exam:civil_servant": 38,
"ceval/ceval-exam:college_economics": 36,
"ceval/ceval-exam:accountant": 31,
"ceval/ceval-exam:college_programming": 23,
"ceval/ceval-exam:fire_engineer": 22,
"ceval/ceval-exam:advanced_mathematics": 22,
"ceval/ceval-exam:high_school_chinese": 21,
"ceval/ceval-exam:high_school_biology": 19,
"ceval/ceval-exam:college_chemistry": 18,
"ceval/ceval-exam:college_physics": 18,
"ceval/ceval-exam:high_school_chemistry": 17,
"ceval/ceval-exam:environmental_impact_assessment_engineer": 16,
"ceval/ceval-exam:business_administration": 15,
"ceval/ceval-exam:discrete_mathematics": 14,
"ceval/ceval-exam:education_science": 13,
"ceval/ceval-exam:clinical_medicine": 13,
"ceval/ceval-exam:electrical_engineer": 12,
"ceval/ceval-exam:computer_network": 11,
"ceval/ceval-exam:art_studies": 11,
"ceval/ceval-exam:computer_architecture": 11,
"ceval/ceval-exam:chinese_language_and_literature": 10,
"ceval/ceval-exam:basic_medicine": 9
}
}
}
}

65
evaluation_reporting/README.md Normal file
View File

@@ -0,0 +1,65 @@
# Evaluation And Report Generation
This folder owns benchmark evaluation, shard merging, and report/summary generation.
It does not train models and does not build model checkpoints.
## Main Files
```text
eval_tokenizer_swap_benchmark.py
merge_tokenizer_swap_benchmark_shards.py
summarize_heldout_capabilities.py
run_public_heldout_eval_8gpu.sh
```
## Input
The default benchmark is the latest public heldout 2K set kept in `dataset_building/`:
```text
dataset_building/heldout_public_mcq_2k_20260607/heldout_public_mcq_2k.jsonl
```
The evaluation script expects a Hugging Face checkpoint directory as `MODEL`.
## Run
```bash
ROOT=/ssd/yi/Tokenizer_Swap \
MODEL=/path/to/checkpoint \
LABEL=my_model \
bash evaluation_reporting/run_public_heldout_eval_8gpu.sh
```
The wrapper launches one shard per GPU, then merges shard outputs.
## Output
Generated outputs go under:
```text
evaluation_reporting/outputs/
```
This directory is ignored by git.
Each evaluation writes per-item JSONL plus summary JSON. The key reported metrics are:
- MCQ accuracy with average-normalized choice logprob
- MCQ accuracy with summed choice logprob
- perplexity
- NLL per token
- bits per byte
## Final Public Heldout 2K Reference
| Model | MCQ acc avg-norm | MCQ acc sum | PPL | NLL/token |
|---|---:|---:|---:|---:|
| Native Qwen3-0.6B tokenizer baseline | 0.2960 | 0.2510 | 61.08 | 3.5399 |
| Remap v2, no training | 0.2940 | 0.2410 | 313.62 | 4.6920 |
| Remap v2 + CPT 1B | 0.3005 | 0.2540 | 71.90 | 3.6580 |
| Remap v2 + CPT 5B | 0.3020 | 0.2615 | 66.95 | 3.5806 |
| Remap v2 + SFT 1M | 0.3105 | 0.2590 | 114.75 | 3.9400 |
| Remap v2 + SFT 1M + v4 continuation | 0.3165 | 0.2595 | 117.33 | 3.9755 |
| Remap v2 + CPT 5B + SFT 1M | 0.3280 | 0.2740 | 88.76 | 3.7899 |

207
evaluation_reporting/eval_tokenizer_swap_benchmark.py Executable file
View File

@@ -0,0 +1,207 @@
#!/usr/bin/env python3
import argparse
import json
import math
from collections import defaultdict
from pathlib import Path
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
def load_jsonl(path, max_items=0, num_shards=1, shard_id=0):
rows = []
with open(path, "r", encoding="utf-8") as f:
for idx, line in enumerate(f):
if num_shards > 1 and idx % num_shards != shard_id:
continue
if max_items and len(rows) >= max_items:
break
rows.append(json.loads(line))
return rows
def encode(tokenizer, text):
return tokenizer.encode(text, add_special_tokens=False)
@torch.inference_mode()
def continuation_nll(model, tokenizer, prompt, continuation, max_length):
prompt_ids = encode(tokenizer, prompt)
cont_ids = encode(tokenizer, continuation)
if not cont_ids:
return None
if len(cont_ids) >= max_length:
cont_ids = cont_ids[: max_length - 1]
prompt_ids = []
else:
prompt_budget = max_length - len(cont_ids)
prompt_ids = prompt_ids[-prompt_budget:]
ids = prompt_ids + cont_ids
labels = [-100] * len(prompt_ids) + cont_ids
if len(ids) < 2:
return None
input_ids = torch.tensor([ids], device=model.device, dtype=torch.long)
label_ids = torch.tensor([labels], device=model.device, dtype=torch.long)
logits = model(input_ids).logits
shift_logits = logits[:, :-1, :].contiguous()
shift_labels = label_ids[:, 1:].contiguous()
mask = shift_labels.ne(-100)
if not mask.any():
return None
log_probs = F.log_softmax(shift_logits, dim=-1)
safe_labels = shift_labels.masked_fill(~mask, 0)
token_log_probs = log_probs.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1)
nll = -token_log_probs[mask].sum().item()
tokens = int(mask.sum().item())
return {"nll": nll, "tokens": tokens}
def score_ppl(model, tokenizer, text, max_length):
ret = continuation_nll(model, tokenizer, "", text, max_length)
if ret is None:
return None
byte_len = max(1, len(text.encode("utf-8")))
ret["nll_per_token"] = ret["nll"] / max(1, ret["tokens"])
ret["ppl"] = math.exp(min(50, ret["nll_per_token"]))
ret["nll_per_byte"] = ret["nll"] / byte_len
ret["bits_per_byte"] = ret["nll"] / (byte_len * math.log(2))
ret["bytes"] = byte_len
return ret
def score_mcq(model, tokenizer, prompt, choices, max_length):
scores = []
for choice in choices:
sep = "" if prompt.endswith((" ", "\n")) else " "
ret = continuation_nll(model, tokenizer, prompt + sep, choice, max_length)
if ret is None:
scores.append({"sum_logprob": -float("inf"), "avg_logprob": -float("inf"), "tokens": 0})
continue
sum_logprob = -ret["nll"]
avg_logprob = sum_logprob / max(1, ret["tokens"])
scores.append(
{
"sum_logprob": sum_logprob,
"avg_logprob": avg_logprob,
"tokens": ret["tokens"],
}
)
pred_avg = max(range(len(scores)), key=lambda i: scores[i]["avg_logprob"])
pred_sum = max(range(len(scores)), key=lambda i: scores[i]["sum_logprob"])
return scores, pred_avg, pred_sum
def mean(xs):
xs = [x for x in xs if x is not None]
return sum(xs) / len(xs) if xs else None
def summarize(results):
groups = defaultdict(list)
for row in results:
groups[row["category"]].append(row)
groups["all"] = results
out = {}
for cat, rows in groups.items():
out[cat] = {
"n": len(rows),
"mcq_acc_avg_norm": mean([x["mcq_correct_avg_norm"] for x in rows]),
"mcq_acc_sum": mean([x["mcq_correct_sum"] for x in rows]),
"ppl_token_mean": mean([x["ppl"]["ppl"] for x in rows if x.get("ppl")]),
"nll_per_token_mean": mean([x["ppl"]["nll_per_token"] for x in rows if x.get("ppl")]),
"bits_per_byte_mean": mean([x["ppl"]["bits_per_byte"] for x in rows if x.get("ppl")]),
"choice_tokens_mean": mean(
[s["tokens"] for x in rows for s in x.get("mcq_scores", []) if s["tokens"]]
),
}
return out
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--benchmark", required=True)
ap.add_argument("--out-dir", required=True)
ap.add_argument("--model-label", default="")
ap.add_argument("--max-items", type=int, default=0)
ap.add_argument("--max-length", type=int, default=2048)
ap.add_argument("--num-shards", type=int, default=1)
ap.add_argument("--shard-id", type=int, default=0)
ap.add_argument("--dtype", choices=["auto", "float16", "bfloat16", "float32"], default="bfloat16")
args = ap.parse_args()
dtype = {
"auto": "auto",
"float16": torch.float16,
"bfloat16": torch.bfloat16,
"float32": torch.float32,
}[args.dtype]
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
label = args.model_label or Path(args.model).name
out_label = label
if args.num_shards > 1:
if args.shard_id < 0 or args.shard_id >= args.num_shards:
raise ValueError("--shard-id must be in [0, --num-shards)")
out_label = f"{label}.shard{args.shard_id:02d}of{args.num_shards:02d}"
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
args.model,
torch_dtype=dtype,
device_map={"": 0},
trust_remote_code=True,
)
model.eval()
rows = load_jsonl(args.benchmark, args.max_items, args.num_shards, args.shard_id)
results = []
for i, row in enumerate(rows, 1):
ppl = score_ppl(model, tokenizer, row["ppl_text"], args.max_length)
mcq_scores, pred_avg, pred_sum = score_mcq(
model, tokenizer, row["mcq_prompt"], row["choices"], args.max_length
)
result = {
"id": row["id"],
"category": row["category"],
"source": row.get("source", ""),
"answer_idx": row["answer_idx"],
"pred_avg_norm": pred_avg,
"pred_sum": pred_sum,
"mcq_correct_avg_norm": int(pred_avg == row["answer_idx"]),
"mcq_correct_sum": int(pred_sum == row["answer_idx"]),
"ppl": ppl,
"mcq_scores": mcq_scores,
}
results.append(result)
if i % 100 == 0:
print(f"[{label}] evaluated {i}/{len(rows)}")
summary = {
"model": args.model,
"model_label": label,
"benchmark": args.benchmark,
"max_items": args.max_items,
"max_length": args.max_length,
"num_shards": args.num_shards,
"shard_id": args.shard_id,
"summary": summarize(results),
}
with (out_dir / f"{out_label}.per_item.jsonl").open("w", encoding="utf-8") as f:
for row in results:
f.write(json.dumps(row, ensure_ascii=False) + "\n")
with (out_dir / f"{out_label}.summary.json").open("w", encoding="utf-8") as f:
json.dump(summary, f, ensure_ascii=False, indent=2)
print(json.dumps(summary, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()

72
evaluation_reporting/merge_tokenizer_swap_benchmark_shards.py Executable file
View File

@@ -0,0 +1,72 @@
#!/usr/bin/env python3
import argparse
import json
import math
import re
from collections import defaultdict
from pathlib import Path
SHARD_RE = re.compile(r"^(?P<label>.+)\.shard\d+of\d+\.per_item\.jsonl$")
def mean(xs):
xs = [x for x in xs if x is not None]
return sum(xs) / len(xs) if xs else None
def summarize(results):
groups = defaultdict(list)
for row in results:
groups[row["category"]].append(row)
groups["all"] = results
out = {}
for cat, rows in groups.items():
out[cat] = {
"n": len(rows),
"mcq_acc_avg_norm": mean([x["mcq_correct_avg_norm"] for x in rows]),
"mcq_acc_sum": mean([x["mcq_correct_sum"] for x in rows]),
"ppl_token_mean": mean([x["ppl"]["ppl"] for x in rows if x.get("ppl")]),
"nll_per_token_mean": mean([x["ppl"]["nll_per_token"] for x in rows if x.get("ppl")]),
"bits_per_byte_mean": mean([x["ppl"]["bits_per_byte"] for x in rows if x.get("ppl")]),
"choice_tokens_mean": mean(
[s["tokens"] for x in rows for s in x.get("mcq_scores", []) if s["tokens"]]
),
}
return out
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--eval-dir", required=True)
args = ap.parse_args()
eval_dir = Path(args.eval_dir)
by_label = defaultdict(list)
for path in sorted(eval_dir.glob("*.shard*of*.per_item.jsonl")):
m = SHARD_RE.match(path.name)
if not m:
continue
label = m.group("label")
with path.open("r", encoding="utf-8") as f:
for line in f:
by_label[label].append(json.loads(line))
for label, rows in sorted(by_label.items()):
rows.sort(key=lambda x: x["id"])
with (eval_dir / f"{label}.per_item.jsonl").open("w", encoding="utf-8") as f:
for row in rows:
f.write(json.dumps(row, ensure_ascii=False) + "\n")
summary = {
"model_label": label,
"merged_from_shards": True,
"n_items": len(rows),
"summary": summarize(rows),
}
with (eval_dir / f"{label}.summary.json").open("w", encoding="utf-8") as f:
json.dump(summary, f, ensure_ascii=False, indent=2)
print(json.dumps(summary, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()

22
evaluation_reporting/run_public_heldout_eval_8gpu.sh Executable file
View File

@@ -0,0 +1,22 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT=${ROOT:-/ssd/yi/Tokenizer_Swap}
NPROC=${NPROC:-8}
MODEL=${MODEL:?Set MODEL to the checkpoint directory to evaluate}
LABEL=${LABEL:-$(basename "$MODEL")}
OUT=${OUT:-$ROOT/evaluation_reporting/outputs/heldout_public_2k_eval/$LABEL}
DATA=${DATA:-$ROOT/dataset_building/heldout_public_mcq_2k_20260607/heldout_public_mcq_2k.jsonl}
mkdir -p "$OUT"
for SHARD in $(seq 0 $((NPROC - 1))); do
CUDA_VISIBLE_DEVICES=$SHARD python "$ROOT/evaluation_reporting/eval_tokenizer_swap_benchmark.py" \
--model "$MODEL" \
--model-label "$LABEL" \
--benchmark "$DATA" \
--out-dir "$OUT" \
--num-shards "$NPROC" \
--shard-id "$SHARD" &
done
wait
python "$ROOT/evaluation_reporting/merge_tokenizer_swap_benchmark_shards.py" --eval-dir "$OUT"

95
evaluation_reporting/summarize_heldout_capabilities.py Normal file
View File

@@ -0,0 +1,95 @@
#!/usr/bin/env python3
import argparse
import collections
import json
import math
CODE_SUBSTRINGS = [
"computer",
"programming",
"operating_system",
"architecture",
"network",
"security",
]
MATH_SUBSTRINGS = [
"math",
"physics",
"chemistry",
"biology",
"statistics",
"probability",
"astronomy",
"anatomy",
"medical",
"electrical",
"engineering",
"machine_learning",
]
def load_jsonl(path):
rows = {}
with open(path, encoding="utf-8") as f:
for line in f:
if line.strip():
row = json.loads(line)
rows[row["id"]] = row
return rows
def bucket(row):
category = row.get("category")
subset = str(row.get("subset") or "").lower()
meta = row.get("metadata") or {}
domain = str(meta.get("high_level_domain") or "").lower()
subdomain = str(meta.get("subdomain") or "").lower()
text = " ".join([subset, domain, subdomain])
if any(x in text for x in CODE_SUBSTRINGS):
return "coding_or_cs"
if category == "gpqa" or any(x in text for x in MATH_SUBSTRINGS):
return "math_or_science_reasoning"
if category == "ceval":
return "chinese_exam"
if category == "mmlu":
return "english_general_mmlu"
return "other"
def mean(xs):
xs = [x for x in xs if x is not None]
return sum(xs) / len(xs) if xs else None
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--benchmark", required=True)
ap.add_argument("--eval", required=True)
args = ap.parse_args()
bench = load_jsonl(args.benchmark)
eval_rows = load_jsonl(args.eval)
groups = collections.defaultdict(list)
for item_id, ev in eval_rows.items():
b = bench[item_id]
groups[bucket(b)].append((b, ev))
out = {}
for name, rows in sorted(groups.items()):
out[name] = {
"n": len(rows),
"mcq_acc_avg_norm": mean([ev.get("mcq_correct_avg_norm") for _, ev in rows]),
"mcq_acc_sum": mean([ev.get("mcq_correct_sum") for _, ev in rows]),
"ppl_token_mean": mean([ev.get("ppl", {}).get("ppl") for _, ev in rows if ev.get("ppl")]),
"nll_per_token_mean": mean([ev.get("ppl", {}).get("nll_per_token") for _, ev in rows if ev.get("ppl")]),
"bits_per_byte_mean": mean([ev.get("ppl", {}).get("bits_per_byte") for _, ev in rows if ev.get("ppl")]),
"subsets_top": dict(collections.Counter(str(b.get("subset")) for b, _ in rows).most_common(20)),
}
print(json.dumps(out, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()

57
model_building/README.md Normal file
View File

@@ -0,0 +1,57 @@
# Model Building
This folder owns construction of the tokenizer-swapped base model.
It contains only the tokenizer swap v2 algorithm. Dataset construction, training, and evaluation live in the other workflow folders.
## Main Files
```text
build_qwen3_dsv4_remap_checkpoint_v2.py
run_remap_v2.sh
```
## Inputs
The remap script needs:
- source Qwen model checkpoint
- source Qwen tokenizer
- target DSV4 tokenizer
Default paths in `run_remap_v2.sh` are environment-variable driven and can be overridden:
```bash
BASE_MODEL=/path/to/Qwen3-0.6B \
DSV_TOKENIZER=/path/to/dsv4_tokenizer \
OUT=/path/to/output_checkpoint \
bash model_building/run_remap_v2.sh
```
## Output
By default, generated checkpoints go to:
```text
model_building/generated_models/
```
This directory is ignored by git. Do not commit checkpoint weights.
## Algorithm Summary
The v2 remap builds DSV4-sized input embedding and LM-head matrices from the source Qwen checkpoint.
For each DSV4 token row, initialization is selected in this priority order:
1. Exact same token surface exists in the Qwen vocab.
2. Functional special-token mapping is available, such as DSV BOS to Qwen `<|im_start|>` and DSV EOS to Qwen EOS.
3. Byte-level token can be decoded, re-tokenized with Qwen, and initialized by averaging the corresponding Qwen rows.
4. Raw token decomposition can be tokenized with Qwen and averaged.
5. Global embedding/head mean fallback.
The script writes the remapped checkpoint plus `tokenizer_remap_v2_report.json` for auditability.
## Output Contract
The output checkpoint is consumed by `model_training/` scripts as `MODEL`.

335
model_building/build_qwen3_dsv4_remap_checkpoint_v2.py Executable file
View File

@@ -0,0 +1,335 @@
#!/usr/bin/env python3
import argparse
import json
import re
import shutil
from collections import Counter, defaultdict
from pathlib import Path
import torch
from tokenizers import Tokenizer
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer
WORKDIR = Path("/ssd/yi/Tokenizer_Swap")
BYTE_FALLBACK_RE = re.compile(r"^<0x[0-9A-Fa-f]{2}>$")
def bytes_to_unicode():
bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
cs = bs[:]
n = 0
for b in range(2**8):
if b not in bs:
bs.append(b)
cs.append(2**8 + n)
n += 1
return dict(zip(bs, [chr(n) for n in cs]))
BYTE_DECODER = {v: k for k, v in bytes_to_unicode().items()}
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--qwen-model", default=str(WORKDIR / "model_building/source_models/Qwen3-0.6B"))
p.add_argument("--dsv-tokenizer", default=str(WORKDIR / "model_building/source_tokenizers/dsv4_flash"))
p.add_argument("--out", default=str(WORKDIR / "model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2"))
p.add_argument("--bos-source", choices=["im_start", "endoftext", "eos"], default="im_start")
p.add_argument("--pad-source", choices=["pad", "eos"], default="eos")
return p.parse_args()
def content_from_cfg(cfg, key):
value = cfg.get(key)
if isinstance(value, dict):
return value.get("content")
return value
def bytelevel_decode(token, special_tokens):
if token in special_tokens or token.startswith(("<|", "<｜")):
return token, False, "special"
if BYTE_FALLBACK_RE.fullmatch(token):
return token, False, "byte_fallback"
try:
bs = bytes(BYTE_DECODER[ch] for ch in token)
except KeyError:
return token, False, "not_bytelevel"
decoded = bs.decode("utf-8", errors="replace")
return decoded, decoded != token, "decoded"
def safe_encode(tokenizer, text, old_rows):
if text is None or text == "":
return []
ids = tokenizer.encode(text, add_special_tokens=False)
return [i for i in ids if 0 <= i < old_rows]
def has_cjk(text):
for ch in text:
code = ord(ch)
if (
0x4E00 <= code <= 0x9FFF
or 0x3400 <= code <= 0x4DBF
or 0x20000 <= code <= 0x2A6DF
or 0x2A700 <= code <= 0x2B73F
or 0x2B740 <= code <= 0x2B81F
or 0x2B820 <= code <= 0x2CEAF
or 0xF900 <= code <= 0xFAFF
):
return True
return False
def token_shape(decoded):
if decoded is None:
return "unknown"
if "<EFBFBD>" in decoded:
return "utf8_fragment"
if decoded == "" or decoded.isspace():
return "whitespace"
if has_cjk(decoded):
return "cjk"
if decoded.isalpha() and decoded.isascii():
return "ascii_word"
if decoded.isnumeric():
return "numeric"
if re.search(r"[A-Za-z]", decoded) and re.search(r"[_./\\{}()[\]<>:=+\-*#@$%^&|;]", decoded):
return "code_like"
if all(not ch.isalnum() for ch in decoded):
return "punct_symbol"
if any(ord(ch) > 127 for ch in decoded):
return "non_ascii"
return "ascii_mixed"
def qwen_special_lookup(qwen_tok, qwen_vocab):
lookup = {}
for name in ["bos_token", "eos_token", "pad_token", "unk_token"]:
tok = getattr(qwen_tok, name, None)
if tok in qwen_vocab:
lookup[name] = tok
for tok in getattr(qwen_tok, "additional_special_tokens", []) or []:
if tok in qwen_vocab:
lookup[tok] = tok
for tok in ["<|im_start|>", "<|im_end|>", "<|endoftext|>", "<think>", "</think>"]:
if tok in qwen_vocab:
lookup[tok] = tok
return lookup
def dsv_special_tokens(dsv_dir, dsv_data):
cfg = json.loads((dsv_dir / "tokenizer_config.json").read_text(encoding="utf-8"))
specials = set()
for key in ["bos_token", "eos_token", "pad_token", "unk_token"]:
tok = content_from_cfg(cfg, key)
if tok:
specials.add(tok)
for item in dsv_data.get("added_tokens", []):
if item.get("special"):
specials.add(str(item.get("content", "")))
return cfg, specials
def build_functional_map(dsv_cfg, qwen_tok, qwen_vocab, args):
qwen_specials = qwen_special_lookup(qwen_tok, qwen_vocab)
mapping = {}
reasons = {}
dsv_bos = content_from_cfg(dsv_cfg, "bos_token")
dsv_eos = content_from_cfg(dsv_cfg, "eos_token")
dsv_pad = content_from_cfg(dsv_cfg, "pad_token") or dsv_eos
def add(dst, src, reason):
if dst and src and src in qwen_vocab:
mapping[dst] = src
reasons[dst] = reason
if args.bos_source == "im_start":
add(dsv_bos, "<|im_start|>", "dsv_bos_to_qwen_im_start")
elif args.bos_source == "endoftext":
add(dsv_bos, "<|endoftext|>", "dsv_bos_to_qwen_endoftext")
else:
add(dsv_bos, getattr(qwen_tok, "eos_token", None), "dsv_bos_to_qwen_eos")
add(dsv_eos, getattr(qwen_tok, "eos_token", None), "dsv_eos_to_qwen_eos")
if args.pad_source == "pad":
add(dsv_pad, getattr(qwen_tok, "pad_token", None), "dsv_pad_to_qwen_pad")
else:
add(dsv_pad, getattr(qwen_tok, "eos_token", None), "dsv_pad_to_qwen_eos")
# Same-surface thinking tags if both tokenizers expose them.
for dst, src, reason in [
("<think>", "<think>", "thinking_start_exact_function"),
("</think>", "</think>", "thinking_end_exact_function"),
]:
add(dst, src, reason)
return mapping, reasons, qwen_specials
def row_mean(old_weight, ids):
return old_weight[ids].float().mean(dim=0).to(old_weight.dtype).cpu()
def build_remap(old_weight, qwen_vocab, qwen_tok, dsv_vocab, dsv_specials, functional_map, functional_reasons, desc):
old_rows, hidden = old_weight.shape
new_rows = max(dsv_vocab.values()) + 1
new_weight = torch.empty((new_rows, hidden), dtype=old_weight.dtype, device="cpu")
global_mean = old_weight.float().mean(dim=0).to(old_weight.dtype).cpu()
stats = Counter()
shape_by_method = defaultdict(Counter)
examples = defaultdict(list)
by_id = sorted(dsv_vocab.items(), key=lambda x: x[1])
for tok, new_id in tqdm(by_id, desc=desc):
decoded, changed, decode_status = bytelevel_decode(tok, dsv_specials)
shape = token_shape(decoded)
old_id = qwen_vocab.get(tok)
if old_id is not None and 0 <= old_id < old_rows:
new_weight[new_id].copy_(old_weight[old_id].cpu())
method = "exact_copy"
elif tok in functional_map and functional_map[tok] in qwen_vocab:
src = functional_map[tok]
new_weight[new_id].copy_(old_weight[qwen_vocab[src]].cpu())
method = "special_function_copy"
else:
ids = []
if decode_status == "decoded" and "<EFBFBD>" not in decoded:
ids = safe_encode(qwen_tok, decoded, old_rows)
if ids:
new_weight[new_id].copy_(row_mean(old_weight, ids))
method = "decoded_decomposition_avg"
else:
raw_ids = safe_encode(qwen_tok, tok, old_rows)
if raw_ids:
new_weight[new_id].copy_(row_mean(old_weight, raw_ids))
method = "raw_decomposition_avg"
else:
new_weight[new_id].copy_(global_mean)
method = "mean_fallback"
stats[method] += 1
shape_by_method[method][shape] += 1
if len(examples[method]) < 40:
ex = {
"token": tok,
"id": new_id,
"decoded": decoded,
"shape": shape,
}
if method == "special_function_copy":
ex["source_token"] = functional_map.get(tok)
ex["reason"] = functional_reasons.get(tok)
examples[method].append(ex)
return new_weight, dict(stats), {k: dict(v) for k, v in shape_by_method.items()}, dict(examples)
def read_dsv_tokenizer(dsv_path):
tokenizer_json = dsv_path / "tokenizer.json"
dsv_raw = Tokenizer.from_file(str(tokenizer_json))
dsv_data = json.loads(tokenizer_json.read_text(encoding="utf-8"))
return dsv_raw, dsv_data
def main():
args = parse_args()
qwen_path = Path(args.qwen_model)
dsv_path = Path(args.dsv_tokenizer)
out = Path(args.out)
out.mkdir(parents=True, exist_ok=True)
qwen_tok = AutoTokenizer.from_pretrained(qwen_path, trust_remote_code=True)
dsv_raw, dsv_data = read_dsv_tokenizer(dsv_path)
dsv_cfg, dsv_specials = dsv_special_tokens(dsv_path, dsv_data)
qwen_vocab = qwen_tok.get_vocab()
dsv_vocab = dsv_raw.get_vocab()
new_vocab_size = max(dsv_vocab.values()) + 1
functional_map, functional_reasons, qwen_specials = build_functional_map(dsv_cfg, qwen_tok, qwen_vocab, args)
model = AutoModelForCausalLM.from_pretrained(
qwen_path,
torch_dtype=torch.bfloat16,
device_map="cpu",
trust_remote_code=True,
)
model.eval()
old_embed = model.get_input_embeddings().weight.detach().cpu()
old_out = model.get_output_embeddings().weight.detach().cpu()
new_embed, embed_stats, embed_shapes, embed_examples = build_remap(
old_embed, qwen_vocab, qwen_tok, dsv_vocab, dsv_specials, functional_map, functional_reasons, "embed-v2"
)
if old_out.data_ptr() == old_embed.data_ptr():
new_out = new_embed
out_stats = embed_stats.copy()
out_shapes = embed_shapes
out_examples = embed_examples
else:
new_out, out_stats, out_shapes, out_examples = build_remap(
old_out, qwen_vocab, qwen_tok, dsv_vocab, dsv_specials, functional_map, functional_reasons, "lm-head-v2"
)
model.resize_token_embeddings(new_vocab_size)
model.get_input_embeddings().weight.data.copy_(new_embed)
model.get_output_embeddings().weight.data.copy_(new_out)
dsv_bos = content_from_cfg(dsv_cfg, "bos_token")
dsv_eos = content_from_cfg(dsv_cfg, "eos_token")
dsv_pad = content_from_cfg(dsv_cfg, "pad_token") or dsv_eos
model.config.vocab_size = new_vocab_size
model.config.bos_token_id = dsv_vocab.get(dsv_bos, 0)
model.config.eos_token_id = dsv_vocab.get(dsv_eos, 1)
model.config.pad_token_id = dsv_vocab.get(dsv_pad, dsv_vocab.get(dsv_eos, 1))
if hasattr(model, "generation_config"):
model.generation_config.bos_token_id = model.config.bos_token_id
model.generation_config.eos_token_id = model.config.eos_token_id
model.generation_config.pad_token_id = model.config.pad_token_id
model.tie_weights()
model.save_pretrained(out, safe_serialization=True)
for name in ["tokenizer.json", "tokenizer_config.json"]:
shutil.copy2(dsv_path / name, out / name)
report = {
"algorithm": "remap_v2_exact_special_decoded_decomposition",
"qwen_model": str(qwen_path),
"dsv_tokenizer": str(dsv_path),
"out": str(out),
"old_embedding_shape": list(old_embed.shape),
"old_lm_head_shape": list(old_out.shape),
"new_vocab_size": new_vocab_size,
"qwen_vocab_size": len(qwen_vocab),
"dsv_vocab_size": len(dsv_vocab),
"common_token_strings": len(set(qwen_vocab) & set(dsv_vocab)),
"functional_map": functional_map,
"functional_reasons": functional_reasons,
"qwen_specials_detected": qwen_specials,
"dsv_specials_detected": sorted(dsv_specials),
"embed_init_stats": embed_stats,
"lm_head_init_stats": out_stats,
"embed_shape_by_method": embed_shapes,
"lm_head_shape_by_method": out_shapes,
"embed_examples": embed_examples,
"lm_head_examples": out_examples,
"bos_token_id": model.config.bos_token_id,
"eos_token_id": model.config.eos_token_id,
"pad_token_id": model.config.pad_token_id,
"bos_token": dsv_bos,
"eos_token": dsv_eos,
"pad_token": dsv_pad,
"bos_source_policy": args.bos_source,
"pad_source_policy": args.pad_source,
}
(out / "tokenizer_remap_v2_report.json").write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
print(json.dumps({k: report[k] for k in ["algorithm", "out", "new_vocab_size", "embed_init_stats", "functional_map", "bos_source_policy", "pad_source_policy"]}, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()

12
model_building/run_remap_v2.sh Executable file
View File

@@ -0,0 +1,12 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT=${ROOT:-/ssd/yi/Tokenizer_Swap}
BASE_MODEL=${BASE_MODEL:-/ssd/yi/tokenizer_swap_cepe/models/Qwen3-0.6B}
DSV_TOKENIZER=${DSV_TOKENIZER:-/ssd/yi/tokenizer_swap_cepe/models/tokenizers/dsv4_flash}
OUT=${OUT:-$ROOT/model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2}
python "$ROOT/model_building/build_qwen3_dsv4_remap_checkpoint_v2.py" \
--qwen-model "$BASE_MODEL" \
--dsv-tokenizer "$DSV_TOKENIZER" \
--out "$OUT"

64
model_training/README.md Normal file
View File

@@ -0,0 +1,64 @@
# Model Training
This folder owns full-parameter training recipes for the final experiments.
It does not build datasets and does not perform tokenizer remapping. It consumes artifacts produced by `dataset_building/` and `model_building/`.
## Training Implementations
```text
train_dsv4_tokenized_full_sft.py
train_cpt_packed_full.py
```
`train_dsv4_tokenized_full_sft.py` trains on pre-tokenized chat JSONL. Prompt tokens are masked with `-100`; assistant tokens are optimized.
`train_cpt_packed_full.py` trains next-token prediction over packed CPT blocks.
## Final Recipe Scripts
```text
run_sft1m_remap_v2_5epoch.sh
run_sft1m_remap_v2_then_v4_noupsample_5epoch_bsz16.sh
run_cpt1b_seed42_train_eval.sh
run_cpt5b_seed42_train_eval.sh
run_cpt5b_then_sft1m_5epoch.sh
```
## Common Environment Variables
The run scripts are configurable through environment variables:
```text
ROOT repo root, default /ssd/yi/Tokenizer_Swap
NPROC number of GPUs/processes, default 8
MODEL input checkpoint
DATA packed CPT dataset directory
TRAIN tokenized SFT train file
EVAL tokenized SFT validation file
OUT output checkpoint directory
```
Example:
```bash
ROOT=/ssd/yi/Tokenizer_Swap \
MODEL=/ssd/yi/Tokenizer_Swap/model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2 \
DATA=/ssd/yi/Tokenizer_Swap/dataset_building/generated/cpt_packed_5b_seq8192_seed42_stratified \
OUT=/ssd/yi/Tokenizer_Swap/model_training/checkpoints/cpt5b \
bash model_training/run_cpt5b_seed42_train_eval.sh
```
## Output
Generated checkpoints go under:
```text
model_training/checkpoints/
```
This directory is ignored by git. Do not commit model weights, optimizer states, or partial checkpoints.
## Output Contract
Trained checkpoints are consumed by `evaluation_reporting/` as `MODEL`.

21
model_training/run_cpt1b_seed42_train_eval.sh Executable file
View File

@@ -0,0 +1,21 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT=${ROOT:-/ssd/yi/Tokenizer_Swap}
NPROC=${NPROC:-8}
MODEL=${MODEL:-$ROOT/model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2}
DATA=${DATA:-$ROOT/dataset_building/generated/cpt_packed_1b_seq8192_seed42_stratified}
OUT=${OUT:-$ROOT/model_training/checkpoints/qwen3_06b_dsv4_remap_v2_cpt_1b_seed42_seq8192_bsz3acc3}
mkdir -p "$(dirname "$OUT")"
torchrun --nproc_per_node "$NPROC" "$ROOT/model_training/train_cpt_packed_full.py" \
--model "$MODEL" \
--data "$DATA" \
--out "$OUT" \
--epochs 1 \
--batch-size 3 \
--grad-accum 3 \
--lr 2e-5 \
--eval-steps 500 \
--save-steps 500 \
--num-workers 2

21
model_training/run_cpt5b_seed42_train_eval.sh Executable file
View File

@@ -0,0 +1,21 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT=${ROOT:-/ssd/yi/Tokenizer_Swap}
NPROC=${NPROC:-8}
MODEL=${MODEL:-$ROOT/model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2}
DATA=${DATA:-$ROOT/dataset_building/generated/cpt_packed_5b_seq8192_seed42_stratified}
OUT=${OUT:-$ROOT/model_training/checkpoints/qwen3_06b_dsv4_remap_v2_cpt_5b_seed42_seq8192_bsz3acc3}
mkdir -p "$(dirname "$OUT")"
torchrun --nproc_per_node "$NPROC" "$ROOT/model_training/train_cpt_packed_full.py" \
--model "$MODEL" \
--data "$DATA" \
--out "$OUT" \
--epochs 1 \
--batch-size 3 \
--grad-accum 3 \
--lr 2e-5 \
--eval-steps 2000 \
--save-steps 2000 \
--num-workers 2

24
model_training/run_cpt5b_then_sft1m_5epoch.sh Executable file
View File

@@ -0,0 +1,24 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT=${ROOT:-/ssd/yi/Tokenizer_Swap}
NPROC=${NPROC:-8}
MODEL=${MODEL:-$ROOT/model_training/checkpoints/qwen3_06b_dsv4_remap_v2_cpt_5b_seed42_seq8192_bsz3acc3}
TRAIN=${TRAIN:-$ROOT/dataset_building/generated/dsv4_chat_tokenized_v4_noupsample_nobbh_921k/train_dsv4_chat_tokenized.jsonl.gz}
EVAL=${EVAL:-$ROOT/dataset_building/generated/dsv4_chat_tokenized_v4_noupsample_nobbh_921k/validation_dsv4_chat_tokenized.jsonl.gz}
OUT=${OUT:-$ROOT/model_training/checkpoints/qwen3_06b_dsv4_remap_v2_cpt5b_then_sft1m_5epoch}
mkdir -p "$(dirname "$OUT")"
torchrun --nproc_per_node "$NPROC" "$ROOT/model_training/train_dsv4_tokenized_full_sft.py" \
--model "$MODEL" \
--train "$TRAIN" \
--eval "$EVAL" \
--out "$OUT" \
--max-length 2048 \
--epochs 5 \
--batch-size 8 \
--grad-accum 16 \
--lr 5e-5 \
--eval-steps 500 \
--save-steps 300 \
--num-workers 2

24
model_training/run_sft1m_remap_v2_5epoch.sh Executable file
View File

@@ -0,0 +1,24 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT=${ROOT:-/ssd/yi/Tokenizer_Swap}
NPROC=${NPROC:-8}
MODEL=${MODEL:-$ROOT/model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2}
TRAIN=${TRAIN:-$ROOT/dataset_building/generated/dsv4_chat_tokenized_alt_sources_1m_fixed_eval_20260607/train_dsv4_chat_tokenized.jsonl.gz}
EVAL=${EVAL:-$ROOT/dataset_building/generated/dsv4_chat_tokenized_alt_sources_1m_fixed_eval_20260607/validation_dsv4_chat_tokenized.jsonl.gz}
OUT=${OUT:-$ROOT/model_training/checkpoints/dsv4_chat_full_sft_remap_v2_alt1m_5epoch_bsz8_accum16}
mkdir -p "$(dirname "$OUT")"
torchrun --nproc_per_node "$NPROC" "$ROOT/model_training/train_dsv4_tokenized_full_sft.py" \
--model "$MODEL" \
--train "$TRAIN" \
--eval "$EVAL" \
--out "$OUT" \
--max-length 2048 \
--epochs 5 \
--batch-size 8 \
--grad-accum 16 \
--lr 5e-5 \
--eval-steps 500 \
--save-steps 500 \
--num-workers 2

24
model_training/run_sft1m_remap_v2_then_v4_noupsample_5epoch_bsz16.sh Executable file
View File

@@ -0,0 +1,24 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT=${ROOT:-/ssd/yi/Tokenizer_Swap}
NPROC=${NPROC:-8}
MODEL=${MODEL:-$ROOT/model_training/checkpoints/dsv4_chat_full_sft_remap_v2_alt1m_5epoch_bsz8_accum16}
TRAIN=${TRAIN:-$ROOT/dataset_building/generated/dsv4_chat_tokenized_v4_noupsample_nobbh_921k/train_dsv4_chat_tokenized.jsonl.gz}
EVAL=${EVAL:-$ROOT/dataset_building/generated/dsv4_chat_tokenized_v4_noupsample_nobbh_921k/validation_dsv4_chat_tokenized.jsonl.gz}
OUT=${OUT:-$ROOT/model_training/checkpoints/dsv4_chat_full_sft_remap_v2_alt1m_then_v4_noupsample_5epoch_bsz16}
mkdir -p "$(dirname "$OUT")"
torchrun --nproc_per_node "$NPROC" "$ROOT/model_training/train_dsv4_tokenized_full_sft.py" \
--model "$MODEL" \
--train "$TRAIN" \
--eval "$EVAL" \
--out "$OUT" \
--max-length 2048 \
--epochs 5 \
--batch-size 8 \
--grad-accum 16 \
--lr 5e-5 \
--eval-steps 500 \
--save-steps 500 \
--num-workers 2

247
model_training/train_cpt_packed_full.py Normal file
View File

@@ -0,0 +1,247 @@
#!/usr/bin/env python3
import argparse
import json
import math
import os
import time
from bisect import bisect_right
from contextlib import nullcontext
from pathlib import Path
import numpy as np
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Dataset, DistributedSampler
from transformers import AutoModelForCausalLM, AutoTokenizer, get_cosine_schedule_with_warmup
def is_dist():
return int(os.environ.get("WORLD_SIZE", "1")) > 1
def rank():
return int(os.environ.get("RANK", "0"))
def local_rank():
return int(os.environ.get("LOCAL_RANK", "0"))
def is_main():
return rank() == 0
class PackedBlockDataset(Dataset):
def __init__(self, data_dir, split):
self.data_dir = Path(data_dir)
manifest = json.loads((self.data_dir / "manifest.json").read_text(encoding="utf-8"))
self.seq_len = int(manifest.get("seq_len", 8192))
shards = manifest[f"{split}_shards"]
if not shards:
raise ValueError(f"no {split} shards in {self.data_dir}")
self.paths = [self.data_dir / x["path"] for x in shards]
self.lengths = [int(x["blocks"]) for x in shards]
self.cum = np.cumsum(self.lengths).tolist()
self._arrays = [None] * len(self.paths)
def __len__(self):
return self.cum[-1]
def _array(self, shard_idx):
arr = self._arrays[shard_idx]
if arr is None:
arr = np.load(self.paths[shard_idx], mmap_mode="r")
self._arrays[shard_idx] = arr
return arr
def __getitem__(self, idx):
shard_idx = bisect_right(self.cum, idx)
prev = 0 if shard_idx == 0 else self.cum[shard_idx - 1]
row_idx = idx - prev
ids = np.asarray(self._array(shard_idx)[row_idx], dtype=np.int64)
return torch.from_numpy(ids)
def collate(batch):
input_ids = torch.stack(batch, dim=0).long()
return {"input_ids": input_ids, "labels": input_ids.clone()}
@torch.no_grad()
def evaluate(model, loader, device, max_batches=0):
model.eval()
total_loss = torch.tensor(0.0, device=device)
total_tokens = torch.tensor(0.0, device=device)
batches = 0
module = model.module if isinstance(model, DDP) else model
for batch in loader:
batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
out = module(**batch)
ntok = batch["labels"].numel()
total_loss += out.loss.float() * ntok
total_tokens += ntok
batches += 1
if max_batches and batches >= max_batches:
break
if is_dist():
dist.all_reduce(total_loss, op=dist.ReduceOp.SUM)
dist.all_reduce(total_tokens, op=dist.ReduceOp.SUM)
loss = (total_loss / total_tokens.clamp_min(1)).item()
model.train()
return {"loss": loss, "ppl": math.exp(min(loss, 20)), "tokens": int(total_tokens.item()), "batches": batches}
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--model", default="/ssd/yi/Tokenizer_Swap/model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
p.add_argument("--data", default="/ssd/yi/Tokenizer_Swap/dataset_building/generated/cpt_packed_1b_seq8192_20260614")
p.add_argument("--out", default="/ssd/yi/Tokenizer_Swap/model_training/checkpoints/qwen3_06b_dsv4_remap_v2_cpt_1b_seq8192_20260614")
p.add_argument("--epochs", type=float, default=1.0)
p.add_argument("--batch-size", type=int, default=2)
p.add_argument("--grad-accum", type=int, default=4)
p.add_argument("--lr", type=float, default=2e-5)
p.add_argument("--warmup-ratio", type=float, default=0.03)
p.add_argument("--eval-steps", type=int, default=100)
p.add_argument("--save-steps", type=int, default=500)
p.add_argument("--max-steps", type=int, default=0)
p.add_argument("--eval-max-batches", type=int, default=32)
p.add_argument("--num-workers", type=int, default=2)
p.add_argument("--gradient-checkpointing", action="store_true")
return p.parse_args()
def main():
args = parse_args()
if is_dist():
dist.init_process_group(backend="nccl")
torch.cuda.set_device(local_rank())
device = torch.device("cuda", local_rank())
else:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
out_dir = Path(args.out)
if is_main():
out_dir.mkdir(parents=True, exist_ok=True)
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.bfloat16, trust_remote_code=True).to(device)
model.config.use_cache = False
if args.gradient_checkpointing:
model.gradient_checkpointing_enable()
for p in model.parameters():
p.requires_grad_(True)
if is_dist():
model = DDP(model, device_ids=[local_rank()], output_device=local_rank(), find_unused_parameters=False)
train_ds = PackedBlockDataset(args.data, "train")
eval_ds = PackedBlockDataset(args.data, "eval")
train_sampler = DistributedSampler(train_ds, shuffle=True) if is_dist() else None
eval_sampler = DistributedSampler(eval_ds, shuffle=False) if is_dist() else None
train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=train_sampler is None, sampler=train_sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=collate)
eval_loader = DataLoader(eval_ds, batch_size=args.batch_size, shuffle=False, sampler=eval_sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=collate)
steps_per_epoch = math.ceil(len(train_loader) / args.grad_accum)
total_steps = int(math.ceil(steps_per_epoch * args.epochs))
if args.max_steps > 0:
total_steps = min(total_steps, args.max_steps)
warmup_steps = int(total_steps * args.warmup_ratio)
optim = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.1)
sched = get_cosine_schedule_with_warmup(optim, warmup_steps, total_steps)
meta = {
"model": args.model,
"data": args.data,
"out": args.out,
"epochs": args.epochs,
"batch_size_per_rank": args.batch_size,
"grad_accum": args.grad_accum,
"world_size": int(os.environ.get("WORLD_SIZE", "1")),
"effective_batch_tokens": args.batch_size * args.grad_accum * int(os.environ.get("WORLD_SIZE", "1")) * train_ds.seq_len,
"lr": args.lr,
"train_blocks": len(train_ds),
"eval_blocks": len(eval_ds),
"total_steps": total_steps,
"warmup_steps": warmup_steps,
}
if is_main():
(out_dir / "train_config.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
print(json.dumps(meta, ensure_ascii=False, indent=2), flush=True)
log_path = out_dir / "train_log.jsonl"
start = time.time()
step = 0
accum_loss = 0.0
accum_tokens = 0
model.train()
optim.zero_grad(set_to_none=True)
initial_eval = evaluate(model, eval_loader, device, args.eval_max_batches)
if is_main():
with log_path.open("a", encoding="utf-8") as f:
f.write(json.dumps({"event": "initial_eval", **initial_eval}, ensure_ascii=False) + "\n")
print(json.dumps({"event": "initial_eval", **initial_eval}, ensure_ascii=False), flush=True)
epoch = 0
while step < total_steps:
if train_sampler is not None:
train_sampler.set_epoch(epoch)
epoch += 1
for batch_idx, batch in enumerate(train_loader):
batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
sync_grad = (batch_idx + 1) % args.grad_accum == 0
sync_context = model.no_sync() if isinstance(model, DDP) and not sync_grad else nullcontext()
with sync_context:
out = model(**batch)
loss = out.loss / args.grad_accum
loss.backward()
ntok = batch["labels"].numel()
accum_loss += float(out.loss.item()) * ntok
accum_tokens += ntok
if (batch_idx + 1) % args.grad_accum == 0:
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optim.step()
sched.step()
optim.zero_grad(set_to_none=True)
step += 1
if is_main() and (step == 1 or step % 10 == 0):
rec = {"event": "train", "step": step, "loss": accum_loss / max(accum_tokens, 1), "tokens": accum_tokens, "lr": sched.get_last_lr()[0], "elapsed_sec": time.time() - start}
print(json.dumps(rec, ensure_ascii=False), flush=True)
with log_path.open("a", encoding="utf-8") as f:
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
accum_loss = 0.0
accum_tokens = 0
if args.eval_steps and step % args.eval_steps == 0:
rec = {"event": "eval", "step": step, **evaluate(model, eval_loader, device, args.eval_max_batches)}
if is_main():
print(json.dumps(rec, ensure_ascii=False), flush=True)
with log_path.open("a", encoding="utf-8") as f:
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
if is_main() and args.save_steps and step % args.save_steps == 0:
ckpt = out_dir / f"checkpoint-{step}"
(model.module if isinstance(model, DDP) else model).save_pretrained(ckpt, safe_serialization=True)
tokenizer.save_pretrained(ckpt)
if step >= total_steps:
break
final_eval = evaluate(model, eval_loader, device, args.eval_max_batches)
if is_main():
module = model.module if isinstance(model, DDP) else model
module.save_pretrained(out_dir, safe_serialization=True)
tokenizer.save_pretrained(out_dir)
(out_dir / "DONE").write_text("ok\n", encoding="utf-8")
with log_path.open("a", encoding="utf-8") as f:
f.write(json.dumps({"event": "final_eval", **final_eval}, ensure_ascii=False) + "\n")
print(json.dumps({"event": "final_eval", **final_eval}, ensure_ascii=False), flush=True)
if is_dist():
dist.destroy_process_group()
if __name__ == "__main__":
main()

303
model_training/train_dsv4_tokenized_full_sft.py Executable file
View File

@@ -0,0 +1,303 @@
#!/usr/bin/env python3
import argparse
import gzip
import json
import math
import os
import time
from pathlib import Path
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Dataset, DistributedSampler
from transformers import AutoModelForCausalLM, AutoTokenizer, get_cosine_schedule_with_warmup
def is_dist():
return int(os.environ.get("WORLD_SIZE", "1")) > 1
def rank():
return int(os.environ.get("RANK", "0"))
def local_rank():
return int(os.environ.get("LOCAL_RANK", "0"))
def is_main():
return rank() == 0
def open_text(path: Path):
if path.suffix == ".gz":
return gzip.open(path, "rt", encoding="utf-8")
return path.open("r", encoding="utf-8")
class TokenizedSFTDataset(Dataset):
def __init__(self, path: str, max_length: int):
self.rows = []
self.max_length = max_length
with open_text(Path(path)) as f:
for line in f:
if line.strip():
row = json.loads(line)
self.rows.append(
{
"input_ids": row["input_ids"][:max_length],
"labels": row["labels"][:max_length],
"id": row.get("id"),
}
)
def __len__(self):
return len(self.rows)
def __getitem__(self, idx):
return self.rows[idx]
def collate(batch, pad_id):
max_len = max(len(x["input_ids"]) for x in batch)
input_ids, labels, attention_mask = [], [], []
for item in batch:
ids = item["input_ids"]
lab = item["labels"]
pad = max_len - len(ids)
input_ids.append(ids + [pad_id] * pad)
labels.append(lab + [-100] * pad)
attention_mask.append([1] * len(ids) + [0] * pad)
return {
"input_ids": torch.tensor(input_ids, dtype=torch.long),
"labels": torch.tensor(labels, dtype=torch.long),
"attention_mask": torch.tensor(attention_mask, dtype=torch.long),
}
def set_trainable_full(model):
for p in model.parameters():
p.requires_grad_(True)
return ["all_parameters"]
@torch.no_grad()
def evaluate(model, loader, device, max_batches=0):
model.eval()
total_loss = torch.tensor(0.0, device=device)
total_tokens = torch.tensor(0.0, device=device)
batches = 0
module = model.module if isinstance(model, DDP) else model
for batch in loader:
batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
out = module(**batch) if isinstance(model, DDP) else model(**batch)
ntok = (batch["labels"] != -100).sum().float()
total_loss += out.loss.float() * ntok
total_tokens += ntok
batches += 1
if max_batches and batches >= max_batches:
break
if is_dist():
dist.all_reduce(total_loss, op=dist.ReduceOp.SUM)
dist.all_reduce(total_tokens, op=dist.ReduceOp.SUM)
loss = (total_loss / total_tokens.clamp_min(1)).item()
model.train()
return {"loss": loss, "ppl": math.exp(min(loss, 20)), "tokens": int(total_tokens.item()), "batches": batches}
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--model", default="/ssd/yi/Tokenizer_Swap/model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
p.add_argument("--train", default="/ssd/yi/Tokenizer_Swap/dataset_building/generated/dsv4_chat_tokenized_fixed100k_20260605/train_dsv4_chat_tokenized.jsonl.gz")
p.add_argument("--eval", default="/ssd/yi/Tokenizer_Swap/dataset_building/generated/dsv4_chat_tokenized_fixed100k_20260605/validation_dsv4_chat_tokenized.jsonl.gz")
p.add_argument("--out", default="/ssd/yi/Tokenizer_Swap/model_training/checkpoints/dsv4_chat_full_sft_fixed100k_20260605")
p.add_argument("--max-length", type=int, default=2048)
p.add_argument("--epochs", type=float, default=3.0)
p.add_argument("--batch-size", type=int, default=2)
p.add_argument("--grad-accum", type=int, default=8)
p.add_argument("--lr", type=float, default=5e-5)
p.add_argument("--warmup-ratio", type=float, default=0.03)
p.add_argument("--eval-steps", type=int, default=200)
p.add_argument("--save-steps", type=int, default=0)
p.add_argument("--max-steps", type=int, default=0)
p.add_argument("--eval-max-batches", type=int, default=0)
p.add_argument("--num-workers", type=int, default=2)
p.add_argument("--gradient-checkpointing", action="store_true")
return p.parse_args()
def main():
args = parse_args()
if is_dist():
dist.init_process_group(backend="nccl")
torch.cuda.set_device(local_rank())
device = torch.device("cuda", local_rank())
else:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
out_dir = Path(args.out)
if is_main():
out_dir.mkdir(parents=True, exist_ok=True)
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
if tokenizer.pad_token_id is None:
tokenizer.pad_token = tokenizer.eos_token
pad_id = tokenizer.pad_token_id or tokenizer.eos_token_id
model = AutoModelForCausalLM.from_pretrained(
args.model,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
).to(device)
model.config.use_cache = False
if args.gradient_checkpointing:
model.gradient_checkpointing_enable()
trainable_names = set_trainable_full(model)
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
total_params = sum(p.numel() for p in model.parameters())
if is_dist():
model = DDP(model, device_ids=[local_rank()], output_device=local_rank(), find_unused_parameters=False)
train_ds = TokenizedSFTDataset(args.train, args.max_length)
eval_ds = TokenizedSFTDataset(args.eval, args.max_length)
train_sampler = DistributedSampler(train_ds, shuffle=True) if is_dist() else None
eval_sampler = DistributedSampler(eval_ds, shuffle=False) if is_dist() else None
train_loader = DataLoader(
train_ds,
batch_size=args.batch_size,
shuffle=train_sampler is None,
sampler=train_sampler,
num_workers=args.num_workers,
pin_memory=True,
collate_fn=lambda b: collate(b, pad_id),
)
eval_loader = DataLoader(
eval_ds,
batch_size=args.batch_size,
shuffle=False,
sampler=eval_sampler,
num_workers=args.num_workers,
pin_memory=True,
collate_fn=lambda b: collate(b, pad_id),
)
steps_per_epoch = math.ceil(len(train_loader) / args.grad_accum)
total_steps = int(math.ceil(steps_per_epoch * args.epochs))
if args.max_steps > 0:
total_steps = min(total_steps, args.max_steps)
warmup_steps = int(total_steps * args.warmup_ratio)
optim = torch.optim.AdamW((p for p in model.parameters() if p.requires_grad), lr=args.lr, weight_decay=0.01)
sched = get_cosine_schedule_with_warmup(optim, warmup_steps, total_steps)
meta = {
"model": args.model,
"train": args.train,
"eval": args.eval,
"out": args.out,
"max_length": args.max_length,
"epochs": args.epochs,
"batch_size_per_rank": args.batch_size,
"grad_accum": args.grad_accum,
"world_size": int(os.environ.get("WORLD_SIZE", "1")),
"effective_batch_examples": args.batch_size * args.grad_accum * int(os.environ.get("WORLD_SIZE", "1")),
"lr": args.lr,
"train_examples": len(train_ds),
"eval_examples": len(eval_ds),
"total_steps": total_steps,
"warmup_steps": warmup_steps,
"trainable_names": trainable_names,
"trainable_params": trainable_params,
"total_params": total_params,
"trainable_fraction": trainable_params / total_params,
}
if is_main():
(out_dir / "train_config.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
print(json.dumps(meta, ensure_ascii=False, indent=2), flush=True)
log_path = out_dir / "train_log.jsonl"
start = time.time()
step = 0
accum_loss = 0.0
accum_tokens = 0
model.train()
optim.zero_grad(set_to_none=True)
initial_eval = evaluate(model, eval_loader, device, max_batches=args.eval_max_batches)
if is_main():
with log_path.open("a", encoding="utf-8") as f:
f.write(json.dumps({"event": "initial_eval", **initial_eval}, ensure_ascii=False) + "\n")
epoch = 0
while step < total_steps:
if train_sampler is not None:
train_sampler.set_epoch(epoch)
epoch += 1
for batch_idx, batch in enumerate(train_loader):
batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
out = model(**batch)
loss = out.loss / args.grad_accum
loss.backward()
ntok = int((batch["labels"] != -100).sum().item())
accum_loss += float(out.loss.item()) * ntok
accum_tokens += ntok
if (batch_idx + 1) % args.grad_accum == 0:
torch.nn.utils.clip_grad_norm_((p for p in model.parameters() if p.requires_grad), 1.0)
optim.step()
sched.step()
optim.zero_grad(set_to_none=True)
step += 1
if is_main() and (step == 1 or step % 10 == 0):
rec = {
"event": "train",
"step": step,
"loss": accum_loss / max(accum_tokens, 1),
"tokens": accum_tokens,
"lr": sched.get_last_lr()[0],
"elapsed_sec": time.time() - start,
}
print(json.dumps(rec, ensure_ascii=False), flush=True)
with log_path.open("a", encoding="utf-8") as f:
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
accum_loss = 0.0
accum_tokens = 0
if args.eval_steps and step % args.eval_steps == 0:
rec = {"event": "eval", "step": step, **evaluate(model, eval_loader, device, args.eval_max_batches)}
if is_main():
print(json.dumps(rec, ensure_ascii=False), flush=True)
with log_path.open("a", encoding="utf-8") as f:
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
if is_main() and args.save_steps and step % args.save_steps == 0:
ckpt = out_dir / f"checkpoint-{step}"
(model.module if isinstance(model, DDP) else model).save_pretrained(ckpt, safe_serialization=True)
tokenizer.save_pretrained(ckpt)
if step >= total_steps:
break
final_eval = evaluate(model, eval_loader, device, max_batches=args.eval_max_batches)
if is_main():
module = model.module if isinstance(model, DDP) else model
module.save_pretrained(out_dir, safe_serialization=True)
tokenizer.save_pretrained(out_dir)
(out_dir / "DONE").write_text("ok\n", encoding="utf-8")
with log_path.open("a", encoding="utf-8") as f:
f.write(json.dumps({"event": "final_eval", **final_eval}, ensure_ascii=False) + "\n")
print(json.dumps({"event": "final_eval", **final_eval}, ensure_ascii=False), flush=True)
if is_dist():
dist.destroy_process_group()
if __name__ == "__main__":
main()