wehub-resource-sync
129 lines
4.3 KiB
Python
129 lines
4.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""MMLU (Massive Multitask Language Understanding) benchmark.
|
|
|
|
Tests knowledge across 57 subjects using 5-shot multiple choice.
|
|
Dataset bundled from cais/mmlu on HuggingFace.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from .base import BaseBenchmark
|
|
from .datasets import load_jsonl, stratified_sample
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DATA_DIR = Path(__file__).parent / "data"
|
|
ANSWER_MAP = {0: "A", 1: "B", 2: "C", 3: "D"}
|
|
|
|
|
|
def _format_subject_name(subject: str) -> str:
|
|
"""Convert subject slug to readable name."""
|
|
return subject.replace("_", " ").title()
|
|
|
|
|
|
def _format_question(item: dict) -> str:
|
|
"""Format a single MMLU question with choices."""
|
|
question = item["question"]
|
|
choices = item["choices"]
|
|
parts = [question]
|
|
for i, choice in enumerate(choices):
|
|
parts.append(f"{ANSWER_MAP[i]}. {choice}")
|
|
return "\n".join(parts)
|
|
|
|
|
|
def _parse_choices(choices_field):
|
|
"""Parse choices field which may be a list or a string repr of a list."""
|
|
if isinstance(choices_field, list):
|
|
return choices_field
|
|
if isinstance(choices_field, str):
|
|
try:
|
|
parsed = json.loads(choices_field.replace("'", '"'))
|
|
if isinstance(parsed, list):
|
|
return parsed
|
|
except (json.JSONDecodeError, ValueError):
|
|
pass
|
|
return []
|
|
|
|
|
|
class MMLUBenchmark(BaseBenchmark):
|
|
"""MMLU: 5-shot multiple choice across 57 academic subjects."""
|
|
|
|
name = "mmlu"
|
|
quick_size = 300
|
|
|
|
def __init__(self):
|
|
self._few_shot_examples: dict[str, list[dict]] = {}
|
|
|
|
async def load_dataset(self, sample_size: int = 0) -> list[dict]:
|
|
"""Load MMLU from bundled data files."""
|
|
test_items = load_jsonl(DATA_DIR / "mmlu_test.jsonl")
|
|
all_items = []
|
|
for item in test_items:
|
|
choices = _parse_choices(item.get("choices", []))
|
|
answer_idx = item.get("answer", 0)
|
|
answer_letter = ANSWER_MAP.get(answer_idx, str(answer_idx))
|
|
all_items.append({
|
|
"question": item["question"],
|
|
"choices": choices,
|
|
"answer": answer_letter,
|
|
"subject": item.get("subject", "unknown"),
|
|
})
|
|
|
|
# Load dev examples for few-shot
|
|
dev_items = load_jsonl(DATA_DIR / "mmlu_dev.jsonl")
|
|
for item in dev_items:
|
|
subject = item.get("subject", "unknown")
|
|
choices = _parse_choices(item.get("choices", []))
|
|
answer_idx = item.get("answer", 0)
|
|
answer_letter = ANSWER_MAP.get(answer_idx, str(answer_idx))
|
|
if subject not in self._few_shot_examples:
|
|
self._few_shot_examples[subject] = []
|
|
if len(self._few_shot_examples[subject]) < 5:
|
|
self._few_shot_examples[subject].append({
|
|
"question": item["question"],
|
|
"choices": choices,
|
|
"answer": answer_letter,
|
|
})
|
|
|
|
logger.info(f"MMLU: loaded {len(all_items)} questions")
|
|
|
|
if sample_size == 0:
|
|
return all_items
|
|
|
|
return stratified_sample(all_items, sample_size, key="subject")
|
|
|
|
def format_prompt(self, item: dict) -> list[dict[str, str]]:
|
|
"""Format with 5-shot examples from the same subject."""
|
|
subject = item["subject"]
|
|
subject_name = _format_subject_name(subject)
|
|
|
|
parts = [
|
|
f"The following are multiple choice questions about {subject_name}. "
|
|
f"Answer with just the letter (A, B, C, or D).\n"
|
|
]
|
|
|
|
# Add few-shot examples
|
|
examples = self._few_shot_examples.get(subject, [])
|
|
for ex in examples:
|
|
parts.append(_format_question(ex))
|
|
parts.append(f"Answer: {ex['answer']}\n")
|
|
|
|
# Add the actual question
|
|
parts.append(_format_question(item))
|
|
parts.append("Answer:")
|
|
|
|
return [{"role": "user", "content": "\n".join(parts)}]
|
|
|
|
def extract_answer(self, response: str, item: dict) -> str:
|
|
return self._extract_mc_answer(response, ["A", "B", "C", "D"])
|
|
|
|
def check_answer(self, predicted: str, item: dict) -> bool:
|
|
return predicted == item["answer"]
|
|
|
|
def get_category(self, item: dict) -> Optional[str]:
|
|
return item.get("subject")
|