OpenCompass/opencompass/datasets/PubMedQA.py

from datasets import Dataset, DatasetDict, load_dataset

from opencompass.registry import LOAD_DATASET

from .base import BaseDataset


@LOAD_DATASET.register_module()
class PubMedQADataset(BaseDataset):

    @staticmethod
    def load_single():
        dataset = []
        ds = load_dataset('qiaojin/PubMedQA', 'pqa_labeled')
        for data in ds['train']:
            data['question'] = (f"CONTEXTS: {data['context']}\n"
                                f"QUESTION: {data['question']}")
            choices = 'A. yes\nB. no\nC. maybe'
            data['choices'] = choices
            if data['final_decision'] == 'yes':
                data['label'] = 'A. yes'
            elif data['final_decision'] == 'no':
                data['label'] = 'B. no'
            else:
                data['label'] = 'C. maybe'

            dataset.append(data)

        return Dataset.from_list(dataset)

    @staticmethod
    def load(path):
        train_dataset = Dataset.from_list([])
        val_dataset = PubMedQADataset.load_single()
        dataset = DatasetDict({
            'train': train_dataset,
            'validation': val_dataset
        })
        return dataset