Clean tokenizer swap migration

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