LTA_openwebtext_dualt/scripts/eval_train8_overfit_ckpts.py

#!/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()