from opencompass.openicl.icl_prompt_template import PromptTemplate
from opencompass.openicl.icl_retriever import BM25Retriever
from opencompass.openicl.icl_inferencer import GenInferencer
from opencompass.openicl.icl_evaluator import BleuEvaluator
from opencompass.datasets import IWSLT2017Dataset

iwslt2017_reader_cfg = dict(
    input_columns='en', output_column='de', train_split='validation')

iwslt2017_infer_cfg = dict(
    ice_template=dict(type='PromptTemplate',
                    template=dict(
                        begin=[
                            dict(role='SYSTEM', fallback_role="HUMAN", prompt='Please translate the following English statements to German:'),
                            '</E>',
                        ],
                        round=[
                            dict(role='HUMAN', prompt='{en}'),
                            dict(role='BOT', prompt='{de}'),
                        ]
                    ),
                    ice_token='</E>'),
    retriever=dict(type=BM25Retriever, ice_num=1),
    inferencer=dict(type=GenInferencer))

iwslt2017_eval_cfg = dict(
    evaluator=dict(type=BleuEvaluator), 
    pred_role='BOT', 
    pred_postprocessor=dict(type='general_cn'),
    dataset_postprocessor=dict(type='general_cn'))

iwslt2017_datasets = [
    dict(
        type=IWSLT2017Dataset,
        path='iwslt2017',
        name='iwslt2017-en-de',
        reader_cfg=iwslt2017_reader_cfg,
        infer_cfg=iwslt2017_infer_cfg,
        eval_cfg=iwslt2017_eval_cfg)
]