From c958d1cbc07d4f279912c39fe12fdf40a4f13f45 Mon Sep 17 00:00:00 2001
From: jingmingzhuo <jingmingzhuo@gmail.com>
Date: Fri, 1 Dec 2023 21:45:43 +0800
Subject: [PATCH] [Feature] Add Codereval inference

---
 .../CoderEval/CoderEval_gen_a33a69.py         | 47 +++++++++++++++++++
 opencompass/datasets/__init__.py              |  1 +
 opencompass/datasets/codereval.py             | 20 ++++++++
 3 files changed, 68 insertions(+)
 create mode 100644 configs/datasets/CoderEval/CoderEval_gen_a33a69.py
 create mode 100644 opencompass/datasets/codereval.py

diff --git a/configs/datasets/CoderEval/CoderEval_gen_a33a69.py b/configs/datasets/CoderEval/CoderEval_gen_a33a69.py
new file mode 100644
index 00000000..0076f8cb
--- /dev/null
+++ b/configs/datasets/CoderEval/CoderEval_gen_a33a69.py
@@ -0,0 +1,47 @@
+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 import CoderEvalDataset
+from opencompass.utils.text_postprocessors import first_capital_postprocess
+
+
+CoderEval_reader_cfg = dict(
+    input_columns="input",
+    output_column=None,
+)
+
+CoderEval_infer_cfg = dict(
+    prompt_template=dict(
+        type=PromptTemplate,
+        template=dict(
+            round=[
+                dict(role="HUMAN", prompt="Please help me complete the following function.\n**Note: only return the function to me, no other description.**\n```python\n{input}\n```"),
+                dict(role="BOT", prompt="{answer}"),
+            ]
+        ),
+    ),
+    retriever=dict(type=ZeroRetriever),
+    inferencer=dict(type=GenInferencer),
+)
+
+CoderEval_eval_cfg = dict(evaluator=dict(type=AccEvaluator),
+                        pred_role="BOT",
+                        pred_postprocessor=dict(type=first_capital_postprocess))
+
+files = ['CEPythonHumanLabel', 'CEPythonRaw']
+CoderEval_datasets = []
+
+for _file in files:
+    CoderEval_datasets.append(
+        dict(
+            type=CoderEvalDataset,
+            abbr=_file,
+            test_path=f"data/CoderEval/{_file}.jsonl",
+            reader_cfg=CoderEval_reader_cfg,
+            infer_cfg=CoderEval_infer_cfg,
+            eval_cfg=CoderEval_eval_cfg,
+        )
+    )
+
+del _file
\ No newline at end of file
diff --git a/opencompass/datasets/__init__.py b/opencompass/datasets/__init__.py
index effbbf41..65934cd6 100644
--- a/opencompass/datasets/__init__.py
+++ b/opencompass/datasets/__init__.py
@@ -21,6 +21,7 @@ from .cmb import *  # noqa: F401, F403
 from .cmmlu import *  # noqa: F401, F403
 from .cmnli import *  # noqa: F401, F403
 from .cmrc import *  # noqa: F401, F403
+from .codereval import *  # noqa: F401, F403
 from .commonsenseqa import *  # noqa: F401, F403
 from .commonsenseqa_cn import *  # noqa: F401, F403
 from .copa import *  # noqa: F401, F403
diff --git a/opencompass/datasets/codereval.py b/opencompass/datasets/codereval.py
new file mode 100644
index 00000000..0d8b07ea
--- /dev/null
+++ b/opencompass/datasets/codereval.py
@@ -0,0 +1,20 @@
+import json
+
+from datasets import Dataset
+
+from opencompass.registry import LOAD_DATASET
+
+from .base import BaseDataset
+
+
+@LOAD_DATASET.register_module()
+class CoderEvalDataset(BaseDataset):
+
+    @staticmethod
+    def load(test_path):
+        datasets = []
+        with open(test_path, 'r', encoding='utf-8') as file:
+            for line in file:
+                dataset = json.loads(line.strip())
+                datasets.append(dataset)
+        return Dataset.from_list(datasets)