Clean tokenizer swap migration

evaluation_reporting/README.md

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

evaluation_reporting/eval_tokenizer_swap_benchmark.py

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

evaluation_reporting/merge_tokenizer_swap_benchmark_shards.py

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

evaluation_reporting/run_public_heldout_eval_8gpu.sh

@@ -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"

evaluation_reporting/summarize_heldout_capabilities.py

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