mirror of
https://github.com/open-compass/opencompass.git
synced 2025-05-30 16:03:24 +08:00
29 lines
692 B
Python
29 lines
692 B
Python
|
from abc import abstractstaticmethod
|
||
|
|
from typing import Dict, Optional, Union
|
||
|
|
|
||
|
|
from datasets import Dataset, DatasetDict
|
||
|
|
|
||
|
|
from opencompass.openicl import DatasetReader
|
||
|
|
|
||
|
|
|
||
|
|
class BaseDataset:
|
||
|
|
|
||
|
|
def __init__(self, reader_cfg: Optional[Dict] = {}, **kwargs):
|
||
|
|
self.dataset = self.load(**kwargs)
|
||
|
|
self._init_reader(**reader_cfg)
|
||
|
|
|
||
|
|
def _init_reader(self, **kwargs):
|
||
|
|
self.reader = DatasetReader(self.dataset, **kwargs)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def train(self):
|
||
|
|
return self.reader.dataset['train']
|
||
|
|
|
||
|
|
@property
|
||
|
|
def test(self):
|
||
|
|
return self.reader.dataset['test']
|
||
|
|
|
||
|
|
@abstractstaticmethod
|
||
|
|
def load(**kwargs) -> Union[Dataset, DatasetDict]:
|
||
|
|
pass
|