mirror of
https://github.com/open-compass/opencompass.git
synced 2025-05-30 16:03:24 +08:00
34 lines
946 B
Python
34 lines
946 B
Python
from datasets import load_dataset
|
|
|
|
from opencompass.registry import LOAD_DATASET
|
|
|
|
from .base import BaseDataset
|
|
|
|
|
|
@LOAD_DATASET.register_module()
|
|
class Earth_Silver_MCQDataset(BaseDataset):
|
|
|
|
name = 'msearth_mcq'
|
|
|
|
@staticmethod
|
|
def load(path: str, prompt_mode: str = 'zero-shot', **kwargs):
|
|
"""
|
|
Args:
|
|
path : HF 标识, 固定写 'MSEarth/MSEarth_MCQ'
|
|
split: 'train' / 'validation' / 'test'
|
|
prompt_mode: 'zero-shot' 或 'few-shot'
|
|
"""
|
|
dataset = load_dataset(path=path)
|
|
|
|
dataset = dataset.map(lambda item: {
|
|
'question': item['question'],
|
|
'answer': item['answer']
|
|
})
|
|
|
|
if prompt_mode == 'zero-shot':
|
|
return dataset
|
|
elif prompt_mode == 'few-shot':
|
|
raise NotImplementedError('few-shot prompt 尚未实现')
|
|
else:
|
|
raise ValueError(f'Unsupported prompt_mode: {prompt_mode}')
|