HelioAI's picture
LFM3-1.7B inference modules: experimental release (part 2)
cd61333 verified
Raw History Blame Contribute Delete
10.7 kB
"""Grades generations produced by gen.py against the frozen eval suite.
Only the final answer (text after </think>) is graded; a response whose thinking never closed counts as wrong.
Writes results/<tag>/graded.jsonl (one row per sample) and prints per-task accuracy.
usage: grade.py <tag>
"""
import glob
import json
import os
import re
import subprocess
import sandbox # noqa: E402 (same folder)
import sys
import tempfile
import threading
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
TMP = os.path.join(tempfile.gettempdir(), "lfm3_codeexec")
os.makedirs(TMP, exist_ok=True)
def final_answer(completion):
if "</think>" not in completion:
return None
return completion.split("</think>", 1)[1].strip()
def last_boxed(text):
idx = text.rfind("\\boxed")
if idx < 0:
idx = text.rfind("\\fbox")
if idx < 0:
return None
i = text.find("{", idx)
if i < 0:
return None
depth = 0
for j in range(i, len(text)):
if text[j] == "{":
depth += 1
elif text[j] == "}":
depth -= 1
if depth == 0:
return text[i + 1: j]
return None
def grade_math(ans_text, gold):
from math_verify import parse, verify
if ans_text is None:
return False
pred = last_boxed(ans_text)
if pred is None:
return False
try:
g = parse("$" + gold + "$", parsing_timeout=None)
p = parse("$" + pred + "$", parsing_timeout=None)
if verify(g, p, timeout_seconds=None):
return True
except Exception:
pass
norm = lambda s: re.sub(r"\\text\{(.*?)\}|\\!|\s|\$|\\left|\\right|\.$", r"\1", s).replace("dfrac", "frac")
return norm(pred) == norm(gold)
def grade_gsm8k(ans_text, gold):
if ans_text is None:
return False
pred = last_boxed(ans_text)
if pred is None:
nums = re.findall(r"-?\d[\d,]*\.?\d*", ans_text)
if not nums:
return False
pred = nums[-1]
pred = re.sub(r"\\text\{.*?\}|\\\$|\$|,|\\%|%|\\!|\s", "", pred)
m = re.search(r"-?\d+(\.\d+)?", pred.replace("{", "").replace("}", ""))
if not m:
return False
try:
return abs(float(m.group(0)) - float(gold)) < 1e-6
except ValueError:
return False
def extract_code(ans_text):
if ans_text is None:
return None
blocks = re.findall(r"```(?:python|py|Python)?\s*\n(.*?)```", ans_text, re.S)
if not blocks:
return None
return max(blocks, key=len) if len(blocks) > 1 and "def " not in blocks[-1] else blocks[-1]
def run_python(src, timeout):
fd, path = tempfile.mkstemp(suffix=".py", dir=TMP)
with os.fdopen(fd, "w", encoding="utf-8") as f:
f.write(src)
try:
r = sandbox.run([sys.executable, "-I", path], capture_output=True, timeout=timeout, cwd=TMP,
creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
return r.returncode == 0
except subprocess.TimeoutExpired:
return False
finally:
try:
os.remove(path)
except OSError:
pass
def grade_humaneval(ans_text, item):
code = extract_code(ans_text)
if code is None:
return False
if f"def {item['entry_point']}" not in code:
code = item["code_prompt"] + code
src = code + "\n\n" + item["test"] + f"\n\ncheck({item['entry_point']})\n"
return run_python(src, 60)
def grade_mbpp(ans_text, item):
code = extract_code(ans_text)
if code is None:
return False
src = "\n".join(item.get("test_imports") or []) + "\n" + code + "\n\n" + item["test"] + "\n"
return run_python(src, 60)
def run_io(code, inp, out, timeout=6):
fd, path = tempfile.mkstemp(suffix=".py", dir=TMP)
with os.fdopen(fd, "w", encoding="utf-8") as f:
f.write(code)
try:
r = sandbox.run([sys.executable, "-I", path], input=inp.encode(), capture_output=True, timeout=timeout,
cwd=TMP, creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
if r.returncode != 0:
return False
return r.stdout.decode(errors="replace").split() == out.split()
except subprocess.TimeoutExpired:
return False
finally:
try:
os.remove(path)
except OSError:
pass
def grade_code_io(ans_text, item):
code = extract_code(ans_text)
if code is None:
return False
return all(run_io(code, i, o) for i, o in zip(item["inputs"], item["outputs"]))
def grade_lcb(ans_text, item):
from lcb_tests import run_lcb
code = extract_code(ans_text)
if code is None:
return False
k, n, _ = run_lcb(code, item, item["hidden"])
return k == n
_IFEVAL_LOCK = threading.Lock()
def grade_ifeval(ans_text, item):
# lm_eval fills empty instruction arguments with random values and langdetect is random too: seed both under a
# lock so that a response always gets the same grade
import random
from langdetect import DetectorFactory
from lm_eval.tasks.ifeval.utils import process_results
doc = {"key": 0, "prompt": item["prompt"], "instruction_id_list": item["instruction_id_list"],
"kwargs": item["kwargs"]}
with _IFEVAL_LOCK:
random.seed(0)
DetectorFactory.seed = 0
res = process_results(doc, [ans_text or ""])
return {"strict": res["prompt_level_strict_acc"], "loose": res["prompt_level_loose_acc"],
"inst_strict": res["inst_level_strict_acc"]}
def grade_mmlupro(ans_text, gold):
if ans_text is None:
return False
m = re.findall(r"answer is:?\s*\(?\**\s*([A-J])\b", ans_text)
if not m:
m = re.findall(r"\\boxed\{\(?([A-J])\)?\}", ans_text)
return bool(m) and m[-1] == gold
def grade_bbh(ans_text, gold, prompt=None):
"""Exact match after extraction. With the prompt (flexible mode) it also accepts answers that state the same
thing in an unambiguous other form: the text of the chosen option instead of its letter, a word list separated by
commas, the whole Dyck sequence (input + completion) instead of the completion, plausible/not plausible for
yes/no. Each rule maps the answer to exactly one gold-comparable form, so it cannot accept a different answer."""
if ans_text is None:
return False
m = re.findall(r"answer is:?\s*(.+)", ans_text, re.I)
pred = m[-1] if m else (last_boxed(ans_text) or "")
pred = re.sub(r"\\boxed\{(.*)\}", r"\1", pred).strip().strip("*").strip().rstrip(".").strip()
gold = gold.strip()
norm = lambda s: re.sub(r"\s+", " ", s.lower().replace("`", "").replace("*", "")).strip().rstrip(".").strip()
if re.fullmatch(r"\([A-R]\)", gold):
mm = re.search(r"\(?([A-R])\)", pred) or re.fullmatch(r"([A-R])", pred)
if mm or prompt is None:
return bool(mm) and mm.group(1) == gold[1]
opts = {norm(t): k for k, t in re.findall(r"^\(([A-R])\)\s*(.+?)\s*$", prompt, re.M)}
return opts.get(norm(pred)) == gold[1]
if norm(pred) == norm(gold):
return True
if prompt is None:
return False
if "Sort the following words alphabetically" in prompt:
return re.findall(r"[^\s,]+", norm(pred)) == norm(gold).split()
if "parentheses are closed properly" in prompt:
inp = re.search(r"Input:\s*(.+)", prompt).group(1)
sp = lambda s: re.sub(r"\s+", "", norm(s))
return sp(pred) == sp(inp) + sp(gold)
yn = {"plausible": "yes", "not plausible": "no", "implausible": "no"}
return norm(gold) in ("yes", "no") and yn.get(norm(pred)) == norm(gold)
def load_suite(pattern=os.path.join(ROOT, "data", "eval", "*.jsonl")):
items = {}
for path in glob.glob(pattern):
with open(path, encoding="utf-8") as f:
for l in f:
r = json.loads(l)
items[r["id"]] = r
return items
def grade_one(rec, item):
a = final_answer(rec["completion"])
t = item["task"]
extra = {}
if t in ("math500", "aime"):
ok = grade_math(a, item["answer"])
elif t == "gsm8k":
ok = grade_gsm8k(a, item["answer"])
elif t == "humaneval":
ok = grade_humaneval(a, item)
elif t in ("mbpp", "code_fn"):
ok = grade_mbpp(a, item)
elif t == "code_io":
ok = grade_code_io(a, item)
elif t == "lcb":
ok = grade_lcb(a, item)
elif t == "math":
ok = grade_math(a, item["answer"])
elif t == "ifeval":
extra = grade_ifeval(a, item)
ok = extra["strict"]
elif t == "mmlupro":
ok = grade_mmlupro(a, item["answer"])
elif t == "bbh":
ok = grade_bbh(a, item["answer"], item["prompt"])
extra = {"strict": grade_bbh(a, item["answer"])}
else:
raise ValueError(t)
return {"id": rec["id"], "task": t, "sample": rec["sample"], "correct": bool(ok), "n_tokens": rec["n_tokens"],
"truncated": rec["truncated"], "closed_think": a is not None,
"subtask": item.get("subtask") or item.get("category") or item.get("level"), **extra}
def main(tag, pattern=None):
items = load_suite(pattern) if pattern else load_suite()
recs = []
for path in glob.glob(os.path.join(ROOT, "results", tag, "*.jsonl")):
if path.endswith("graded.jsonl"):
continue
with open(path, encoding="utf-8") as f:
recs += [json.loads(l) for l in f]
with ThreadPoolExecutor(12) as ex:
graded = list(ex.map(lambda r: grade_one(r, items[r["id"]]), recs))
with open(os.path.join(ROOT, "results", tag, "graded.jsonl"), "w", encoding="utf-8") as f:
for g in graded:
f.write(json.dumps(g) + "\n")
agg = defaultdict(list)
for g in graded:
agg[g["task"]].append(g)
print(f"{'task':10s} {'n':>5s} {'acc':>7s} {'trunc':>6s} {'no</think>':>10s} {'avg_tok':>8s}")
for t in sorted(agg):
rows = agg[t]
n = len(rows)
print(f"{t:10s} {n:5d} {100 * sum(r['correct'] for r in rows) / n:7.2f} "
f"{100 * sum(r['truncated'] for r in rows) / n:6.1f} {100 * sum(not r['closed_think'] for r in rows) / n:10.1f} "
f"{sum(r['n_tokens'] for r in rows) / n:8.0f}")
if t == "ifeval":
print(f"{' loose':10s} {'':5s} {100 * sum(r['loose'] for r in rows) / n:7.2f} inst_strict "
f"{100 * sum(sum(r['inst_strict']) for r in rows) / sum(len(r['inst_strict']) for r in rows):.2f}")
if __name__ == "__main__":
main(*sys.argv[1:])