"""MMLU (official source: cais/mmlu)."""

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

_LETTERS = 'ABCDEFGHIJ'


@register_dataset(
    DatasetSpec(
        name='mmlu',
        source='cais/mmlu',  # official: https://huggingface.co/datasets/cais/mmlu
        subset='all',  # 57 subjects; override with --subset <subject>
        split='test',
        prompt_style='cot_letter',  # es contract: CoT + last-line ANSWER
        few_shot_split='dev',
        few_shot_num=5,
        task_type='mcq',
        tags=['knowledge'],
        description='Massive Multitask Language Understanding (official).',
    )
)
def mmlu():
    def to_sample(record: dict) -> Sample:
        return Sample(
            input=record['question'],
            choices=list(record['choices']),
            target=_LETTERS[int(record['answer'])],
            metadata={'subject': record.get('subject')},
        )

    return to_sample