#!/usr/bin/env python3 from __future__ import annotations import argparse import csv import json import math import re import sys from pathlib import Path import numpy as np import torch import torch.nn.functional as F REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) from eval import build_model_from_ckpt from flowtext_lab.decode import state_for_model from flowtext_lab.tokenization import BpeTextTokenizer from train import make_bridge, masked_linear_soft_kl def parse_step(path: Path) -> int: if path.name == "latest.pt": return 10**18 m = re.search(r"step_(\d+)\.pt$", path.name) return int(m.group(1)) if m else 10**17 def load_gold(cache_dir: Path, max_len: int, limit: int) -> torch.Tensor: meta = json.loads((cache_dir / "meta.json").read_text()) n = int(meta["num_chunks"]) cache_len = int(meta["max_len"]) if cache_len != max_len: raise ValueError(f"cache max_len={cache_len} != requested {max_len}") arr = np.memmap(cache_dir / "chunks.i32.bin", mode="r", dtype=np.int32, shape=(n, max_len)) count = min(limit, n) return torch.from_numpy(np.asarray(arr[:count], dtype=np.int64).copy()) def masked_hard_ce(logits: torch.Tensor, ids: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: per = F.cross_entropy(logits.flatten(0, 1), ids.flatten(), reduction="none").view_as(ids) weight = mask.float() return (per * weight).sum() / weight.sum().clamp_min(1.0) @torch.no_grad() def eval_ckpt( ckpt_path: Path, tok: BpeTextTokenizer, gold_cpu: torch.Tensor, t_values: list[float], seeds: list[int], device: torch.device, ) -> list[dict[str, float | int | str]]: ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False, mmap=True) args = argparse.Namespace(**ckpt["args"]) max_len = int(gold_cpu.shape[1]) model = build_model_from_ckpt(ckpt, tok.vocab_size, max_len, device).eval() ids = gold_cpu.to(device) attn = torch.ones_like(ids, dtype=torch.bool, device=device) rows: list[dict[str, float | int | str]] = [] for t_value in t_values: for seed in seeds: torch.manual_seed(int(seed)) force_t = torch.full((ids.size(0),), float(t_value), dtype=torch.float32, device=device) bridge = make_bridge(args, ids, attn, tok.vocab_size, force_t=force_t) logits = model(state_for_model(model, bridge.state, float(args.eps)), force_t, bridge.attn_mask).float() mask = bridge.corrupt_mask & attn ce = masked_hard_ce(logits, bridge.ids, mask) if str(args.target_loss) == "linear_soft_kl": obj = masked_linear_soft_kl( logits, bridge.ids, mask, force_t, target_prob=float(args.target_prob), min_conf=float(args.linear_soft_target_min_conf), max_conf=float(args.linear_soft_target_max_conf), power=float(args.linear_soft_target_power), ) else: obj = ce pred = logits.argmax(dim=-1) acc = ((pred == bridge.ids) & mask).float().sum() / mask.float().sum().clamp_min(1.0) probs = F.softmax(logits.float(), dim=-1) maxp = probs.max(dim=-1).values[mask].mean() entropy = -(probs.clamp_min(1e-12) * probs.clamp_min(1e-12).log()).sum(dim=-1)[mask].mean() rows.append( { "checkpoint": str(ckpt_path), "step": int(ckpt.get("step", parse_step(ckpt_path))), "target_loss": str(args.target_loss), "clean_state_mode": str(args.clean_state_mode), "lowk": str(args.mask_mixture_lowk_clean_tokens), "t": float(t_value), "seed": int(seed), "objective": float(obj.detach().cpu()), "gold_ce": float(ce.detach().cpu()), "gold_acc": float(acc.detach().cpu()), "mean_maxp": float(maxp.detach().cpu()), "dist_entropy": float(entropy.detach().cpu()), "corrupt_frac": float(mask.float().mean().detach().cpu()), } ) del model, ckpt if device.type == "cuda": torch.cuda.empty_cache() return rows def summarize(rows: list[dict[str, float | int | str]], fit_acc: float, fit_objective: float) -> list[dict[str, float | int | str]]: grouped: dict[int, list[dict[str, float | int | str]]] = {} for row in rows: grouped.setdefault(int(row["step"]), []).append(row) out = [] for step in sorted(grouped): rs = grouped[step] rec = { "step": step, "target_loss": rs[0]["target_loss"], "clean_state_mode": rs[0]["clean_state_mode"], "lowk": rs[0]["lowk"], "objective_mean": sum(float(r["objective"]) for r in rs) / len(rs), "gold_ce_mean": sum(float(r["gold_ce"]) for r in rs) / len(rs), "gold_acc_mean": sum(float(r["gold_acc"]) for r in rs) / len(rs), "mean_maxp": sum(float(r["mean_maxp"]) for r in rs) / len(rs), "dist_entropy": sum(float(r["dist_entropy"]) for r in rs) / len(rs), "fit": 0, } rec["fit"] = int(float(rec["gold_acc_mean"]) >= fit_acc and float(rec["objective_mean"]) <= fit_objective) out.append(rec) return out def write_tsv(path: Path, rows: list[dict[str, float | int | str]]) -> None: path.parent.mkdir(parents=True, exist_ok=True) if not rows: path.write_text("", encoding="utf-8") return with path.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()), delimiter="\t") writer.writeheader() writer.writerows(rows) def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--run_dir", required=True) ap.add_argument("--cache_dir", required=True) ap.add_argument("--tokenizer_path", required=True) ap.add_argument("--out_dir", required=True) ap.add_argument("--max_len", type=int, default=1024) ap.add_argument("--limit", type=int, default=8) ap.add_argument("--t_values", default="0.125,0.25,0.5,0.75,1.0") ap.add_argument("--seeds", default="123,456,789") ap.add_argument("--fit_acc", type=float, default=0.95) ap.add_argument("--fit_objective", type=float, default=0.02) args = ap.parse_args() run_dir = Path(args.run_dir) out_dir = Path(args.out_dir) ckpts = sorted(run_dir.glob("step_*.pt"), key=parse_step) if not ckpts: raise FileNotFoundError(f"no step_*.pt checkpoints under {run_dir}") tok = BpeTextTokenizer.from_file(args.tokenizer_path) gold = load_gold(Path(args.cache_dir), args.max_len, args.limit) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") t_values = [float(x) for x in args.t_values.split(",") if x.strip()] seeds = [int(x) for x in args.seeds.split(",") if x.strip()] all_rows: list[dict[str, float | int | str]] = [] for ckpt in ckpts: rows = eval_ckpt(ckpt, tok, gold, t_values, seeds, device) all_rows.extend(rows) summary = summarize(rows, args.fit_acc, args.fit_objective)[0] print( f"{run_dir.name} step={summary['step']} obj={summary['objective_mean']:.4f} " f"ce={summary['gold_ce_mean']:.4f} acc={summary['gold_acc_mean']:.4f}", flush=True, ) summary_rows = summarize(all_rows, args.fit_acc, args.fit_objective) first_fit = next((r for r in summary_rows if int(r["fit"]) == 1), None) result = { "run_dir": str(run_dir), "first_fit_step": None if first_fit is None else int(first_fit["step"]), "fit_acc": args.fit_acc, "fit_objective": args.fit_objective, "best_acc": max(float(r["gold_acc_mean"]) for r in summary_rows), "best_objective": min(float(r["objective_mean"]) for r in summary_rows), "last": summary_rows[-1], } out_dir.mkdir(parents=True, exist_ok=True) write_tsv(out_dir / "per_t.tsv", all_rows) write_tsv(out_dir / "summary.tsv", summary_rows) (out_dir / "result.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8") print(json.dumps(result, ensure_ascii=False, indent=2), flush=True) if __name__ == "__main__": main()