Blame for src/part3/evaluate.py

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