roles_map = {
    'system': 'system',
    'user': 'user',
    'human': 'user',
    'assistant': 'assistant',
    'gpt': 'assistant',
    'AI': 'assistant',
}


pretrain_reasoning_datasets = [
    #
    # basic reasoning
    #
    # 10.8 MB, 15,770
    {'kind': 'instruct', 'path': 'AtlasUnified/Atlas-Reasoning', 'data_files': 'reasoning.csv', 'transform': lambda r: [
        {'role': 'user', 'content': r['Prompt']},
        {'role': 'assistant', 'content': r['Step-by-step reasoning'] + '\n' + r['Solution']},
    ]},
    # 1.23 GB, 859,594
    *[
        {'kind': 'instruct', 'path': 'AI-MO/NuminaMath-CoT', 'split': f'train[{i}%:{i + 10}%]', 'field': 'messages'}
        for i in range(0, 100, 10)
    ],
    # 148 MB, 72,540
    *[
        {'kind': 'instruct', 'path': 'AI-MO/NuminaMath-TIR', 'split': f'train[{i}%:{i + 10}%]', 'field': 'messages'}
        for i in range(0, 100, 10)
    ],

    #
    # math reasoning
    #
    # 1.79 MB, 3,963
    {'kind': 'instruct', 'path': 'AlgorithmicResearchGroup/math_reasoning_autoformalization_track', 'transform': lambda r: [
        {'role': 'user', 'content': r['informal_statement']},
        {'role': 'assistant', 'content': r['informal_proof'] + '\n' + r['formal_proof']},
    ]},
    # 307 MB, 19,944
    {'kind': 'instruct', 'path': 'KingNish/reasoning-base-20k', 'transform': lambda r: [
        {'role': 'user', 'content': r['user']},
        {'role': 'assistant', 'content': r['reasoning'] + '\n' + r['assistant']},
    ]},
    # 9.45 MB, 10,000
    {'kind': 'instruct', 'path': 'Aarushhh/math-reasoning-10k', 'transform': lambda r: [
        {'role': 'user', 'content': r['problem']},
        {'role': 'assistant', 'content': r['plan'] + '\n' + r['solution']},
    ]},

    #
    # cot reasoning
    #
    # 11.7 GB, 1,850,809
    *[
        {'kind': 'instruct', 'path': 'ServiceNow-AI/R1-Distill-SFT', 'data_dir': 'v0', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
            {'role': 'user', 'content': r['problem']},
            {'role': 'assistant', 'content': r['reannotated_assistant_content']},
        ]}
        for i in range(0, 100, 10)
    ],
    *[
        {'kind': 'instruct', 'path': 'ServiceNow-AI/R1-Distill-SFT', 'data_dir': 'v1', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: r['reannotated_messages']}
        for i in range(0, 100, 10)
    ],
    # 3.85 GB, 300k (3.98 GB, 814,334)
    *[
        {'kind': 'instruct', 'path': 'cognitivecomputations/dolphin-r1', 'data_files': 'dolphin-r1-reasoning-deepseek.jsonl', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
            *r['messages'],
            # {'role': 'assistant', 'content': (('<think>\n' + r['reasoning'] + '\n</think>\n') if r.get('reasoning') else '') + r['answer']},
            {'role': 'assistant', 'content': (r.get('reasoning') or '') + (r.get('answer') or '')},
        ]}
        for i in range(0, 100, 10)
    ],
    # 3.49 GB, 300k (3.98 GB, 814,334)
    *[
        {'kind': 'instruct', 'path': 'cognitivecomputations/dolphin-r1', 'data_files': 'dolphin-r1-reasoning-flash.jsonl', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
            *r['messages'],
            # {'role': 'assistant', 'content': (('<think>\n' + r['reasoning'] + '\n</think>\n') if r.get('reasoning') else '') + r['answer']},
            {'role': 'assistant', 'content': (r.get('reasoning') or '') + (r.get('answer') or '')},
        ]}
        for i in range(0, 100, 10)
    ],
    # 1.08 GB, 113,957
    {'kind': 'instruct', 'path': 'open-thoughts/OpenThoughts-114k', 'split': 'train', 'field': 'conversations', 'transform': lambda msgs: [
        {'role': roles_map[m['from']], 'content': m['value']}
        for m in msgs
    ]},
    # 384 MB, 77,685
    {'kind': 'instruct', 'path': 'O1-OPEN/OpenO1-SFT', 'split': 'train', 'transform': lambda r: [
        {'role': 'user', 'content': r['instruction']},
        {'role': 'assistant', 'content': r['output']},
    ]},
    # 6.88 MB, 1,000
    {'kind': 'instruct', 'path': 'simplescaling/s1K', 'split': 'train', 'transform': lambda r: [
        {'role': 'user', 'content': r['question']},
        {'role': 'assistant', 'content': '<think>\n' + '\n'.join(r['thinking_trajectories']) + '\n</think>\n' + r['solution']},
    ]},
]