"""MMLU-Pro (official source: TIGER-Lab/MMLU-Pro)."""

from ..sample import Sample
from ..registry import register_dataset
from ..spec import DatasetSpec


@register_dataset(
    DatasetSpec(
        name='mmlu_pro',
        source='TIGER-Lab/MMLU-Pro',  # official: https://huggingface.co/datasets/TIGER-Lab/MMLU-Pro
        split='test',
        few_shot_split='validation',
        few_shot_num=5,
        task_type='mcq',
        tags=['knowledge'],
        description='MMLU-Pro: 10-option harder MMLU (official).',
    )
)
def mmlu_pro():
    def to_sample(record: dict) -> Sample:
        return Sample(
            input=record['question'],
            choices=list(record['options']),
            target=str(record['answer']).strip(),  # already a letter
            metadata={'category': record.get('category'), 'question_id': record.get('question_id')},
        )

    return to_sample