| """Evaluate one arm (or Gemma 3 270M as the reference) and save EVERY per-question result, so error bars can be
|
| computed afterwards by bars.py without re-running anything.
|
|
|
| python eval.py --arm B --ckpt ckpt/B/step_07820.pt [--limit N]
|
| python eval.py --gemma [--limit N] # google/gemma-3-270m, the shape our body copies
|
|
|
| Benchmarks (EleutherAI lm-evaluation-harness, standard zero-shot settings): lambada_openai (acc), hellaswag (acc_norm),
|
| arc_easy (acc_norm), piqa (acc_norm). Our arms are wrapped as an lm-eval model: every request is scored by the
|
| log-probability of the continuation given the context, exactly as the harness does for Hugging Face models. Documents
|
| in our training data are separated by <eos> (no <bos>), so every context starts with <eos> for our arms; Gemma
|
| 3 270M was trained with <bos>, so its contexts start with <bos>.
|
| Validation loss: mean next-token loss per 1,024-token sequence of the held-out split (data/val.npy), one number per
|
| sequence, so its error bar is over sequences.
|
|
|
| Output: reports/eval_<name>.json = {summary, per_question: {task: [0/1 ...]}, val_loss_per_seq: [...], meta}
|
| """
|
| import argparse, json, math, os, pathlib, time
|
| import numpy as np
|
| import torch
|
| import torch.nn.functional as F
|
|
|
| TASKS = {"lambada_openai": "acc", "hellaswag": "acc_norm", "arc_easy": "acc_norm", "piqa": "acc_norm"}
|
| SEQ = 1024
|
|
|
|
|
| def load_model(a):
|
| from transformers import AutoTokenizer
|
| tok = AutoTokenizer.from_pretrained("google/gemma-3-270m")
|
| if a.gemma:
|
| from transformers import AutoModelForCausalLM
|
| g = AutoModelForCausalLM.from_pretrained("google/gemma-3-270m", torch_dtype=torch.float32).cuda().eval()
|
| def logits_fn(ids):
|
| return g(input_ids=ids, use_cache=False).logits
|
| return "gemma3_270m", tok, logits_fn
|
| import train
|
| counts = np.load(pathlib.Path(a.data) / "unigram_counts.npy")
|
| m = train.Model(a.arm, counts).cuda().eval()
|
| st = torch.load(a.ckpt, map_location="cuda", weights_only=False)
|
| m.load_state_dict(st["model"])
|
| def logits_fn(ids):
|
| return m.logits(m.hidden(ids))
|
| return f"arm{a.arm}_step{st.get('step', 0)}", tok, logits_fn
|
|
|
|
|
| def make_lm(tok, logits_fn, bs, sep):
|
| from lm_eval.api.model import LM
|
|
|
| class Wrapped(LM):
|
| def __init__(self):
|
| super().__init__()
|
|
|
| @torch.no_grad()
|
| def _score(self, ctx_ids, cont_ids):
|
| ids = (ctx_ids + cont_ids)[-SEQ:]
|
| n = len(cont_ids)
|
| x = torch.tensor([ids], device="cuda")
|
| with torch.autocast("cuda", dtype=torch.bfloat16):
|
| lg = logits_fn(x[:, :-1]).float()
|
| lp = F.log_softmax(lg[0, -n:], -1)
|
| tgt = x[0, -n:]
|
| ll = float(lp.gather(-1, tgt[:, None]).sum())
|
| greedy = bool((lp.argmax(-1) == tgt).all())
|
| return ll, greedy
|
|
|
| def loglikelihood(self, requests):
|
| out = []
|
| for r in requests:
|
| ctx, cont = r.args
|
| ctx_ids = [sep] + tok(ctx, add_special_tokens=False).input_ids
|
| whole = tok(ctx + cont, add_special_tokens=False).input_ids
|
| cont_ids = whole[len(ctx_ids) - 1:] if whole[:len(ctx_ids) - 1] == ctx_ids[1:] else tok(cont, add_special_tokens=False).input_ids
|
| out.append(self._score(ctx_ids, cont_ids))
|
| return out
|
|
|
| def loglikelihood_rolling(self, requests):
|
| raise NotImplementedError
|
|
|
| def generate_until(self, requests):
|
| raise NotImplementedError
|
|
|
| return Wrapped()
|
|
|
|
|
| @torch.no_grad()
|
| def val_losses(logits_fn, val, n_seq):
|
| x = torch.from_numpy(val[:n_seq * (SEQ + 1)].astype(np.int64).reshape(n_seq, SEQ + 1))
|
| out = []
|
| for i in range(n_seq):
|
| ids = x[i:i + 1].cuda()
|
| with torch.autocast("cuda", dtype=torch.bfloat16):
|
| lg = logits_fn(ids[:, :-1]).float()
|
| out.append(float(F.cross_entropy(lg[0], ids[0, 1:])))
|
| return out
|
|
|
|
|
| def main():
|
| ap = argparse.ArgumentParser()
|
| ap.add_argument("--arm", choices=["A", "B", "C", "D"])
|
| ap.add_argument("--ckpt")
|
| ap.add_argument("--gemma", action="store_true")
|
| ap.add_argument("--data", default="data")
|
| ap.add_argument("--limit", type=int, default=0, help="questions per task (0 = all); for smoke tests only")
|
| ap.add_argument("--val-seqs", type=int, default=512)
|
| ap.add_argument("--name", default="")
|
| a = ap.parse_args()
|
| t0 = time.time()
|
| name, tok, logits_fn = load_model(a)
|
| name = a.name or name
|
| import lm_eval
|
| res = lm_eval.simple_evaluate(model=make_lm(tok, logits_fn, 1, tok.bos_token_id if a.gemma else tok.eos_token_id), tasks=list(TASKS), num_fewshot=0,
|
| limit=a.limit or None, log_samples=True, bootstrap_iters=0)
|
| per_q, summary = {}, {}
|
| for task, metric in TASKS.items():
|
| samples = sorted(res["samples"][task], key=lambda s: s["doc_id"])
|
| per_q[task] = [float(s[metric]) for s in samples]
|
| summary[task] = {"metric": metric, "n": len(per_q[task]), "score": float(np.mean(per_q[task]))}
|
| val = np.load(pathlib.Path(a.data) / "val.npy", mmap_mode="r")
|
| vl = val_losses(logits_fn, val, a.val_seqs)
|
| summary["val_loss"] = {"metric": "nats/token", "n": len(vl), "score": float(np.mean(vl))}
|
| out = {"name": name, "summary": summary, "per_question": per_q, "val_loss_per_seq": vl,
|
| "meta": {"ckpt": a.ckpt, "gemma": a.gemma, "limit": a.limit, "seconds": round(time.time() - t0),
|
| "lm_eval": lm_eval.__version__, "date": time.strftime("%Y-%m-%d %H:%M")}}
|
| pathlib.Path("reports").mkdir(exist_ok=True)
|
| p = pathlib.Path("reports") / f"eval_{name}.json"
|
| json.dump(out, open(p, "w"))
|
| print(json.dumps(summary, indent=1))
|
| print("saved", p)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
