from opencompass.openicl.icl_prompt_template import PromptTemplate
from opencompass.openicl.icl_retriever import ZeroRetriever
from opencompass.openicl.icl_inferencer import GenInferencer
from opencompass.openicl.icl_evaluator import AccEvaluator
from opencompass.datasets.OpenFinData import OpenFinDataDataset, OpenFinDataKWEvaluator
from opencompass.utils.text_postprocessors import last_capital_postprocess

OpenFinData_datasets = []
OpenFinData_3choices_list = ['emotion_identification', 'entity_disambiguation', 'financial_facts']
OpenFinData_4choices_list = ['data_inspection', 'financial_terminology', 'metric_calculation', 'value_extraction']
OpenFinData_5choices_list = ['intent_understanding']
OpenFinData_keyword_list = ['entity_recognition']
OpenFinData_all_list = OpenFinData_3choices_list + OpenFinData_4choices_list + OpenFinData_5choices_list + OpenFinData_keyword_list

OpenFinData_eval_cfg = dict(evaluator=dict(type=AccEvaluator), pred_postprocessor=dict(type=last_capital_postprocess))
OpenFinData_KW_eval_cfg = dict(evaluator=dict(type=OpenFinDataKWEvaluator))

for _name in OpenFinData_all_list:
    if _name in OpenFinData_3choices_list:
        OpenFinData_infer_cfg = dict(
            ice_template=dict(type=PromptTemplate, template=dict(begin='</E>', round=[
                        dict(role='HUMAN', prompt=f'{{question}}\nA. {{A}}\nB. {{B}}\nC. {{C}}\n答案: '),
                        dict(role='BOT', prompt='{answer}')]),
                        ice_token='</E>'), retriever=dict(type=ZeroRetriever), inferencer=dict(type=GenInferencer))

        OpenFinData_datasets.append(
            dict(
                type=OpenFinDataDataset,
                path='./data/openfindata_release',
                name=_name,
                abbr='OpenFinData-' + _name,
                reader_cfg=dict(
                    input_columns=['question', 'A', 'B', 'C'],
                    output_column='answer'),
                infer_cfg=OpenFinData_infer_cfg,
                eval_cfg=OpenFinData_eval_cfg,
            ))

    if _name in OpenFinData_4choices_list:
        OpenFinData_infer_cfg = dict(
            ice_template=dict(type=PromptTemplate, template=dict(begin='</E>', round=[
                        dict(role='HUMAN', prompt=f'{{question}}\nA. {{A}}\nB. {{B}}\nC. {{C}}\nD. {{D}}\n答案: '),
                        dict(role='BOT', prompt='{answer}')]),
                        ice_token='</E>'), retriever=dict(type=ZeroRetriever), inferencer=dict(type=GenInferencer))

        OpenFinData_datasets.append(
            dict(
                type=OpenFinDataDataset,
                path='./data/openfindata_release',
                name=_name,
                abbr='OpenFinData-' + _name,
                reader_cfg=dict(
                    input_columns=['question', 'A', 'B', 'C', 'D'],
                    output_column='answer'),
                infer_cfg=OpenFinData_infer_cfg,
                eval_cfg=OpenFinData_eval_cfg,
            ))

    if _name in OpenFinData_5choices_list:
        OpenFinData_infer_cfg = dict(
            ice_template=dict(type=PromptTemplate, template=dict(begin='</E>', round=[
                        dict(role='HUMAN', prompt=f'{{question}}\nA. {{A}}\nB. {{B}}\nC. {{C}}\nD. {{D}}\nE. {{E}}\n答案: '),
                        dict(role='BOT', prompt='{answer}')]),
                        ice_token='</E>'), retriever=dict(type=ZeroRetriever), inferencer=dict(type=GenInferencer))

        OpenFinData_datasets.append(
            dict(
                type=OpenFinDataDataset,
                path='./data/openfindata_release',
                name=_name,
                abbr='OpenFinData-' + _name,
                reader_cfg=dict(
                    input_columns=['question', 'A', 'B', 'C', 'D', 'E'],
                    output_column='answer'),
                infer_cfg=OpenFinData_infer_cfg,
                eval_cfg=OpenFinData_eval_cfg,
            ))

    if _name in OpenFinData_keyword_list:
        OpenFinData_infer_cfg = dict(
            ice_template=dict(type=PromptTemplate, template=dict(begin='</E>', round=[
                        dict(role='HUMAN', prompt=f'{{question}}\n答案: '),
                        dict(role='BOT', prompt='{answer}')]),
                        ice_token='</E>'), retriever=dict(type=ZeroRetriever), inferencer=dict(type=GenInferencer))

        OpenFinData_datasets.append(
            dict(
                type=OpenFinDataDataset,
                path='./data/openfindata_release',
                name=_name,
                abbr='OpenFinData-' + _name,
                reader_cfg=dict(
                    input_columns=['question'],
                    output_column='answer'),
                infer_cfg=OpenFinData_infer_cfg,
                eval_cfg=OpenFinData_KW_eval_cfg,
            ))

del _name