Add MP-01 (9500) 100M base: weights, config, tokenizer, model code, model card and eval results
11e2c7c verified | """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 | |
| 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 | |