Files
trainlib/trainlib/datasets/__init__.py

13 lines
312 B
Python

from trainlib.dataset import BatchedDataset
from trainlib.utils.text import camel_to_snake
from trainlib.datasets.memory import RandomXYDataset
_datasets = [
RandomXYDataset,
]
dataset_map: dict[str, type[BatchedDataset]] = {
camel_to_snake(_dataset.__name__): _dataset
for _dataset in _datasets
}