File size: 1,205 Bytes
256a159 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
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)
|