wehub-resource-sync
239 lines
7.9 KiB
Python
239 lines
7.9 KiB
Python
"""Self-contained LLM-as-judge grader for the PixelRAG reproduction.
|
|
|
|
Migrated from the paper's evaluation/worldvqa_eval/worldvqa_eval.py + evaluate.py
|
|
(the encyclopedic_vqa / mmsearch / worldvqa path) so the eval pipeline does not
|
|
depend on the old dr-agent (Vis-RAG) repo. Behaviour is byte-faithful to the
|
|
paper grader:
|
|
|
|
- Judge prompt = JUDGE_WORLDQA_PROMPT_EN (verbatim from MoonshotAI/WorldVQA),
|
|
loaded from eval/repro_assets/judge_worldvqa_prompt.txt.
|
|
- Ground truth:
|
|
* encyclopedic_vqa -> "Any of: " + " | ".join(reference_list) (ANY match = correct)
|
|
* mmsearch / worldvqa -> gt_answer (single string)
|
|
- The model response has <think>...</think> stripped before judging.
|
|
- Judge model gpt-4.1-2025-04-14, temperature=0; verdict parsed from a
|
|
`Label: Correct|Incorrect|Unattempted` line.
|
|
- score = #Correct / N.
|
|
|
|
CLI:
|
|
python -m lib.grader <task> <responses.jsonl> [--grader-model gpt-4.1-2025-04-14]
|
|
Requires OPENAI_API_KEY (+ optional OPENAI_BASE_URL) in the environment.
|
|
"""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import re
|
|
import string
|
|
from pathlib import Path
|
|
|
|
_ASSETS = Path(__file__).resolve().parent.parent / "repro_assets"
|
|
JUDGE_WORLDQA_PROMPT_EN = (_ASSETS / "judge_worldvqa_prompt.txt").read_text()
|
|
SIMPLEQA_GRADER_TEMPLATE = (_ASSETS / "simpleqa_grader_template.txt").read_text()
|
|
|
|
# Which grader each task uses (matches paper scripts/evaluate.py dispatch).
|
|
WORLDVQA_TASKS = {
|
|
"encyclopedic_vqa",
|
|
"mmsearch",
|
|
"worldvqa",
|
|
"factualvqa",
|
|
"webqa",
|
|
"multimodalqa",
|
|
}
|
|
EXACT_MATCH_TASKS = {"nq", "nq_tables", "triviaqa"}
|
|
SIMPLEQA_TASKS = {"simpleqa", "simpleqa_verified"}
|
|
|
|
DEFAULT_GRADER_MODEL = "gpt-4.1-2025-04-14"
|
|
# Match the paper grader sampler (scripts/evaluate.py -> ChatCompletionSampler):
|
|
# system message "You are a helpful assistant.", temperature=0, max_tokens=1000, seed=42.
|
|
GRADER_SYSTEM_MESSAGE = "You are a helpful assistant."
|
|
GRADER_MAX_TOKENS = 1000
|
|
GRADER_SEED = 42
|
|
|
|
|
|
def strip_think(text: str) -> str:
|
|
# Verbatim from paper worldvqa_eval.strip_think_tags.
|
|
if text is None:
|
|
return ""
|
|
if "<think>" in text and "</think>" in text:
|
|
return text.split("</think>")[-1].strip()
|
|
elif "think>" in text:
|
|
return text.split("think>")[-1].strip()
|
|
return text
|
|
|
|
|
|
def build_ground_truth(task: str, original_data: dict) -> str:
|
|
"""Match evaluate.py convert_to_evaluate_format."""
|
|
if task == "encyclopedic_vqa":
|
|
refs = original_data.get("reference_list") or []
|
|
if refs:
|
|
return "Any of: " + " | ".join(refs)
|
|
return original_data.get("answer", "") or original_data.get("gt_answer", "")
|
|
# mmsearch / worldvqa / simplevqa / factualvqa
|
|
return original_data.get("gt_answer", "") or original_data.get("answer", "")
|
|
|
|
|
|
def parse_label(judge_text: str) -> str:
|
|
m = re.search(
|
|
r"Label:\s*(Correct|Incorrect|Unattempted)", judge_text, re.IGNORECASE
|
|
)
|
|
if m:
|
|
return m.group(1).lower()
|
|
tl = judge_text.lower()
|
|
if "incorrect" in tl:
|
|
return "incorrect"
|
|
if "unattempted" in tl:
|
|
return "unattempted"
|
|
if "correct" in tl:
|
|
return "correct"
|
|
return "incorrect"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# NQ / NQ-Tables exact-match (verbatim from short_form_qa_eval.short_form_eval)
|
|
# ---------------------------------------------------------------------------
|
|
def _normalize_text(s: str) -> str:
|
|
s = re.sub(
|
|
r"\b(a|an|the)\b",
|
|
" ",
|
|
s.lower().translate(str.maketrans("", "", string.punctuation)),
|
|
)
|
|
return " ".join(s.split())
|
|
|
|
|
|
def is_exact_match(prediction: str, golds) -> bool:
|
|
prediction = (prediction or "").replace("Exact Answer: ", "").strip()
|
|
pred_norm = _normalize_text(prediction)
|
|
return any(_normalize_text(str(g)) == pred_norm for g in golds)
|
|
|
|
|
|
def _golds_for(task: str, od: dict):
|
|
if task in EXACT_MATCH_TASKS:
|
|
g = (
|
|
od.get("answers")
|
|
or od.get("reference_list")
|
|
or od.get("answer")
|
|
or od.get("gt_answer")
|
|
)
|
|
return g if isinstance(g, list) else [g]
|
|
return None
|
|
|
|
|
|
def grade_exact_match(path: str) -> dict:
|
|
rows = [json.loads(l) for l in open(path)]
|
|
c = 0
|
|
for d in rows:
|
|
golds = _golds_for("nq", d.get("original_data", {}))
|
|
if is_exact_match(strip_think(d.get("final_response")), golds):
|
|
c += 1
|
|
n = len(rows)
|
|
return {
|
|
"task": "exact_match",
|
|
"file": path,
|
|
"n": n,
|
|
"correct": c,
|
|
"incorrect": n - c,
|
|
"unattempted": 0,
|
|
"errors": 0,
|
|
"score": c / n if n else 0.0,
|
|
}
|
|
|
|
|
|
async def grade_file(
|
|
task: str,
|
|
path: str,
|
|
grader_model: str = DEFAULT_GRADER_MODEL,
|
|
concurrency: int = 16,
|
|
) -> dict:
|
|
if task in EXACT_MATCH_TASKS:
|
|
return grade_exact_match(path)
|
|
from openai import AsyncOpenAI
|
|
|
|
client = AsyncOpenAI(
|
|
api_key=os.environ["OPENAI_API_KEY"], base_url=os.environ.get("OPENAI_BASE_URL")
|
|
)
|
|
rows = [json.loads(l) for l in open(path)]
|
|
sem = asyncio.Semaphore(concurrency)
|
|
labels = [None] * len(rows)
|
|
|
|
is_sqa = task in SIMPLEQA_TASKS
|
|
|
|
async def judge(i, d):
|
|
od = d.get("original_data", {})
|
|
answer = strip_think(d.get("final_response"))
|
|
if is_sqa:
|
|
target = od.get("answer", "") or od.get("gt_answer", "")
|
|
prompt = SIMPLEQA_GRADER_TEMPLATE.format(
|
|
question=d.get("problem", ""), target=target, predicted_answer=answer
|
|
)
|
|
else:
|
|
gt = build_ground_truth(task, od)
|
|
prompt = JUDGE_WORLDQA_PROMPT_EN.format(
|
|
question=d.get("problem", ""),
|
|
model_answer=answer,
|
|
ground_truth_answer=gt,
|
|
)
|
|
async with sem:
|
|
try:
|
|
r = await client.chat.completions.create(
|
|
model=grader_model,
|
|
temperature=0,
|
|
max_tokens=GRADER_MAX_TOKENS,
|
|
seed=GRADER_SEED,
|
|
messages=[
|
|
{"role": "system", "content": GRADER_SYSTEM_MESSAGE},
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
)
|
|
out = r.choices[0].message.content
|
|
if is_sqa:
|
|
m = re.search(r"(A|B|C)", out or "")
|
|
letter = m.group(0) if m else "C"
|
|
labels[i] = {"A": "correct", "B": "incorrect", "C": "unattempted"}[
|
|
letter
|
|
]
|
|
else:
|
|
labels[i] = parse_label(out)
|
|
except Exception as e:
|
|
labels[i] = ("__error__", str(e))
|
|
|
|
await asyncio.gather(*[judge(i, d) for i, d in enumerate(rows)])
|
|
errs = [l for l in labels if isinstance(l, tuple)]
|
|
verdicts = [l for l in labels if isinstance(l, str)]
|
|
n = len(verdicts)
|
|
c = verdicts.count("correct")
|
|
inc = verdicts.count("incorrect")
|
|
una = verdicts.count("unattempted")
|
|
return {
|
|
"task": task,
|
|
"file": path,
|
|
"n": n,
|
|
"correct": c,
|
|
"incorrect": inc,
|
|
"unattempted": una,
|
|
"errors": len(errs),
|
|
"score": c / n if n else 0.0,
|
|
}
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("task", help="encyclopedic_vqa | mmsearch | worldvqa | ...")
|
|
ap.add_argument("jsonl", help="responses jsonl from run_bench.py")
|
|
ap.add_argument("--grader-model", default=DEFAULT_GRADER_MODEL)
|
|
ap.add_argument("--concurrency", type=int, default=16)
|
|
args = ap.parse_args()
|
|
res = asyncio.run(
|
|
grade_file(args.task, args.jsonl, args.grader_model, args.concurrency)
|
|
)
|
|
print(
|
|
f"{Path(res['file']).name}: {res['correct']}/{res['n']} = {res['score']:.4f} "
|
|
f"(C={res['correct']} I={res['incorrect']} U={res['unattempted']} err={res['errors']})"
|
|
)
|
|
print(f"Score: {res['score']:.3f}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|