File size: 3,985 Bytes
11e2c7c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """Zero-shot multiple-choice evaluation (lm-evaluation-harness conventions) that fits next to a training job."""
from __future__ import annotations
import json
import random
import re
from pathlib import Path
import torch
EVAL_DIR = Path(__file__).resolve().parent.parent / "data" / "evals"
# chance level, for the dashboard
CHANCE = {"hellaswag": 25.0, "arc_easy": 25.0, "arc_challenge": 25.0, "piqa": 50.0, "sciq": 25.0, "openbookqa": 25.0, "winogrande": 50.0}
METRIC = {"hellaswag": "acc_norm", "arc_easy": "acc", "arc_challenge": "acc_norm", "piqa": "acc_norm", "sciq": "acc_norm", "openbookqa": "acc_norm", "winogrande": "acc"}
def load_tasks(limit: int | None = None, seed: int = 0) -> dict[str, list[dict]]:
tasks = {}
for f in sorted(EVAL_DIR.glob("*.json")):
items = json.loads(f.read_text())
if limit and len(items) > limit:
items = random.Random(seed).sample(items, limit)
tasks[f.stem] = items
return tasks
@torch.no_grad()
def loglik_batch(model, pairs, device, pad_id=1, max_ctx=480, bs=48):
"""pairs: list of (ctx_ids, cont_ids). Returns (sum logprob, n_tokens) per pair."""
res = [None] * len(pairs)
order = sorted(range(len(pairs)), key=lambda i: len(pairs[i][0]) + len(pairs[i][1]))
w = model.out_weight
for s in range(0, len(order), bs):
idxs = order[s : s + bs]
seqs = []
for i in idxs:
c, k = pairs[i]
c = c[-max_ctx:]
seqs.append((c + k, len(c), len(k)))
L = max(len(x[0]) for x in seqs)
inp = torch.full((len(seqs), L), pad_id, dtype=torch.long)
for r, (ids, _, _) in enumerate(seqs):
inp[r, : len(ids)] = torch.tensor(ids)
inp = inp.to(device)
with torch.autocast(device.type, torch.bfloat16):
h = model.hidden(inp)
rows, cols, tgts = [], [], []
for r, (ids, nc, nk) in enumerate(seqs):
for j in range(nk):
rows.append(r); cols.append(nc + j - 1); tgts.append(ids[nc + j])
hs = h[torch.tensor(rows, device=device), torch.tensor(cols, device=device)]
logits = (hs @ w.to(hs.dtype).t()).float()
lp = torch.log_softmax(logits, -1).gather(1, torch.tensor(tgts, device=device)[:, None]).squeeze(1).cpu()
p = 0
for r, i in enumerate(idxs):
nk = seqs[r][2]
res[i] = (float(lp[p : p + nk].sum()), nk); p += nk
return res
def evaluate(model, tok, tasks: dict[str, list[dict]], device, pad_id=1, bs=48) -> dict[str, dict]:
model.eval()
out = {}
for name, items in tasks.items():
pairs, spans, labels, chs = [], [], [], []
for it in items:
start = len(pairs)
if "alt_ctx" in it: # WinoGrande partial scoring: same continuation, two candidate contexts
tail = tok.encode(it["choices"][0]).ids
pairs += [(tok.encode(it["ctx"]).ids, tail), (tok.encode(it["alt_ctx"]).ids, tail)]
labels.append(it["wino_label"]); chs.append([it["choices"][0]] * 2)
else:
cids = tok.encode(it["ctx"]).ids
for ch in it["choices"]:
pairs.append((cids, tok.encode(ch).ids or [0]))
labels.append(it["label"]); chs.append(it["choices"])
spans.append((start, len(pairs) - start))
ll = loglik_batch(model, pairs, device, pad_id=pad_id, bs=bs)
acc = accn = 0
for (st, n), lab, ch in zip(spans, labels, chs):
raw = [ll[st + k][0] for k in range(n)]
norm = [raw[k] / max(1, len(ch[k].encode())) for k in range(n)]
acc += int(max(range(n), key=raw.__getitem__) == lab)
accn += int(max(range(n), key=norm.__getitem__) == lab)
out[name] = {"acc": 100 * acc / len(items), "acc_norm": 100 * accn / len(items), "n": len(items)}
out[name]["score"] = out[name][METRIC.get(name, "acc")]
model.train()
return out
|