File size: 5,046 Bytes
17a55cc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | """Evaluation of a Groot checkpoint on decision rows and public benchmarks.
groot-eval --model runs/groot-1-s/final --rows data/test.jsonl [--max-rows 2000] [--out eval.json]
groot-eval --model ... --bench typed_decisions # LocalLLaMA/typed-decisions test split
Reports accuracy per type, NLL, Brier, ECE (with and without the stored
temperatures), coverage/accuracy at confidence thresholds, and per-source
breakdowns. Rows are the JSONL format of data.py.
"""
from __future__ import annotations
import argparse
import json
from collections import defaultdict
from pathlib import Path
import numpy as np
import torch
from .calibration import Record, evaluate_records, softmax
from .config import QTYPES
from .data import CONVERTERS, read_jsonl
from .model import GrootDecisionModel
from .schema import confidence_choice, confidence_noul, confidence_score
from .serialize import Serializer
from .train import evaluate
def confidence_of(probabilities: list[float], qtype: str) -> float:
if qtype == "choice":
return confidence_choice(probabilities)
if qtype == "score":
return confidence_score(probabilities)
return confidence_noul(probabilities[1])
def coverage_table(records: list[Record], temps: list[float], thresholds=(0.3, 0.5, 0.7, 0.9)) -> dict:
"""Accuracy of answers whose System One confidence passes each threshold."""
rows = []
for r in records:
p = softmax(r.logits, temps[r.qtype]).tolist()
c = confidence_of(p, QTYPES[r.qtype])
rows.append((c, int(np.argmax(p) == int(r.target.argmax()))))
out = {}
for th in thresholds:
kept = [a for c, a in rows if c >= th]
out[str(th)] = {"coverage": len(kept) / max(1, len(rows)), "accuracy": float(np.mean(kept)) if kept else None}
return out
def per_source(rows: list[dict], records: list[Record], temps: list[float]) -> dict:
groups: dict[str, list[Record]] = defaultdict(list)
for row, rec in zip(rows, records):
groups[row.get("source", "?")].append(rec)
return {src: evaluate_records(recs, temps).to_dict() for src, recs in groups.items()}
def main(argv: list[str] | None = None) -> None:
ap = argparse.ArgumentParser(prog="groot-eval")
ap.add_argument("--model", required=True)
ap.add_argument("--rows", default=None)
ap.add_argument("--bench", default=None, help="name of a converter in data.py (uses its test split)")
ap.add_argument("--max-rows", type=int, default=None)
ap.add_argument("--batch-size", type=int, default=16)
ap.add_argument("--device", default=None)
ap.add_argument("--out", default=None)
ap.add_argument("--set", action="append", default=[], help="config override key=value (e.g. state_retrieval=true)")
ap.add_argument("--calibrate-on", default=None, help="fit per-type temperatures on this JSONL before evaluating")
ap.add_argument("--save-calibration", action="store_true", help="write the fitted temperatures into the checkpoint")
args = ap.parse_args(argv)
from transformers import AutoTokenizer
device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"))
model = GrootDecisionModel.from_pretrained(args.model, device=device)
tok = AutoTokenizer.from_pretrained(Path(args.model) / "tokenizer")
if args.set:
from .config import GrootConfig, parse_overrides
model.cfg = GrootConfig.from_dict({**model.cfg.to_dict(), **parse_overrides(args.set)})
ser = Serializer(tok, model.cfg)
if args.calibrate_on:
from .calibration import fit_temperatures
cal_rows = list(read_jsonl(args.calibrate_on))
cal_records, _ = evaluate(model, cal_rows, ser, device, args.batch_size, None)
temps = fit_temperatures(cal_records, n_types=len(model.cfg.qtypes))
model.temperature.copy_(torch.tensor(temps, dtype=model.temperature.dtype, device=model.temperature.device))
print(json.dumps({"fitted_temperatures": temps, "on": args.calibrate_on, "rows": len(cal_rows)}), flush=True)
if args.save_calibration:
(Path(args.model) / "calibration.json").write_text(json.dumps({"temperature": temps}, indent=2) + "\n")
if args.rows:
rows = list(read_jsonl(args.rows, args.max_rows))
elif args.bench:
rows = [r for r in CONVERTERS[args.bench](limit=args.max_rows) if r["split"] == "test"]
else:
raise SystemExit("pass --rows or --bench")
records, raw = evaluate(model, rows, ser, device, args.batch_size, None)
temps = model.temperature.tolist()
report = {
"model": args.model,
"rows": len(rows),
"raw": raw,
"calibrated": evaluate_records(records, temps).to_dict(),
"coverage": coverage_table(records, temps),
"per_source": per_source(rows, records, temps),
}
text = json.dumps(report, indent=2)
print(text)
if args.out:
Path(args.out).write_text(text + "\n")
if __name__ == "__main__":
main()
|