from opencompass.datasets import KaoshiDataset, KaoshiEvaluator
from opencompass.openicl.icl_inferencer import GenInferencer
from opencompass.openicl.icl_prompt_template import PromptTemplate
from opencompass.openicl.icl_retriever import ZeroRetriever

prompts = {
        '单选题' : '请你做一道单项选择题\n请你一步一步思考并将思考过程写在【解析】和<eoe>之间。你将从A，B，C，D中选出正确的答案，并写在【答案】和<eoa>之间，答案应只包含最终结果，不要添加额外词语。\n例如：【答案】: A <eoa>\n完整的题目回答的格式如下：\n【解析】 ... <eoe>\n【答案】 ... <eoa>\n请你严格按照上述格式作答。\n题目如下：',
        '多选题' : '请你做一道多项选择题\n请你一步一步思考并将思考过程写在【解析】和<eoe>之间。你将从多个选项中选出正确的答案，答案可能是一个到多个选项，奇怪将其写在【答案】和<eoa>之间，答案应只包含最终结果，不要添加额外词语。\n例如：【答案】: A D <eoa>\n完整的题目回答的格式如下：\n【解析】 ... <eoe>\n【答案】 ... <eoa>\n请你严格按照上述格式作答。\n题目如下：',
        '填空题' : '请解答下面的填空题\n仔细阅读题目，解答其中的问题，请你一步步思考并将思考过程写在【解析】和<eoe>之间。请把你的答案写在【答案】和<eoa>之间，答案应只包含最终结果，不要添加额外词语。\n完整的题目回答格式如下：\n【解析】 ... <eoe>\n【答案】... <eoa>\n请你严格按照上述格式作答。\n题目如下:',
        '完形填空' : '请你做一道英语完形填空题,其中包含二十个小题。\n请你一步一步思考。每一题你将从A，B，C，D中选出正确的答案，并写在【答案】和<eoa>之间。\n例如：（1）【答案】 A <eoa>\n（2）【答案】 B <eoa>\n请你严格按照上述格式作答。\n',
        '七选五': '请回答下面的问题，将符合题意的五个选项的字母写在【答案】和<eoa>之间，例如：【答案】 A B C D E <eoa>\n请严格按照上述格式作答。题目如下：\n',
        '判断题' : '请回答下面的判断题，将你的判断结果写在【答案】和<eoa>之间，若给定表述正确时回答：\n【答案】正确 <eoa>\n 表述错误时回答：\n【答案】错误 <eoa>\n请严格按照上述格式作答。题目如下：\n',
}

splits_with_type = {'单选题': ['职业-消防', '职业-测绘', '考研-经济', '职业-安全工程', '考研-政治', '职业-建筑', '考研-英语', '职业-教师资格', '职业-证券', '职业-会计', '职业-公务员', '考研-数学', '职业-高项', '考研-临床医学', '职业-银行', '考研-管理类综合', '职业-基金'],
                    '多选题': ['职业-消防', '职业-测绘', '考研-政治', '职业-建筑', '职业-证券', '职业-会计', '考研-临床医学', '职业-银行'],
                    '完形填空': ['考研-英语'],
                    '七选五': ['考研-英语'],
                    '判断题': ['职业-证券'],
                    '填空题': ['考研-数学']}

zh2en = {'单选题': 'single_choice', '多选题': 'multi_choice', '完形填空': 'multi_question_choice', '判断题': 'judgment', '填空题': 'cloze', '七选五': 'five_out_of_seven'}

kaoshi_datasets = []

for _type in list(splits_with_type.keys()):
    for _split in splits_with_type[_type]:
        _folder = _split.replace('-' + _type, '')
        _p = prompts[_type]
        _reader_cfg = {
            'input_columns': ['question'],
            'output_column': 'answer',
        }
        _infer_cfg = {
            'ice_template': {
                'type': PromptTemplate,
                'template': {
                    'round': [{
                        'role': 'HUMAN',
                        'prompt': _p + '{question}'
                    }]
                },
                'ice_token': '</E>'
            },
            'retriever': {
                'type': ZeroRetriever
            },
            'inferencer': {
                'type': GenInferencer,
                'max_out_len': 1024,
            }
        }
        _eval_cfg = {
            'evaluator': {
                'type': KaoshiEvaluator,
                  'question_type': zh2en[_type],
            },
            'pred_role': 'BOT',
        }
        _base_path = './data/Kaoshi'
        _dataset = {
            'type': KaoshiDataset,
            'abbr': 'Kaoshi' + _split + '-' + _type,
            'path': _base_path + '/' + _folder + '/' + _type + '.jsonl',
            'name': zh2en[_type],
            'reader_cfg': _reader_cfg,
            'infer_cfg': _infer_cfg,
            'eval_cfg': _eval_cfg,
        }

        kaoshi_datasets.append(_dataset)

_temporary_variables = [k for k in globals() if k.startswith('_')]
for _t in _temporary_variables:
    del globals()[_t]
del _temporary_variables, _t