# Copyright (c) Alibaba, Inc. and its affiliates.
from mmengine.config import read_base
from opencompass.partitioners import NaivePartitioner
from opencompass.runners import LocalRunner
from opencompass.tasks import OpenICLInferTask


with read_base():
    # from opencompass.configs.summarizers.medium import summarizer
    # from opencompass.configs.summarizers.PMMEval import summarizer
    from evalscope.backend.opencompass.tasks.eval_datasets import datasets

# 1. Get datasets
# Note: OpenAI API format evaluation needs a special humaneval postprocessor
for _dataset in datasets:
    if _dataset['path'] == 'openai_humaneval':
        from opencompass.datasets.humaneval import humaneval_gpt_postprocess
        _dataset['eval_cfg']['pred_postprocessor']['type'] = humaneval_gpt_postprocess


# 2. Get models, only for placeholder, you should fill in the real model information from command line
# See more templates in `opencompass.cli.arguments.ApiModelConfig`
models = []

# 3. Get infer config
infer = dict(
    partitioner=dict(type=NaivePartitioner),
    runner=dict(
        type=LocalRunner,
        max_num_workers=4,
        task=dict(type=OpenICLInferTask)),
)
