ToolOrchestratorEnv / env /answer_grading.py
Andrew Lara
Initial implementation of ToolOrchestratorEnv
8ca3a35
Raw
History Blame Contribute Delete
3.19 kB
"""Answer grading utilities: exact match + token F1.
Ported from SearchEconomicsEnv/env/answer_grading.py and adapted for
multi-domain use (HotpotQA-style EM/F1 + code/math fallback).
"""
from __future__ import annotations
import json
import re
import string
from collections import Counter
from typing import Tuple
# ---------------------------------------------------------------------------
# Normalisation
# ---------------------------------------------------------------------------
def normalize_answer(text: str) -> list[str]:
"""Lowercase, strip articles/punctuation, tokenise."""
text = text.lower().strip()
# Remove articles
text = re.sub(r"\b(a|an|the)\b", " ", text)
# Remove punctuation
text = text.translate(str.maketrans("", "", string.punctuation))
return text.split()
# ---------------------------------------------------------------------------
# Metrics
# ---------------------------------------------------------------------------
def exact_match(pred: str, gold: str) -> bool:
return normalize_answer(pred) == normalize_answer(gold)
def token_f1(pred: str, gold: str) -> float:
pred_tokens = normalize_answer(pred)
gold_tokens = normalize_answer(gold)
if not pred_tokens or not gold_tokens:
return float(pred_tokens == gold_tokens)
common = Counter(pred_tokens) & Counter(gold_tokens)
num_common = sum(common.values())
if num_common == 0:
return 0.0
precision = num_common / len(pred_tokens)
recall = num_common / len(gold_tokens)
return 2 * precision * recall / (precision + recall)
# ---------------------------------------------------------------------------
# Answer extraction
# ---------------------------------------------------------------------------
def extract_answer(raw: str) -> str:
"""Pull the answer string out of various agent output formats."""
# Strip markdown fences
raw = re.sub(r"```[a-z]*\n?", "", raw).strip()
# Try JSON {"answer": ...}
try:
parsed = json.loads(raw)
if isinstance(parsed, dict):
for key in ("answer", "Answer", "result", "Result"):
if key in parsed:
return str(parsed[key]).strip()
except (json.JSONDecodeError, ValueError):
pass
# Prefix patterns
for prefix in ("Answer:", "Final answer:", "Result:", "Output:"):
idx = raw.lower().find(prefix.lower())
if idx != -1:
return raw[idx + len(prefix):].strip().split("\n")[0].strip()
# Last non-empty line
lines = [line.strip() for line in raw.splitlines() if line.strip()]
return lines[-1] if lines else raw.strip()
# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------
def grade(predicted: str, ground_truth: str) -> Tuple[bool, float, float]:
"""Return (exact_match, f1, quality) where quality ∈ [0, 1]."""
extracted = extract_answer(predicted)
em = exact_match(extracted, ground_truth)
f1 = token_f1(extracted, ground_truth)
quality = 1.0 if em else f1
return em, f1, quality