"""DROP (official source: ucinlp/drop)."""

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


@register_dataset(
    DatasetSpec(
        name='drop',
        source='ucinlp/drop',  # official: https://huggingface.co/datasets/ucinlp/drop
        split='validation',
        task_type='qa',
        tags=['reading_comprehension'],
        description='DROP reading comprehension; target = answer spans list.',
    )
)
def drop():
    def to_sample(record: dict) -> Sample:
        spans = record['answers_spans']['spans']
        return Sample(
            input=record['question'],
            target=list(spans) if len(spans) > 1 else spans[0],
            metadata={
                'passage': record['passage'],  # required reading context, kept out of input
                'query_id': record.get('query_id'),
                'section_id': record.get('section_id'),
            },
        )

    return to_sample