OpenCompass/opencompass/datasets/winogrande.py

import json
import os

from datasets import Dataset

from opencompass.registry import LOAD_DATASET

from .base import BaseDataset


@LOAD_DATASET.register_module()
class winograndeDataset(BaseDataset):
    """Disconnect from Huggingface, winograndeDataset."""

    @staticmethod
    def load(path):
        path = os.path.join(path, 'dev.jsonl')
        dataset_list = []
        with open(path, 'r', encoding='utf-8') as f:
            for line in f:
                line = json.loads(line)
                prompt = line['sentence']
                dataset_list.append({
                    'opt1':
                    prompt.replace('_', line['option1']),
                    'opt2':
                    prompt.replace('_', line['option2']),
                    'answer':
                    line['answer']
                })
        dataset_list = Dataset.from_list(dataset_list)
        return dataset_list


@LOAD_DATASET.register_module()
class winograndeDataset_V2(BaseDataset):
    """Disconnect from Huggingface, winograndeDataset_V2."""

    @staticmethod
    def load(path):
        path = os.path.join(path, 'dev.jsonl')
        dataset_list = []
        with open(path, 'r', encoding='utf-8') as f:
            for line in f:
                line = json.loads(line)
                prompt = line['sentence']
                answer = line['answer']
                answer = ' AB'[int(answer)] if answer != '' else 'NULL'
                dataset_list.append({
                    'opt1':
                    prompt.replace('_', line['option1']),
                    'opt2':
                    prompt.replace('_', line['option2']),
                    'answer':
                    answer
                })
        dataset_list = Dataset.from_list(dataset_list)
        return dataset_list