"""HellaSwag (official source: rowanz/hellaswag)."""

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

_LETTERS = 'ABCD'


@register_dataset(
    DatasetSpec(
        name='hellaswag',
        source='Rowan/hellaswag',  # parquet conversion of the official rowanz/hellaswag
        split='validation',
        task_type='mcq',
        tags=['commonsense'],
        description='HellaSwag commonsense sentence completion (official).',
    )
)
def hellaswag():
    def to_sample(record: dict) -> Sample:
        return Sample(
            input=record['ctx'],
            choices=list(record['endings']),
            target=_LETTERS[int(record['label'])],
            metadata={'activity_label': record.get('activity_label')},
        )

    return to_sample