81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""JMMLU (Japanese MMLU) benchmark.
|
|
|
|
Tests knowledge across 112 Japanese subjects using 0-shot multiple choice.
|
|
Includes Japan-specific topics like Japanese history and culture.
|
|
Dataset bundled from nlp-waseda/JMMLU on HuggingFace.
|
|
"""
|
|
|
|
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"
|
|
|
|
|
|
def _format_subject_name(subject: str) -> str:
|
|
return subject.replace("_", " ").replace("-", " ").title()
|
|
|
|
|
|
def _format_question(item: dict) -> str:
|
|
question = item["question"]
|
|
choices = item["choices"]
|
|
parts = [question]
|
|
for letter, choice in zip(["A", "B", "C", "D"], choices):
|
|
parts.append(f"{letter}. {choice}")
|
|
return "\n".join(parts)
|
|
|
|
|
|
class JMMLUBenchmark(BaseBenchmark):
|
|
"""JMMLU: 0-shot Japanese multiple choice across 112 subjects."""
|
|
|
|
name = "jmmlu"
|
|
quick_size = 300
|
|
|
|
async def load_dataset(self, sample_size: int = 0) -> list[dict]:
|
|
test_items = load_jsonl(DATA_DIR / "jmmlu_test.jsonl")
|
|
all_items = []
|
|
for item in test_items:
|
|
answer = item.get("answer", "A")
|
|
all_items.append({
|
|
"question": item["question"],
|
|
"choices": item["choices"],
|
|
"answer": answer, # Already A/B/C/D
|
|
"subject": item.get("subject", "unknown"),
|
|
})
|
|
|
|
logger.info(f"JMMLU: 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]]:
|
|
subject = item["subject"]
|
|
subject_name = _format_subject_name(subject)
|
|
|
|
parts = [
|
|
f"以下は{subject_name}に関する選択問題です。"
|
|
f"正解のアルファベット(A、B、C、D)だけを答えてください。\n"
|
|
]
|
|
|
|
parts.append(_format_question(item))
|
|
parts.append("答え:")
|
|
|
|
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")
|