openjev / code /eval.py
AlexWortega's picture
openjev v2: results, videos (Doom v2, Minecraft), code, card
0988d66 verified
Raw
History Blame Contribute Delete
22.9 kB
#!/usr/bin/env python
"""Evaluate NLI cross-encoders (ours + dleemiller/ModernCE-large-nli) on MNLI, GPQA-diamond and GSM8K.
Modes follow https://huggingface.co/blog/dleemiller/nli-xenc-ways-to-use:
* QA rerank (#3): premise=question, hypothesis=candidate; pick argmax P(entailment)
* grading (#6): premise=question+reference answer, hypothesis=candidate; entailment <=> correct
Usage:
python eval.py --models ckpt/qwen3.5-0.8b-nli dleemiller/ModernCE-large-nli --out results/qwen0.8b.json
"""
import argparse
import collections
import json
import os
import re
import time
import numpy as np
import requests
import torch
from datasets import load_dataset
from sklearn.metrics import f1_score
from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer
CON, ENT, NEU = 0, 1, 2
# Some HF configs mislabel the logit order. ModernCE's config.json says {0 ent, 1 neu, 2 con} but the actual
# logits are [con, ent, neu] (as the model card says; verified: trusting config gives 3% MNLI accuracy).
LABEL_ORDER_OVERRIDES = {"dleemiller/ModernCE-large-nli": ["contradiction", "entailment", "neutral"]}
LLAMA_URL = "http://127.0.0.1:18085/v1/chat/completions"
class NLIScorer:
"""Returns softmax probs [contradiction, entailment, neutral] for (premise, hypothesis) pairs."""
def __init__(self, path, device="cuda", bs=32, max_len=1024):
self.cfg = AutoConfig.from_pretrained(path)
self.tok = AutoTokenizer.from_pretrained(path)
cls = AutoModelForSequenceClassification
if getattr(self.cfg, "model_type", "") == "qwen3_5_moe":
from modeling_qwen35_moe_seqcls import Qwen3_5MoeForSequenceClassification as cls
self.model = cls.from_pretrained(path, dtype=torch.bfloat16).to(device).eval()
self.device, self.bs, self.max_len = device, bs, max_len
self.template = getattr(self.cfg, "nli_template", None) # set by train.py for Qwen models
if self.tok.pad_token is None:
self.tok.pad_token = self.tok.eos_token
if self.template:
self.tok.padding_side = "right"
self.model.config.get_text_config().pad_token_id = self.tok.pad_token_id
# models differ in label order (ModernCE config: 0 ent, 1 neu, 2 con) -> permute to [con, ent, neu]
l2i = {v.lower(): int(k) for k, v in self.cfg.id2label.items()}
if path in LABEL_ORDER_OVERRIDES:
l2i = {name: i for i, name in enumerate(LABEL_ORDER_OVERRIDES[path])}
self.perm = [l2i["contradiction"], l2i["entailment"], l2i["neutral"]]
print(f"{path}: id2label={self.cfg.id2label} perm={self.perm} template={'yes' if self.template else 'no'}")
@torch.no_grad()
def predict(self, pairs):
out = []
for i in range(0, len(pairs), self.bs):
chunk = pairs[i : i + self.bs]
if self.template:
texts = [self.template.format(premise=p.strip(), hypothesis=h.strip()) for p, h in chunk]
enc = self.tok(texts, truncation=True, max_length=self.max_len, padding=True, return_tensors="pt")
else:
enc = self.tok([p for p, _ in chunk], [h for _, h in chunk], truncation=True,
max_length=self.max_len, padding=True, return_tensors="pt")
enc = {k: v.to(self.device) for k, v in enc.items()}
logits = self.model(**enc).logits.float()[:, self.perm]
out.append(torch.softmax(logits, -1).cpu().numpy())
return np.concatenate(out, 0)
# ----------------------------------------------------------------------------- MNLI
def eval_mnli(scorer, n=None):
res = {}
native2ours = {0: ENT, 1: NEU, 2: CON}
for split in ["validation_matched", "validation_mismatched"]:
ds = load_dataset("nyu-mll/multi_nli", split=split).filter(lambda x: x["label"] in (0, 1, 2))
if n:
ds = ds.shuffle(seed=0).select(range(n))
probs = scorer.predict(list(zip(ds["premise"], ds["hypothesis"])))
gold = np.array([native2ours[l] for l in ds["label"]])
res[split] = {"acc": float((probs.argmax(-1) == gold).mean()), "n": len(ds)}
return res
# ----------------------------------------------------------------------------- hard NLI test sets (v2)
NLI_SETS = { # name -> (repo, config, split, premise col, hypothesis col, label col)
"anli_r1": ("facebook/anli", None, "test_r1", "premise", "hypothesis", "label"),
"anli_r2": ("facebook/anli", None, "test_r2", "premise", "hypothesis", "label"),
"anli_r3": ("facebook/anli", None, "test_r3", "premise", "hypothesis", "label"),
"wanli": ("alisawuffles/WANLI", None, "test", "premise", "hypothesis", "gold"),
"scitail": ("allenai/scitail", "snli_format", "test", "sentence1", "sentence2", "gold_label"),
"control": ("tasksource/ConTRoL-nli", None, "test", "premise", "hypothesis", "label"),
}
_NAME2ID = {"contradiction": CON, "entailment": ENT, "neutral": NEU, "entails": ENT, "not_entailment": NEU}
def eval_nli_set(scorer, name):
repo, cfg, split, pc, hc, lc = NLI_SETS[name]
ds = load_dataset(repo, cfg, split=split) if cfg else load_dataset(repo, split=split)
feat = ds.features[lc]
pairs, gold = [], []
for ex in ds:
v = ex[lc]
nm = feat.names[v] if isinstance(v, int) and hasattr(feat, "names") else v
if not isinstance(nm, str) or nm.strip().lower() not in _NAME2ID:
continue
pairs.append((ex[pc], ex[hc])); gold.append(_NAME2ID[nm.strip().lower()])
probs = scorer.predict(pairs)
gold = np.array(gold)
pred = probs.argmax(-1)
two_class = len(set(gold.tolist())) == 2 # SciTail: entails/neutral only -> contradiction counts as "not entailed"
if two_class:
pred = np.where(pred == ENT, ENT, NEU)
return {"acc": float((pred == gold).mean()), "n": int(len(gold)), "two_class": two_class}
# ----------------------------------------------------------------------------- multiple choice
# Every MC task is a list of dicts {"q": str, "opts": [str], "gold": int, "hyp": callable|None}.
# rerank (blog #3): premise=q, hypothesis="The correct answer is: {opt}" -> argmax P(ent)
# grading (blog #6): premise=q + "Reference answer: {gold}", hypothesis="Answer: {opt}" -> ent <=> gold
FEWSHOT_HYP = lambda o: f"Answer: {o}"
def with_demos(q, demos):
"""Few-shot premise: k solved (question, gold answer) pairs followed by the question."""
return "\n\n".join(f"{dq}\nAnswer: {da}" for dq, da in demos) + "\n\n" + q
def load_gpqa(path="data/gpqa_diamond.csv", fewshot=0, seed=0):
import csv
import random as _r
rows = list(csv.DictReader(open(path))) if os.path.exists(path) else load_dataset("Idavidrein/gpqa", "gpqa_diamond", split="train")
base = []
for ex in rows:
opts = [ex["Correct Answer"], ex["Incorrect Answer 1"], ex["Incorrect Answer 2"], ex["Incorrect Answer 3"]]
base.append({"q": ex["Question"].strip(), "opts": [o.strip() for o in opts], "gold": 0})
if not fewshot:
return base
rng = _r.Random(seed)
items = []
for i, it in enumerate(base): # leave-one-out demos from other diamond questions
pool = [j for j in range(len(base)) if j != i]
demos = [(base[j]["q"], base[j]["opts"][base[j]["gold"]]) for j in rng.sample(pool, fewshot)]
items.append({"q": with_demos(it["q"], demos), "opts": it["opts"], "gold": it["gold"], "hyp": FEWSHOT_HYP})
return items
def load_mmlu(n=None, seed=0, fewshot=0):
ds = load_dataset("cais/mmlu", "all", split="test")
if n:
ds = ds.shuffle(seed=seed).select(range(n))
dev = {}
if fewshot: # standard MMLU few-shot: dev split, 5 per subject
for ex in load_dataset("cais/mmlu", "all", split="dev"):
dev.setdefault(ex["subject"], []).append((ex["question"].strip(), ex["choices"][int(ex["answer"])].strip()))
items = []
for ex in ds:
it = {"q": ex["question"].strip(), "opts": [c.strip() for c in ex["choices"]], "gold": int(ex["answer"])}
if fewshot:
it["q"] = with_demos(it["q"], dev[ex["subject"]][:fewshot])
it["hyp"] = FEWSHOT_HYP
items.append(it)
return items
def load_arc(cfg):
ds = load_dataset("allenai/ai2_arc", cfg, split="test")
items = []
for ex in ds:
labels = ex["choices"]["label"]
if ex["answerKey"] not in labels:
continue
items.append({"q": ex["question"].strip(), "opts": [t.strip() for t in ex["choices"]["text"]], "gold": labels.index(ex["answerKey"])})
return items
def load_winogrande():
ds = load_dataset("allenai/winogrande", "winogrande_xl", split="validation")
items = []
for ex in ds:
sent = ex["sentence"]
opts = [ex["option1"], ex["option2"]]
# hypothesis = sentence with the blank filled; premise = sentence with the blank left open
items.append({"q": sent, "opts": opts, "gold": int(ex["answer"]) - 1,
"hyp": lambda o, sent=sent: sent.replace("_", o)})
return items
def load_chess(n=500, seed=0):
"""Synthetic 'Chess (4 legal moves)': random position, 4 candidate moves in SAN, exactly one is legal."""
import random as _r
import chess
rng = _r.Random(seed)
items = []
while len(items) < n:
board = chess.Board()
for _ in range(rng.randint(6, 40)):
moves = list(board.legal_moves)
if not moves or board.is_game_over():
break
board.push(rng.choice(moves))
legal = list(board.legal_moves)
if len(legal) < 2 or board.is_game_over():
continue
legal_san = {board.san(m) for m in legal}
good = board.san(rng.choice(legal))
bad = set()
tries = 0
while len(bad) < 3 and tries < 500:
tries += 1
sq = rng.choice([s for s in chess.SQUARES if board.piece_at(s) and board.piece_at(s).color == board.turn])
piece = board.piece_at(sq)
to = rng.choice(chess.SQUARES)
if to == sq or (board.piece_at(to) and board.piece_at(to).color == board.turn):
continue
capture = board.piece_at(to) is not None
if piece.piece_type == chess.PAWN:
san = (chess.square_name(sq)[0] + "x" if capture else "") + chess.square_name(to)
else:
san = chess.piece_symbol(piece.piece_type).upper() + ("x" if capture else "") + chess.square_name(to)
if san not in legal_san and san != good:
bad.add(san)
if len(bad) < 3:
continue
opts = [good] + sorted(bad)
rng.shuffle(opts)
pgn = chess.Board().variation_san(board.move_stack)
q = (f"Chess position after the moves: {pgn}\nFEN: {board.fen()}\n"
f"{'White' if board.turn else 'Black'} to move. Which of the following moves is legal in this position?")
items.append({"q": q, "opts": opts, "gold": opts.index(good)})
return items
def load_hellaswag(n=None, seed=0, split="validation"):
ds = load_dataset("Rowan/hellaswag", split=split)
if n:
ds = ds.shuffle(seed=seed).select(range(n))
items = []
for ex in ds:
ctx = (ex["ctx_a"] + " " + ex["ctx_b"].capitalize()).strip() if ex["ctx_b"] else ex["ctx_a"].strip()
items.append({"q": f"{ex['activity_label']}: {ctx}", "opts": [e.strip() for e in ex["endings"]], "gold": int(ex["label"]),
"hyp": lambda o: o}) # hypothesis = the ending itself
return items
def load_gsm8k_mc(k=4, n=None, seed=0, split="test"):
"""GSM8K as k-way multiple choice: gold final answer + k-1 numeric distractors (deterministic perturbations)."""
import random as _r
rng = _r.Random(seed)
ds = load_dataset("openai/gsm8k", "main", split=split)
if n:
ds = ds.shuffle(seed=seed).select(range(n))
items = []
for ex in ds:
g = extract_number(ex["answer"])
gv = float(g)
cands = set()
gen = [lambda: gv + rng.choice([1, 2, 3, 5, 10]), lambda: gv - rng.choice([1, 2, 3, 5, 10]), lambda: gv * 2, lambda: gv / 2,
lambda: gv + rng.choice([4, 6, 7, 8, 9, 12, 15, 20, 25, 50]), lambda: gv * 10, lambda: gv * rng.choice([3, 4, 5]),
lambda: gv - rng.choice([4, 6, 7, 8, 9, 12, 15, 20, 25, 50]), lambda: abs(gv) + rng.randint(100, 999)]
gi = 0
while len(cands) < k - 1 and gi < 200:
v = gen[gi % len(gen)](); gi += 1
vs = str(int(v)) if float(v) == int(v) else f"{v:.2f}"
if vs != g and vs not in cands and v >= 0:
cands.add(vs)
opts = [g] + sorted(cands, key=lambda x: rng.random())
order = list(range(len(opts))); rng.shuffle(order)
opts = [opts[i] for i in order]
items.append({"q": ex["question"].strip(), "opts": opts, "gold": opts.index(g), "hyp": lambda o: f"The answer is {o}."})
return items
MC_TASKS = {
"hellaswag": lambda a: load_hellaswag(a.mc_n),
"gsm8k_mc4": lambda a: load_gsm8k_mc(4),
"gsm8k_mc10": lambda a: load_gsm8k_mc(10),
"gpqa": lambda a: load_gpqa(),
"mmlu": lambda a: load_mmlu(a.mc_n),
"gpqa_fewshot": lambda a: load_gpqa(fewshot=a.fewshot),
"mmlu_fewshot": lambda a: load_mmlu(a.mc_n, fewshot=a.fewshot),
"arc_easy": lambda a: load_arc("ARC-Easy"),
"arc_challenge": lambda a: load_arc("ARC-Challenge"),
"winogrande": lambda a: load_winogrande(),
"chess": lambda a: load_chess(a.chess_n),
}
def eval_mc(scorer, items):
rerank_pairs, grade_pairs, grade_gold, offsets = [], [], [], []
for it in items:
q, opts, g = it["q"], it["opts"], it["gold"]
hyp = it.get("hyp") or (lambda o: f"The correct answer is: {o}")
offsets.append((len(rerank_pairs), len(opts)))
for j, o in enumerate(opts):
rerank_pairs.append((q, hyp(o)))
grade_pairs.append((f"{q}\nReference answer: {opts[g]}", f"Answer: {o}"))
grade_gold.append(1 if j == g else 0)
pr = scorer.predict(rerank_pairs)
pg = scorer.predict(grade_pairs)
grade_gold = np.array(grade_gold)
rerank_hits, margin_hits, rank_hits, rand = [], [], [], []
for (s, k), it in zip(offsets, items):
g = it["gold"]
rerank_hits.append(pr[s:s+k, ENT].argmax() == g)
margin_hits.append((pr[s:s+k, ENT] - pr[s:s+k, CON]).argmax() == g)
rank_hits.append(pg[s:s+k, ENT].argmax() == g)
rand.append(1.0 / k)
pred_ent = (pg.argmax(-1) == ENT).astype(int)
return {
"n_questions": len(items),
"random_baseline": float(np.mean(rand)),
"rerank_acc": float(np.mean(rerank_hits)),
"rerank_margin_acc": float(np.mean(margin_hits)),
"grade_acc": float((pred_ent == grade_gold).mean()),
"grade_f1": float(f1_score(grade_gold, pred_ent)),
"grade_rank_acc": float(np.mean(rank_hits)),
"label_dist_rerank": np.bincount(pr.argmax(-1), minlength=3).tolist(),
}
# ----------------------------------------------------------------------------- GSM8K
NUM_RE = re.compile(r"-?\d[\d,]*\.?\d*")
def extract_number(text):
m = re.search(r"####\s*(-?[\d,]*\.?\d+)", text)
if m:
s = m.group(1)
else:
nums = NUM_RE.findall(text)
if not nums:
return None
s = nums[-1]
s = s.replace(",", "").rstrip(".")
try:
v = float(s)
except ValueError:
return None
return str(int(v)) if v == int(v) else str(v)
def llama_chat(prompt, temperature, max_tokens=512, retries=3):
body = {
"model": "qwen35",
"messages": [{"role": "user", "content": prompt}],
"max_tokens": max_tokens,
"temperature": temperature,
"top_p": 0.95,
"chat_template_kwargs": {"enable_thinking": False},
}
for _ in range(retries):
try:
r = requests.post(LLAMA_URL, json=body, timeout=300)
r.raise_for_status()
return r.json()["choices"][0]["message"]["content"]
except Exception as e: # noqa: BLE001
print("llama-server error:", e)
time.sleep(5)
return ""
def gen_gsm8k_candidates(cache, n_q, n_samples, seed=0):
if os.path.exists(cache):
return [json.loads(l) for l in open(cache)]
ds = load_dataset("openai/gsm8k", "main", split="test").shuffle(seed=seed).select(range(n_q))
rows = []
t0 = time.time()
for i, ex in enumerate(ds):
prompt = (ex["question"].strip() + "\n\nSolve the problem step by step. "
"Finish with a final line of the form: #### <number>")
gold = extract_number(ex["answer"])
cands = [{"text": llama_chat(prompt, 0.0), "kind": "greedy"}]
cands += [{"text": llama_chat(prompt, 0.7), "kind": "sample"} for _ in range(n_samples)]
for c in cands:
c["pred"] = extract_number(c["text"])
c["correct"] = c["pred"] is not None and c["pred"] == gold
rows.append({"question": ex["question"].strip(), "gold_solution": ex["answer"].strip(), "gold": gold, "cands": cands})
if (i + 1) % 10 == 0:
print(f" gsm8k gen {i+1}/{n_q} {time.time()-t0:.0f}s", flush=True)
os.makedirs(os.path.dirname(cache) or ".", exist_ok=True)
with open(cache, "w") as f:
for r in rows:
f.write(json.dumps(r) + "\n")
return rows
def eval_gsm8k(scorer, rows):
n = len(rows)
greedy = np.mean([r["cands"][0]["correct"] for r in rows])
samples = [r["cands"][1:] for r in rows]
k = len(samples[0])
def maj_vote(cs):
votes = collections.Counter(c["pred"] for c in cs if c["pred"] is not None)
if not votes:
return False
top = votes.most_common(1)[0][0]
return any(c["correct"] for c in cs if c["pred"] == top)
maj = np.mean([maj_vote(cs) for cs in samples])
oracle = np.mean([any(c["correct"] for c in cs) for cs in samples])
pass1 = np.mean([np.mean([c["correct"] for c in cs]) for cs in samples])
# best-of-k rerank: premise=question, hypothesis=candidate solution
pairs = [(r["question"], c["text"]) for r, cs in zip(rows, samples) for c in cs]
pr = scorer.predict(pairs).reshape(n, k, 3)
pick = pr[:, :, ENT].argmax(-1)
rerank = np.mean([samples[i][pick[i]]["correct"] for i in range(n)])
pick_m = (pr[:, :, ENT] - pr[:, :, CON]).argmax(-1)
rerank_margin = np.mean([samples[i][pick_m[i]]["correct"] for i in range(n)])
# grading: premise=question+gold solution, hypothesis="The answer is <pred>"
gpairs, ggold = [], []
for r in rows:
for c in r["cands"]:
if c["pred"] is None:
continue
gpairs.append((f"{r['question']}\nReference solution: {r['gold_solution']}", f"The answer is {c['pred']}."))
ggold.append(int(c["correct"]))
pg = scorer.predict(gpairs)
ggold = np.array(ggold)
pred = (pg.argmax(-1) == ENT).astype(int)
return {
"n_questions": n, "k": k,
"greedy_acc": float(greedy), "sample_pass1": float(pass1),
f"maj@{k}": float(maj), f"oracle@{k}": float(oracle),
f"nli_rerank@{k}": float(rerank), f"nli_rerank_margin@{k}": float(rerank_margin),
"grade_acc": float((pred == ggold).mean()), "grade_f1": float(f1_score(ggold, pred)),
"grade_n_pairs": int(len(ggold)), "grade_pos_rate": float(ggold.mean()),
}
# ----------------------------------------------------------------------------- main
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--models", nargs="+", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--gsm8k-cache", default="data/gsm8k_cands.jsonl")
ap.add_argument("--gsm8k-n", type=int, default=200)
ap.add_argument("--gsm8k-k", type=int, default=4)
ap.add_argument("--mnli-n", type=int, default=None)
ap.add_argument("--tasks", nargs="+", default=["mnli", "gpqa", "gsm8k"],
help="any of: mnli gsm8k " + " ".join(MC_TASKS))
ap.add_argument("--mc-n", type=int, default=None, help="subsample size for MMLU (default: full 14k)")
ap.add_argument("--chess-n", type=int, default=500)
ap.add_argument("--fewshot", type=int, default=5, help="k demos for *_fewshot tasks")
ap.add_argument("--bs", type=int, default=32)
ap.add_argument("--max-len", type=int, default=4096)
ap.add_argument("--gen-only", action="store_true", help="only generate GSM8K candidates and exit")
args = ap.parse_args()
rows = None
if "gsm8k" in args.tasks:
rows = gen_gsm8k_candidates(args.gsm8k_cache, args.gsm8k_n, args.gsm8k_k)
print(f"gsm8k candidates: {len(rows)} questions, greedy acc={np.mean([r['cands'][0]['correct'] for r in rows]):.3f}")
if args.gen_only:
return
mc_items = {t: MC_TASKS[t](args) for t in args.tasks if t in MC_TASKS}
for t, its in mc_items.items():
print(f"{t}: {len(its)} questions; example: {its[0]['q'][:120]!r} opts={its[0]['opts'][:4]}")
results = {}
for m in args.models:
print(f"\n===== {m}")
scorer = NLIScorer(m, bs=args.bs, max_len=args.max_len)
r = {}
if "mnli" in args.tasks:
r["mnli"] = eval_mnli(scorer, args.mnli_n); print("mnli", r["mnli"])
for t in args.tasks:
if t in NLI_SETS:
r[t] = eval_nli_set(scorer, t); print(t, r[t], flush=True)
for t in args.tasks:
if t in MC_TASKS:
r[t] = eval_mc(scorer, mc_items[t]); print(t, r[t], flush=True)
if rows is not None:
r["gsm8k"] = eval_gsm8k(scorer, rows); print("gsm8k", r["gsm8k"])
results[m] = r
del scorer; torch.cuda.empty_cache()
os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
json.dump(results, open(args.out, "w"), indent=2)
# summary table
print("\n| model | task | n | random | rerank acc | grade acc | grade F1 |")
print("|---|---|---|---|---|---|---|")
for m, r in results.items():
for t in args.tasks:
if t in MC_TASKS and t in r:
d = r[t]
print(f"| {m} | {t} | {d['n_questions']} | {d['random_baseline']:.3f} | {d['rerank_acc']:.3f} | {d['grade_acc']:.3f} | {d['grade_f1']:.3f} |")
if "mnli" in r:
print(f"| {m} | mnli m/mm | - | 0.333 | {r['mnli']['validation_matched']['acc']:.3f}/{r['mnli']['validation_mismatched']['acc']:.3f} | - | - |")
if "gsm8k" in r:
gs = r["gsm8k"]; k = gs["k"]
print(f"| {m} | gsm8k (greedy {gs['greedy_acc']:.3f}, maj@{k} {gs[f'maj@{k}']:.3f}, oracle {gs[f'oracle@{k}']:.3f}) | {gs['n_questions']} | - | {gs[f'nli_rerank@{k}']:.3f} | {gs['grade_acc']:.3f} | {gs['grade_f1']:.3f} |")
if __name__ == "__main__":
main()