File size: 835 Bytes
4947683 | 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 | #!/usr/bin/python
# -*- coding:utf-8 -*-
from .mmap_dataset import MMAPDataset
from .dataset_wrapper import MixDatasetWrapper, PretrainDatasetWrapper
from .dataset_prot_sample import ProtSampleDataset
import utils.register as R
def create_dataset(config: dict):
splits = []
for split_name in ['train', 'valid', 'test']:
split_config = config.get(split_name, None)
if split_config is None:
splits.append(None)
continue
if isinstance(split_config, dict) and 'datasets' in split_config:
dataset = PretrainDatasetWrapper(split_config['max_n_vertex'],
*[R.construct(cfg) for cfg in split_config['datasets']]
)
else:
dataset = R.construct(split_config)
splits.append(dataset)
return splits # train/valid/test
|