File size: 692 Bytes
256a159 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
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
|