1"""Loading, tokenising and batching the synthetic task splits (char-level or BPE tokenizer)."""2from __future__ import annotations34import json5from collections import defaultdict6from pathlib import Path78import numpy as np9import torch1011from data.char_tokenizer import CharTokenizer121314class AnyTokenizer:15 def __init__(self, spec: str):16 self.kind = "char" if spec == "char" else "bpe"17 if self.kind == "char":18 self.tok = CharTokenizer(); self.eot_id = self.tok.eot_id; self.pad_id = self.tok.pad_id19 self.vocab_size = self.tok.vocab_size20 else:21 from tokenizers import Tokenizer22 self.tok = Tokenizer.from_file(spec); self.eot_id = self.tok.token_to_id("<|endoftext|>")23 self.pad_id = self.tok.token_to_id("<|pad|>"); self.vocab_size = self.tok.get_vocab_size()2425 def encode(self, s: str) -> list[int]:26 return self.tok.encode(s) if self.kind == "char" else self.tok.encode(s).ids2728 def decode(self, ids) -> str:29 ids = [int(i) for i in ids]30 return self.tok.decode(ids) if self.kind == "char" else self.tok.decode(ids, skip_special_tokens=True)313233def load_items(path: str | Path) -> list[dict]:34 return [json.loads(l) for l in open(path)]353637def encode_item(tok: AnyTokenizer, it: dict) -> dict:38 p = tok.encode(it["prompt"]); a = tok.encode(it["answer"]) + [tok.eot_id]39 return {"ids": p + a, "prompt_len": len(p), "answer_len": len(a), "item": it}404142class SyntheticBatches:43 """Right-padded batches with a loss mask over the answer span (+EOT). Targets elsewhere are -100."""4445 def __init__(self, items: list[dict], tok: AnyTokenizer, max_len: int = 256):46 self.enc = [encode_item(tok, it) for it in items]47 self.enc = [e for e in self.enc if len(e["ids"]) <= max_len]48 self.pad = tok.pad_id4950 def __len__(self):51 return len(self.enc)5253 def collate(self, encs, device, loss_span: str = "answer"):54 L = max(len(e["ids"]) for e in encs)55 x = torch.full((len(encs), L), self.pad, dtype=torch.long)56 y = torch.full((len(encs), L), -100, dtype=torch.long)57 for i, e in enumerate(encs):58 ids = torch.tensor(e["ids"])59 x[i, :len(ids)] = ids60 if loss_span == "answer":61 s = e["prompt_len"] - 1 # position predicting the first answer token62 y[i, s:len(ids) - 1] = ids[s + 1:]63 else:64 y[i, :len(ids) - 1] = ids[1:]65 return x.to(device), y.to(device)6667 def random_batches(self, batch_size: int, seed: int, device, loss_span="answer"):68 rng = np.random.default_rng(seed)69 while True:70 idx = rng.integers(0, len(self.enc), size=batch_size)71 yield self.collate([self.enc[i] for i in idx], device, loss_span)7273 def sequential_batches(self, batch_size: int, device, loss_span="answer", max_items=None):74 encs = self.subset(max_items)75 for i in range(0, len(encs), batch_size):76 yield self.collate(encs[i:i + batch_size], device, loss_span)7778 def subset(self, max_items=None):79 """Stratified deterministic subset: every k-th item so that all families/templates/depths are covered."""80 if not max_items or max_items >= len(self.enc):81 return self.enc82 k = len(self.enc) / max_items83 return [self.enc[int(i * k)] for i in range(max_items)]8485 def prompt_buckets(self, batch_size: int, max_items=None):86 """Groups of encoded items with identical prompt length (for cache-based free generation)."""87 encs = self.subset(max_items)88 by_len = defaultdict(list)89 for e in encs:90 by_len[e["prompt_len"]].append(e)91 for L, group in sorted(by_len.items()):92 for i in range(0, len(group), batch_size):93 yield group[i:i + batch_size]949596@torch.no_grad()97def exact_match(model, batches: SyntheticBatches, tok: AnyTokenizer, device, max_budget: int, cost=None,98 forced_actions_fn=None, batch_size: int = 64, max_items=None, max_new=12, lambda_cost=None, kappa_risk=None):99 """Free greedy generation of the answer; returns dict with accuracy, per-family accuracy, and per-item records."""100 from prisme.generation import generate101 model.eval()102 recs = []103 for group in batches.prompt_buckets(batch_size, max_items):104 prompt = torch.tensor([e["ids"][:e["prompt_len"]] for e in group], device=device)105 out = generate(model, prompt, max_new, max_budget=max_budget, cost=cost, temperature=0.0, eos_id=tok.eot_id,106 forced_budget_actions=forced_actions_fn, lambda_cost=lambda_cost, kappa_risk=kappa_risk)107 gen = out["tokens"][:, prompt.shape[1]:]108 ne = out["n_edits"]109 for i, e in enumerate(group):110 ids = gen[i].tolist()111 if tok.eot_id in ids:112 ids = ids[:ids.index(tok.eot_id)]113 pred = tok.decode(ids).strip()114 truth = e["item"]["answer"].strip()115 recs.append({"uid": e["item"]["uid"], "family": e["item"]["family"], "template": e["item"]["template"],116 "depth": e["item"]["depth"], "pred": pred, "truth": truth, "correct": pred == truth,117 "edits_per_token": float(ne[i].float().mean()) if ne is not None else 0.0})118 acc = float(np.mean([r["correct"] for r in recs])) if recs else float("nan")119 fam = defaultdict(list)120 for r in recs:121 fam[r["family"]].append(r["correct"])122 return {"accuracy": acc, "n": len(recs), "by_family": {k: float(np.mean(v)) for k, v in fam.items()},123 "edits_per_token": float(np.mean([r["edits_per_token"] for r in recs])) if recs else 0.0, "records": recs}124