"""GSM8K dataset plugin (official source: openai/gsm8k).

Offline demo: examples/data/gsm8k_main_test.jsonl ships a tiny subset, e.g.
``evalharness data fetch gsm8k --source examples/data/gsm8k_main_test.jsonl``
"""

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


@register_dataset(
    DatasetSpec(
        name='gsm8k',
        source='openai/gsm8k',  # official: https://huggingface.co/datasets/openai/gsm8k
        subset='main',
        split='test',
        few_shot_split='train',
        few_shot_num=4,
        gen_config={'temperature': 0.0, 'max_tokens': 32768},
        prompt_suffix="\nPlease reason step by step, and put your final answer within \\boxed{}.",
        task_type='math',
        tags=['math', 'cot'],
        description='Grade school math word problems (OpenAI, official).',
    )
)
def gsm8k():
    def to_sample(record: dict) -> Sample:
        parts = record['answer'].split('####')
        target = parts.pop().strip()
        return Sample(
            input=record['question'],
            target=target,
            metadata={'reasoning': '####'.join(parts).strip()},
        )

    return to_sample