|
import json |
|
|
|
from datasets import Dataset, DatasetDict |
|
|
|
from opencompass.registry import LOAD_DATASET |
|
|
|
from .base import BaseDataset |
|
|
|
|
|
@LOAD_DATASET.register_module() |
|
class dropDataset(BaseDataset): |
|
|
|
@staticmethod |
|
def get_answers(validated_answers): |
|
answers = [] |
|
for answer_item in validated_answers: |
|
if answer_item['number']: |
|
answers.append(answer_item['number']) |
|
elif any(answer_item['date'][i] for i in ['day', 'month', 'year']): |
|
d = [answer_item['date'][i] for i in ['day', 'month', 'year']] |
|
answers.append(' '.join(d).strip()) |
|
else: |
|
for span in answer_item['spans']: |
|
answers.append(span) |
|
answers = list(set(answers)) |
|
return answers |
|
|
|
@staticmethod |
|
def load(path, only_number=True): |
|
with open(path, 'r', encoding='utf-8') as f: |
|
lines = json.load(f) |
|
dataset_list = [] |
|
for line in lines.values(): |
|
for qa_pair in line['qa_pairs']: |
|
validated_answers = qa_pair['validated_answers'] |
|
if only_number and not any(i['number'] |
|
for i in validated_answers): |
|
continue |
|
item = { |
|
'prompt': line['passage'], |
|
'question': qa_pair['question'], |
|
'answers': dropDataset.get_answers(validated_answers), |
|
} |
|
dataset_list.append(item) |
|
|
|
dataset_list = Dataset.from_list(dataset_list) |
|
return DatasetDict({'validation': dataset_list}) |
|
|
