| """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" |
|
|
|
|
| |
| |
| |
| 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) |
|
|
|
|
| |
| |
| |
| 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): |
| 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 |
|
|
| |
| 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() |
|
|
