OpenCompass/opencompass/datasets/clozeTest_maxmin.py
jingmingzhuo b3cbef3226
[Feature] Add py150 and maxmin (#562)
* [feat] add clozeTesst_maxmin dataset

* [feat] add py150 datasets

* [feat] change __init__.py in opencompass/datasets

* [fix] pre-commit check

* [fix] rename py150 and masxmin datasets in configs

* [feat] add gen.py of py150 and maxmin in configs/datasets
2023-11-09 22:05:25 +08:00

36 lines
1.2 KiB
Python

import json
from datasets import Dataset
from opencompass.registry import LOAD_DATASET
from .base import BaseDataset
@LOAD_DATASET.register_module()
class MaxminDataset(BaseDataset):
@staticmethod
def load(test_path, answer_path=None):
if answer_path is not None:
with open(answer_path, 'r', encoding='utf-8') as answer_f:
answers = {}
for line in answer_f.readlines():
line = line.strip()
answers[line.split('<CODESPLIT>')[0]] = line.split(
'<CODESPLIT>')[1]
datasets = []
with open(test_path, 'r') as test_f:
test_data = json.load(test_f)
for item in test_data:
dataset = dict()
dataset['nl_tokens'] = ' '.join(item['nl_tokens'])
dataset['pl_tokens'] = ' '.join(item['pl_tokens'])
if answer_path is not None:
dataset['answer'] = 'A' if answers[
item['idx']] == 'max' else 'B'
else:
dataset['answer'] = ''
datasets.append(dataset)
return Dataset.from_list(datasets)