"""Part 3: benchmark decoding strategies on ROCStories with Pythia-160M. python -m src.part3.evaluate --n_samples 1000 Metric definitions (free-form generation has no single "correct" answer, so these are stated explicitly rather than assumed): * **Reference perplexity** -- the model's perplexity on the *gold* continuation under teacher forcing. It measures the model, not the decoder, so it is identical for every strategy and acts as a control. * **Generation perplexity** -- the model's perplexity on *its own* output. This does vary by strategy and quantifies how "safe" each one plays: greedy sits in the high-probability region by construction, sampling deliberately leaves it. * **Precision / recall / F1** -- unigram overlap against the reference, computed as a multiset intersection and micro-averaged over the corpus. Precision is matched/|generated|, recall is matched/|reference|. * **Accuracy** -- position-wise next-token agreement with the reference over the overlapping prefix, plus the exact-match rate (almost always ~0, and reported to make the point that it is the wrong metric here). * **Distinct-1/2 and repetition rate** -- diversity, which is the axis the overlap metrics cannot see and the reason sampling exists at all. """ from __future__ import annotations import argparse import time from collections import Counter import numpy as np import torch import torch.nn.functional as F from ..utils import RUNS, save_json, set_seed from .decoding import beam_search, greedy_decode, sample_decode MODEL_ID = "EleutherAI/pythia-160m" DATASET = "hamishivi/ROCStories" # --------------------------------------------------------------------------- # Metrics # --------------------------------------------------------------------------- def overlap_prf(hyp_ids, ref_ids): """Micro-averaged unigram precision / recall / F1 over the corpus.""" tp = hyp_n = ref_n = 0 for h, r in zip(hyp_ids, ref_ids): ch, cr = Counter(h), Counter(r) tp += sum((ch & cr).values()) hyp_n += len(h) ref_n += len(r) p = tp / max(hyp_n, 1) r = tp / max(ref_n, 1) f1 = 2 * p * r / max(p + r, 1e-9) return p, r, f1 def positional_accuracy(hyp_ids, ref_ids): """Token agreement at matching positions, plus the exact-match rate.""" match = total = exact = 0 for h, r in zip(hyp_ids, ref_ids): n = min(len(h), len(r)) match += sum(1 for i in range(n) if h[i] == r[i]) total += max(len(r), 1) exact += int(h == r) return match / max(total, 1), exact / max(len(hyp_ids), 1) def diversity(hyp_ids): """distinct-1, distinct-2 and the share of tokens that repeat a bigram.""" u1, u2, n1, n2 = set(), set(), 0, 0 rep = tot = 0 for h in hyp_ids: u1.update(h) n1 += len(h) bg = list(zip(h, h[1:])) u2.update(bg) n2 += len(bg) seen = set() for g in bg: if g in seen: rep += 1 seen.add(g) tot += len(bg) return (len(u1) / max(n1, 1), len(u2) / max(n2, 1), rep / max(tot, 1)) @torch.no_grad() def sequence_nll(model, prompt_ids, prompt_mask, cont_ids, device): """Mean NLL per token of ``cont_ids`` given the prompt (teacher forcing).""" tot_nll, tot_tok = 0.0, 0 for i in range(len(cont_ids)): cont = cont_ids[i] if len(cont) == 0: continue p = prompt_ids[i][prompt_mask[i].bool()] seq = torch.cat([p, torch.tensor(cont, device=device)]).unsqueeze(0) logits = model(input_ids=seq).logits.float() tgt = seq[:, 1:] lp = F.log_softmax(logits[:, :-1], dim=-1) nll = -lp.gather(2, tgt.unsqueeze(-1)).squeeze(-1) tot_nll += nll[0, len(p) - 1:].sum().item() tot_tok += len(cont) return tot_nll / max(tot_tok, 1) # --------------------------------------------------------------------------- # Runner # --------------------------------------------------------------------------- def build_batches(tok, rows, batch_size, device, max_prompt=128): """Left-padded prompt batches plus the tokenised references.""" batches = [] for s in range(0, len(rows), batch_size): chunk = rows[s:s + batch_size] enc = [tok(r["prompt"], truncation=True, max_length=max_prompt)["input_ids"] for r in chunk] refs = [tok(r["continuation"], truncation=True, max_length=128)["input_ids"] for r in chunk] L = max(len(e) for e in enc) pad = tok.pad_token_id or 0 ids = torch.full((len(enc), L), pad, dtype=torch.long) mask = torch.zeros((len(enc), L), dtype=torch.long) for i, e in enumerate(enc): # LEFT padding ids[i, L - len(e):] = torch.tensor(e) mask[i, L - len(e):] = 1 batches.append((ids.to(device), mask.to(device), refs)) return batches def run_strategy(name, fn, batches, model, device, max_new, eos_id): hyps, refs_all, scores = [], [], [] t0 = time.perf_counter() for ids, mask, refs in batches: out, sc = fn(model, ids, mask, max_new=max_new, eos_id=eos_id) hyps.extend(out) refs_all.extend(refs) scores.extend(sc) torch.cuda.synchronize() if device.type == "cuda" else None elapsed = time.perf_counter() - t0 p, r, f1 = overlap_prf(hyps, refs_all) acc, exact = positional_accuracy(hyps, refs_all) d1, d2, rep = diversity(hyps) gen_nll = -np.mean([s / max(len(h), 1) for s, h in zip(scores, hyps) if len(h) > 0]) return { "strategy": name, "precision": p, "recall": r, "f1": f1, "accuracy": acc, "exact_match": exact, "gen_ppl": float(np.exp(min(gen_nll, 20))), "distinct_1": d1, "distinct_2": d2, "repetition_rate": rep, "mean_len": float(np.mean([len(h) for h in hyps])), "seconds": elapsed, "sec_per_sample": elapsed / max(len(hyps), 1), }, hyps, refs_all def main(): ap = argparse.ArgumentParser() ap.add_argument("--n_samples", type=int, default=1000) ap.add_argument("--batch_size", type=int, default=32) ap.add_argument("--max_new", type=int, default=48) ap.add_argument("--split", default="test") ap.add_argument("--seed", type=int, default=1337) args = ap.parse_args() from datasets import load_dataset from transformers import AutoModelForCausalLM, AutoTokenizer set_seed(args.seed) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") tok = AutoTokenizer.from_pretrained(MODEL_ID) if tok.pad_token is None: tok.pad_token = tok.eos_token model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=torch.float32).to(device).eval() ds = load_dataset(DATASET, split=args.split) rows = [ds[i] for i in range(min(args.n_samples, len(ds)))] print(f"[part3] {len(rows)} samples from {DATASET}:{args.split}", flush=True) batches = build_batches(tok, rows, args.batch_size, device) eos_id = tok.eos_token_id # Control: perplexity of the gold continuation. Decoder-independent. nll = np.mean([sequence_nll(model, ids, mask, refs, device) for ids, mask, refs in batches[:8]]) ref_ppl = float(np.exp(nll)) print(f"[part3] reference perplexity (teacher forcing) = {ref_ppl:.2f}", flush=True) import functools strategies = [ ("greedy", greedy_decode), ("top_k_10", functools.partial(sample_decode, top_k=10)), ("top_k_50", functools.partial(sample_decode, top_k=50)), ("top_p_0.90", functools.partial(sample_decode, top_p=0.90)), ("top_p_0.95", functools.partial(sample_decode, top_p=0.95)), ("beam_1", functools.partial(beam_search, beam_width=1)), ("beam_2", functools.partial(beam_search, beam_width=2)), ("beam_4", functools.partial(beam_search, beam_width=4)), ] results, examples = [], {} for name, fn in strategies: res, hyps, refs = run_strategy(name, fn, batches, model, device, args.max_new, eos_id) res["ref_ppl"] = ref_ppl results.append(res) examples[name] = [tok.decode(h) for h in hyps[:5]] print(f" {name:11s} P {res['precision']:.3f} R {res['recall']:.3f} " f"F1 {res['f1']:.3f} acc {res['accuracy']:.3f} " f"ppl {res['gen_ppl']:7.2f} d2 {res['distinct_2']:.3f} " f"rep {res['repetition_rate']:.3f} {res['seconds']:6.1f}s", flush=True) out = RUNS / "part3" out.mkdir(parents=True, exist_ok=True) save_json({ "model": MODEL_ID, "dataset": DATASET, "split": args.split, "n_samples": len(rows), "max_new": args.max_new, "batch_size": args.batch_size, "reference_ppl": ref_ppl, "results": results, "prompts": [r["prompt"] for r in rows[:5]], "references": [r["continuation"] for r in rows[:5]], "examples": examples, }, out / "metrics.json") print(f"[part3] wrote {out/'metrics.json'}") if __name__ == "__main__": main()