File size: 2,312 Bytes
9f5fbb5 | 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 | #!/usr/bin/env python3
"""Rescore harness samples with format-robust extraction.
Why: thinking models break stock filters two ways — (1) they echo the
answer prefix back (`The best answer is C` vs target `C`), failing exact
match on a correct answer; (2) long thought can leave `content` empty on a
tight token budget (counted separately as UNANSWERED, not wrong).
This script scores the raw `resps` text both sides can audit.
Letter tasks (ARC-chat, MMLU-generative): last bare A-D token vs target.
GSM8K: `#### N` if present, else last integer vs target number.
Usage: python3 rescore.py <samples.jsonl> [letter|gsm8k]
"""
import json
import re
import sys
path, kind = sys.argv[1], (sys.argv[2] if len(sys.argv) > 2 else "letter")
rows = [json.loads(l) for l in open(path)]
correct = wrong = unanswered = 0
for r in rows:
raw = r.get("resps") or []
text = raw[0] if raw and isinstance(raw[0], str) else (
raw[0][0] if raw and raw[0] else "")
if not (text or "").strip():
unanswered += 1
continue
if kind == "gsm8k":
m = re.search(r"####\s*(-?\d[\d,]*)", text)
if not m:
nums = re.findall(r"-?\d[\d,]*", text)
m = [None, nums[-1]] if nums else None
pred = (m[1] if isinstance(m, list) else m.group(1)).replace(",", "") \
if m else None
tm = re.search(r"####\s*(-?\d[\d,]*)", r.get("target", ""))
tgt = tm.group(1).replace(",", "") if tm else None
correct += pred is not None and tgt is not None and pred == tgt
wrong += not (pred is not None and tgt is not None and pred == tgt)
else: # letter tasks: explicit "answer is X" wins, else last bare letter
m = re.search(r"[Aa]nswer is\s*\**([A-D])", text)
pred = m.group(1) if m else None
if pred is None:
letters = re.findall(r"\b([A-D])\b", text)
if not letters:
unanswered += 1
continue
pred = letters[-1]
if pred == (r.get("target") or "").strip():
correct += 1
else:
wrong += 1
n = len(rows)
print(f"{path.split('/')[-1]}: n={n} correct={correct} wrong={wrong} "
f"unanswered={unanswered} acc={correct / n:.3f} "
f"(answered-acc={correct / max(1, correct + wrong):.3f})")
|