File size: 3,889 Bytes
3270dae | 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 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | """Factory for creating datasets based on configuration."""
from taoTrain.config import TrainingConfig, TrainingModeEnum
from taoTrain.data.pretrain_jsonl import PretrainJSONLDataset
from taoTrain.data.sft_jsonl import SFTJSONLDataset
from taoTrain.data.rl_jsonl import RLJSONLDataset
try:
from taoTrain.data.hf_pretrain import PretrainDataset
from taoTrain.data.hf_sft import SFTDataset
from taoTrain.data.hf_rl import RLDataset
except ImportError:
PretrainDataset = None
SFTDataset = None
RLDataset = None
class DatasetFactory:
"""Factory for creating datasets based on configuration."""
# Registry of dataset classes by mode and backend
DATASETS = {
(TrainingModeEnum.PRETRAIN, "jsonl"): PretrainJSONLDataset,
(TrainingModeEnum.SFT, "jsonl"): SFTJSONLDataset,
(TrainingModeEnum.RL, "jsonl"): RLJSONLDataset,
}
if PretrainDataset is not None:
DATASETS.update({
(TrainingModeEnum.PRETRAIN, "huggingface"): PretrainDataset,
(TrainingModeEnum.SFT, "huggingface"): SFTDataset,
(TrainingModeEnum.RL, "huggingface"): RLDataset,
})
@staticmethod
def create_dataset(
config: TrainingConfig,
split: str = "train",
):
"""
Create dataset instance based on configuration.
Args:
config: Training configuration
split: Dataset split (train, validation, test) - primarily for HuggingFace datasets
Returns:
Dataset instance matching the configured mode and backend
Raises:
ValueError: If configuration is invalid or unsupported mode/backend combination
"""
# Determine backend: JSONL or HuggingFace
if config.dataset.local:
backend = "jsonl"
else:
backend = "huggingface"
# Get mode
mode = config.mode
# Look up dataset class
key = (mode, backend)
if key not in DatasetFactory.DATASETS:
if backend == "huggingface":
raise ImportError(
"HuggingFace dataset support requires the optional 'datasets' dependency. "
"Install project dependencies before using dataset.local=false."
)
raise ValueError(
f"Unsupported dataset configuration: mode={mode.value}, backend={backend}. "
f"Supported: {list(DatasetFactory.DATASETS.keys())}"
)
dataset_class = DatasetFactory.DATASETS[key]
# Instantiate dataset
if backend == "jsonl":
# JSONL datasets don't use split parameter
return dataset_class(config)
else:
# HuggingFace datasets use split parameter
return dataset_class(config, split=split)
@staticmethod
def register_dataset(mode: TrainingModeEnum, backend: str, dataset_class):
"""
Register a custom dataset class.
Args:
mode: Training mode (e.g., TrainingModeEnum.PRETRAIN)
backend: Backend name (e.g., "jsonl", "huggingface")
dataset_class: Dataset class to register
"""
DatasetFactory.DATASETS[(mode, backend)] = dataset_class
@staticmethod
def list_available_datasets():
"""List all available dataset configurations."""
configs = {}
for (mode, backend), dataset_class in DatasetFactory.DATASETS.items():
key = f"{mode.value}_{backend}"
configs[key] = {
"mode": mode.value,
"backend": backend,
"class": dataset_class.__name__,
}
return configs
|