from opencompass.openicl.icl_prompt_template import PromptTemplate
from opencompass.openicl.icl_retriever import ZeroRetriever
from opencompass.openicl.icl_inferencer import PPLInferencer
from opencompass.openicl.icl_evaluator import AccEvaluator
from opencompass.datasets import CommonsenseQADataset_CN

commonsenseqacn_reader_cfg = dict(
    input_columns=["question", "A", "B", "C", "D", "E"],
    output_column="answerKey",
    test_split="validation",
)

_ice_template = dict(
    type=PromptTemplate,
    template={
        ans: dict(
            begin="</E>",
            round=[
                dict(role="HUMAN", prompt="问题: {question}\n答案: "),
                dict(role="BOT", prompt=ans_token),
            ],
        )
        for ans, ans_token in [
            ["A", "{A}"],
            ["B", "{B}"],
            ["C", "{C}"],
            ["D", "{D}"],
            ["E", "{E}"],
        ]
    },
    ice_token="</E>",
)


commonsenseqacn_infer_cfg = dict(
    prompt_template=_ice_template,
    retriever=dict(type=ZeroRetriever),
    inferencer=dict(type=PPLInferencer),
)

commonsenseqacn_eval_cfg = dict(evaluator=dict(type=AccEvaluator))

commonsenseqacn_datasets = [
    dict(
        abbr="commonsenseqa_cn",
        type=CommonsenseQADataset_CN,
        path="./data/commonsenseqa_cn/validation.jsonl",
        reader_cfg=commonsenseqacn_reader_cfg,
        infer_cfg=commonsenseqacn_infer_cfg,
        eval_cfg=commonsenseqacn_eval_cfg,
    )
]