spb/wpai0001_uqo

1"""Exact-match evaluation of a toy model on a Split under greedy decoding, with per-item outputs."""2from __future__ import annotations34from dataclasses import dataclass56import numpy as np7import torch89from wp_ai_0001.datasets.common import Split, Vocab, encode_split10from wp_ai_0001.models.decoding import greedy_decode11from wp_ai_0001.models.toy_transformer import ToyTransformer121314@dataclass15class EvalResult:16    split: str17    split_hash: str18    n: int19    accuracy: float  # final-answer exact match20    trace_accuracy: float  # all scratchpad tokens exact match21    correct: np.ndarray  # [N] bool (final answer)22    trace_correct: np.ndarray  # [N] bool23    clusters: list[str]24    predictions: list[list[str]]25    answers: list[list[str]]2627    def to_summary(self) -> dict:28        return {"split": self.split, "split_hash": self.split_hash, "n": self.n, "accuracy": self.accuracy, "trace_accuracy": self.trace_accuracy}293031@torch.no_grad()32def evaluate(model: ToyTransformer, split: Split, vocab: Vocab, device: torch.device, bs: int = 1024) -> EvalResult:33    enc = encode_split(split, vocab, model.cfg.max_seq_len)34    K = enc["answers"].shape[1]35    preds = []36    for i in range(0, len(enc["ids"]), bs):37        ids = torch.from_numpy(enc["ids"][i : i + bs]).to(device)38        pl = torch.from_numpy(enc["prompt_len"][i : i + bs]).to(device)39        na = torch.from_numpy(enc["n_answer"][i : i + bs]).to(device)40        preds.append(greedy_decode(model, ids, pl, na, K).cpu().numpy())41    pred = np.concatenate(preds)42    na = enc["n_answer"]43    ar = np.arange(len(na))44    final_ok = pred[ar, na - 1] == enc["answers"][ar, na - 1]45    valid = np.arange(K)[None] < na[:, None]46    trace_ok = ((pred == enc["answers"]) | ~valid).all(-1)47    return EvalResult(split.name, split.hash, len(na), float(final_ok.mean()), float(trace_ok.mean()), final_ok, trace_ok,48                      [e.cluster for e in split.examples], [vocab.decode(pred[i, : na[i]]) for i in ar],49                      [list(e.answers) for e in split.examples])505152def chance_level(split: Split) -> float:53    """Majority-class rate of the answer distribution (the relevant chance baseline)."""54    return max(split.answer_distribution().values())55