# SPDX-License-Identifier: Apache-2.0
"""Base classes and data models for accuracy benchmarks."""
import asyncio
import logging
import re
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Callable, Optional
logger = logging.getLogger(__name__)
# Token budget for thinking/reasoning models (industry reference: OpenCompass 8K~32K)
THINKING_MIN_TOKENS = 8192
THINKING_MAX_TOKENS = 32768
@dataclass
class QuestionResult:
"""Result for a single benchmark question."""
question_id: str
correct: bool
expected: str
predicted: str
time_seconds: float
question_text: str = ""
raw_response: str = ""
category: Optional[str] = None
@dataclass
class BenchmarkResult:
"""Aggregated result for a complete benchmark run."""
benchmark_name: str
accuracy: float
total_questions: int
correct_count: int
time_seconds: float
question_results: list[QuestionResult] = field(default_factory=list)
category_scores: Optional[dict[str, float]] = None
thinking_used: bool = False
class BaseBenchmark(ABC):
"""Abstract base class for accuracy benchmarks."""
name: str = ""
quick_size: int = 100
@abstractmethod
async def load_dataset(self, sample_size: int = 0) -> list[dict]:
"""Load dataset items.
Args:
sample_size: Number of questions to sample. 0 = full dataset.
Returns:
List of dataset items (format varies by benchmark).
"""
pass
@abstractmethod
def format_prompt(self, item: dict) -> list[dict[str, str]]:
"""Format a dataset item into chat messages for the engine.
Returns:
List of message dicts with 'role' and 'content' keys.
"""
pass
@abstractmethod
def extract_answer(self, response: str, item: dict) -> str:
"""Extract the predicted answer from model response text."""
pass
@abstractmethod
def check_answer(self, predicted: str, item: dict) -> bool:
"""Check if the predicted answer is correct."""
pass
def get_max_tokens(self) -> int:
"""Max tokens to generate per question. Override for longer answers."""
return 128
def get_category(self, item: dict) -> Optional[str]:
"""Return category/subject for per-category scoring. None if N/A."""
return None
def get_question_text(self, item: dict) -> str:
"""Return a human-readable question text for result export."""
return item.get("question", item.get("description", item.get("context", "")))
@staticmethod
def _extract_mc_answer(response: str, valid_letters: list[str]) -> str:
"""Extract multiple choice answer from response.
Strategy:
1. Look for explicit "answer is X" / "answer: X" patterns (last match)
2. Fall back to last valid letter in response
3. Case-insensitive
"""
response_upper = response.strip().upper()
pattern_letters = "".join(valid_letters)
# 1. Look for "answer is X", "answer: X", "answer X" patterns — use LAST match
answer_patterns = re.findall(
r"(?:answer\s*(?:is|:)\s*)([" + pattern_letters + r"])\b",
response_upper,
)
if answer_patterns:
return answer_patterns[-1]
# 2. Fall back to last valid letter with word boundary
all_matches = re.findall(
r"\b([" + pattern_letters + r"])\b",
response_upper,
)
if all_matches:
return all_matches[-1]
# 3. Check first character
if response.strip() and response.strip()[0].upper() in valid_letters:
return response.strip()[0].upper()
return ""
@staticmethod
def _extract_last_code_block(response: str) -> str:
"""Extract the LAST code block from model response.
Uses last match to avoid picking up drafts/examples.
Falls back to line-by-line detection if no code blocks found.
"""
response = response.strip()
# Find ALL python code blocks, use LAST
blocks = re.findall(r"```python\s*\n(.*?)```", response, re.DOTALL)
if blocks:
return blocks[-1].strip()
# Generic code blocks
blocks = re.findall(r"```\s*\n(.*?)```", response, re.DOTALL)
if blocks:
return blocks[-1].strip()
# Line-by-line fallback
lines = response.split("\n")
code_lines = []
in_code = False
for line in lines:
if not in_code and (
line.startswith("def ")
or line.startswith("class ")
or line.startswith("import ")
or line.startswith("from ")
or line.startswith("#")
):
in_code = True
if in_code:
code_lines.append(line)
return "\n".join(code_lines) if code_lines else response
@staticmethod
def _strip_think_tags(text: str) -> str:
"""Remove ... blocks from model output."""
return re.sub(r".*?", "", text, flags=re.DOTALL).strip()
async def _eval_single(
self, engine: Any, item: dict, index: int,
sampling_kwargs: Optional[dict] = None,
enable_thinking: bool = False,
) -> tuple[int, dict, str, str, str]:
"""Evaluate a single item.
Returns (index, item, response_text, prompt_text, raw_text).
raw_text is the unstripped output for auto-detection of thinking tags.
"""
messages = self.format_prompt(item)
prompt_text = "\n".join(m.get("content", "") for m in messages)
kwargs = dict(sampling_kwargs or {})
# max_tokens is always benchmark-controlled — a model's small configured
# limit must not truncate long answers and corrupt scores.
max_tokens = self.get_max_tokens()
# Harmony models (gpt_oss) use analysis + final channels;
# analysis can consume the entire budget before final is emitted
if getattr(engine, "model_type", None) == "gpt_oss":
max_tokens = max(max_tokens * 4, 8192)
elif enable_thinking:
max_tokens = min(
max(max_tokens, THINKING_MIN_TOKENS), THINKING_MAX_TOKENS
)
kwargs["max_tokens"] = max_tokens
# Greedy/neutral defaults keep scores reproducible. setdefault (not
# force-set) lets the caller's "model_settings" sampling profile supply
# its own temperature/penalties; the default profile passes none, so
# these fall through to deterministic values.
kwargs.setdefault("temperature", 0.0)
kwargs.setdefault("presence_penalty", 0.0)
kwargs.setdefault("repetition_penalty", 1.0)
# Merge enable_thinking into any existing chat_template_kwargs
ct_kwargs = kwargs.pop("chat_template_kwargs", {}) or {}
ct_kwargs["enable_thinking"] = enable_thinking
kwargs["chat_template_kwargs"] = ct_kwargs
try:
output = await engine.chat(
messages=messages,
**kwargs,
)
raw_text = output.text
text = self._strip_think_tags(raw_text)
return index, item, text, prompt_text, raw_text
except Exception as e:
logger.warning(f"Engine error on question {index}: {e}")
return index, item, "", prompt_text, ""
async def run(
self,
engine: Any,
items: list[dict],
on_progress: Optional[Callable[[int, int], Any]] = None,
batch_size: int = 1,
sampling_kwargs: Optional[dict] = None,
enable_thinking: bool = False,
) -> BenchmarkResult:
"""Run the benchmark on all items.
Args:
engine: oMLX engine instance with chat() method.
items: Dataset items to evaluate.
on_progress: Callback(current, total) for progress reporting.
batch_size: Number of concurrent requests (1 = sequential).
enable_thinking: Enable thinking mode for reasoning models.
When False, auto-detects if the model outputs tags
and re-runs the first batch with thinking enabled.
Returns:
BenchmarkResult with accuracy and per-question details.
"""
results: list[QuestionResult] = []
correct = 0
category_correct: dict[str, int] = {}
category_total: dict[str, int] = {}
start_time = time.time()
completed = 0
thinking_used = enable_thinking
auto_switched = False
# Process in batches
for batch_start in range(0, len(items), batch_size):
batch_end = min(batch_start + batch_size, len(items))
batch = items[batch_start:batch_end]
batch_start_time = time.time()
# Launch concurrent requests
tasks = [
self._eval_single(
engine, item, batch_start + j, sampling_kwargs, thinking_used
)
for j, item in enumerate(batch)
]
batch_results = await asyncio.gather(*tasks)
# Auto-detection: check first batch for tags
if not thinking_used and not auto_switched and batch_start == 0:
auto_switched = True
has_think_tags = any(
"" in raw for _, _, _, _, raw in batch_results
)
if has_think_tags:
logger.warning(
f"{self.name}: model outputs tags with "
"enable_thinking=False, auto-switching to thinking mode"
)
thinking_used = True
# Re-run first batch with increased token budget
tasks = [
self._eval_single(
engine, item, batch_start + j, sampling_kwargs, True
)
for j, item in enumerate(batch)
]
batch_results = await asyncio.gather(*tasks)
batch_elapsed = time.time() - batch_start_time
# Process results in order
for idx, item, response_text, prompt_text, _raw in sorted(batch_results, key=lambda x: x[0]):
predicted = self.extract_answer(response_text, item)
is_correct = self.check_answer(predicted, item)
if is_correct:
correct += 1
cat = self.get_category(item)
if cat is not None:
category_total[cat] = category_total.get(cat, 0) + 1
if is_correct:
category_correct[cat] = category_correct.get(cat, 0) + 1
q_id = item.get("id", str(idx))
expected = item.get("answer", "")
results.append(
QuestionResult(
question_id=str(q_id),
correct=is_correct,
expected=str(expected),
predicted=predicted,
time_seconds=batch_elapsed / len(batch),
question_text=prompt_text,
raw_response=response_text,
category=cat,
)
)
completed += len(batch)
if on_progress:
await on_progress(completed, len(items))
total_time = time.time() - start_time
total = len(items)
accuracy = correct / total if total > 0 else 0.0
cat_scores = None
if category_total:
cat_scores = {}
for cat in sorted(category_total.keys()):
cat_scores[cat] = (
category_correct.get(cat, 0) / category_total[cat]
if category_total[cat] > 0
else 0.0
)
return BenchmarkResult(
benchmark_name=self.name,
accuracy=accuracy,
total_questions=total,
correct_count=correct,
time_seconds=total_time,
question_results=results,
category_scores=cat_scores,
thinking_used=thinking_used,
)