"""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(
            # es adapter parity: ctx_a + ' ' + ctx_b.capitalize() -- the
            # mirror's pre-joined `ctx` keeps ctx_b lowercase, es capitalizes
            input=str(record['ctx_a']).strip() + ' ' + str(record['ctx_b']).strip().capitalize(),
            choices=list(record['endings']),
            target=_LETTERS[int(record['label'])],
            metadata={'activity_label': record.get('activity_label')},
        )

    return to_sample