"""Winogrande (official source: allenai/winogrande, winogrande_xl)."""

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


@register_dataset(
    DatasetSpec(
        name='winogrande',
        source='allenai/winogrande',  # official: https://huggingface.co/datasets/allenai/winogrande
        subset='winogrande_xl',
        split='validation',
        task_type='mcq',
        tags=['commonsense', 'coreference'],
        description='Winogrande XL binary coreference (official).',
    )
)
def winogrande():
    def to_sample(record: dict) -> Sample:
        return Sample(
            input=record['sentence'],
            choices=[record['option1'], record['option2']],
            target={'1': 'A', '2': 'B'}[record['answer']],
        )

    return to_sample