# Copyright (c) ModelScope Contributors. All rights reserved. import numpy as np import os from contextlib import nullcontext from datasets import Dataset as HfDataset from datasets import IterableDataset as HfIterableDataset from datasets import load_dataset as hf_load_dataset from functools import partial from modelscope.hub.utils.utils import get_cache_dir from typing import Dict, List, Literal, Optional, Tuple, Union from swift.hub import get_hub from swift.utils import get_logger, get_seed, safe_ddp_context, use_hf_hub from .dataset_meta import DATASET_TYPE, BaseDatasetLoader from .dataset_syntax import DatasetSyntax from .preprocessor import RowPreprocessor from .register import DATASET_MAPPING, DatasetMeta, SubsetDataset logger = get_logger() class DatasetLoader(BaseDatasetLoader): def __init__( self, num_proc: int = 1, load_from_cache_file: bool = True, streaming: bool = False, hub_token: Optional[str] = None, strict: bool = False, download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists', columns: Optional[Dict[str, str]] = None, remove_unused_columns: bool = True, disable_auto_column_mapping: bool = False, ): self.num_proc = num_proc self.load_from_cache_file = load_from_cache_file self.streaming = streaming self.hub_token = hub_token self.strict = strict self.download_mode = download_mode self.columns = columns self.remove_unused_columns = remove_unused_columns self.disable_auto_column_mapping = disable_auto_column_mapping def _load_dataset_path( self, dataset_path: str, dataset_meta: DatasetMeta, ) -> HfDataset: ext = os.path.splitext(dataset_path)[1].lstrip('.') file_type = {'jsonl': 'json', 'txt': 'text'}.get(ext) or ext kwargs = {'split': 'train', 'streaming': self.streaming, 'num_proc': self.num_proc} if file_type == 'csv': kwargs['na_filter'] = False with safe_ddp_context(None, True): kwargs['cache_dir'] = os.path.join(get_cache_dir(), 'datasets') dataset = hf_load_dataset(file_type, data_files=dataset_path, **kwargs) if self.columns: dataset = RowPreprocessor.safe_rename_columns(dataset, self.columns) dataset = dataset_meta.preprocess_func( dataset, num_proc=self.num_proc, load_from_cache_file=self.load_from_cache_file, strict=self.strict, enable_auto_mapping=not self.disable_auto_column_mapping) if self.remove_unused_columns: dataset = RowPreprocessor.remove_useless_columns(dataset) return dataset def _load_repo_dataset( self, dataset_id: str, subset: SubsetDataset, *, use_hf: Optional[bool] = None, revision: Optional[str] = None, ) -> HfDataset: datasets = [] if os.path.isdir(dataset_id): retry = 1 load_context = nullcontext use_hf = True dataset_str = f'Use local folder, dataset_dir: {dataset_id}' # The dataset downloaded from modelscope will have an additional dataset_infos.json file. with safe_ddp_context('dataset_infos_rename'): dataset_infos_path = os.path.join(dataset_id, 'dataset_infos.json') if os.path.isfile(dataset_infos_path): os.rename(dataset_infos_path, f'{dataset_infos_path}_bak') elif dataset_id.startswith('/'): raise ValueError(f'The local path does not exist, dataset_id: `{dataset_id}`. ' f'os.path.exists(dataset_id): {os.path.exists(dataset_id)}') else: retry = 3 load_context = partial(safe_ddp_context, hash_id=dataset_id, use_barrier=True) dataset_str_f = 'Downloading the dataset from {hub}, dataset_id: {dataset_id}' if use_hf: dataset_str = dataset_str_f.format(hub='HuggingFace', dataset_id=dataset_id) else: dataset_str = dataset_str_f.format(hub='ModelScope', dataset_id=dataset_id) logger.info(dataset_str) hub = get_hub(use_hf) for split in subset.split: i = 1 with load_context(): while True: try: dataset = hub.load_dataset( dataset_id, subset.subset, split, streaming=self.streaming, revision=revision, download_mode=self.download_mode, hub_token=self.hub_token, num_proc=self.num_proc) except Exception as e: if i == retry: raise i += 1 logger.error(f'Dataset {dataset_id} load failed: subset_name={subset.subset},' f'split={split} with error: {e}') else: break if hasattr(dataset, '_hf_ds'): dataset = dataset._hf_ds if self.streaming and isinstance(dataset, HfDataset): dataset = dataset.to_iterable_dataset() if self.columns: dataset = RowPreprocessor.safe_rename_columns(dataset, self.columns) dataset = subset.preprocess_func( dataset, num_proc=self.num_proc, load_from_cache_file=self.load_from_cache_file, strict=self.strict, enable_auto_mapping=not self.disable_auto_column_mapping) if self.remove_unused_columns: dataset = RowPreprocessor.remove_useless_columns(dataset) datasets.append(dataset) return self.concat_datasets(datasets) @staticmethod def _select_subsets(subsets: List[str], dataset_meta: DatasetMeta) -> List[SubsetDataset]: subset_mapping = {subset.name: subset for subset in dataset_meta.subsets} subset_names = list(subset_mapping.keys()) if not subsets: if len(subset_names) <= 1: subsets = subset_names elif 'default' in subset_names: subsets = ['default'] else: raise ValueError(f'Please provide subsets. available subsets: {subset_names}') elif len(subsets) == 1 and subsets[0] == 'all' and 'all' not in subset_names: subsets = [subset_name for subset_name in subset_names if not subset_mapping[subset_name].is_weak_subset] subsets = [ subset_mapping[subset_name] if subset_name in subset_mapping else SubsetDataset(subset=subset_name) for subset_name in subsets ] return [subset.set_default(dataset_meta) for subset in subsets] def load( self, dataset_syntax: Optional[DatasetSyntax] = None, dataset_meta: Optional[DatasetMeta] = None, *, use_hf: Optional[bool] = None, ) -> HfDataset: if dataset_syntax.dataset_type == 'path': dataset = self._load_dataset_path( dataset_syntax.dataset, dataset_meta=dataset_meta, ) else: subsets: List[SubsetDataset] = self._select_subsets(dataset_syntax.subsets, dataset_meta) revision = dataset_meta.hf_revision if use_hf else dataset_meta.ms_revision datasets = [] for subset in subsets: dataset = self._load_repo_dataset( dataset_syntax.dataset, subset, use_hf=use_hf, revision=revision, ) datasets.append(dataset) dataset = self.concat_datasets(datasets) return dataset def init_self_cognition_preprocessor( dataset_meta: Optional[DatasetMeta], model_name: Optional[Union[Tuple[str, str], List[str]]] = None, model_author: Optional[Union[Tuple[str, str], List[str]]] = None, ) -> None: from .dataset.llm import SelfCognitionPreprocessor if dataset_meta is None or model_name is None and model_author is None: return kwargs = {} # zh, en for key in ['name', 'author']: val = locals()[f'model_{key}'] if isinstance(val, str): val = [val] if val is not None and val[0] is not None and (len(val) == 1 or val[1] is None): val = (val[0], val[0]) kwargs[key] = val preprocess_funcs = [dataset_meta.preprocess_func] preprocess_funcs += [subset.preprocess_func for subset in dataset_meta.subsets if isinstance(subset, SubsetDataset)] for preprocess_func in preprocess_funcs: if isinstance(preprocess_func, SelfCognitionPreprocessor): preprocess_func.set_name_author(**kwargs) logger.info_once(f"SelfCognitionPreprocessor has been successfully configured with name: {kwargs['name']}, " f"author: {kwargs['author']}.") def _inject_dataset_routing_tag(dataset: DATASET_TYPE, ds_name: str) -> DATASET_TYPE: """Inject ``dataset`` column for multi-teacher routing (constant per source dataset).""" if isinstance(dataset, HfIterableDataset): return dataset.map(lambda example: {**example, 'dataset': ds_name}) return dataset.add_column('dataset', [ds_name] * len(dataset)) def load_dataset( datasets: Union[List[str], str], *, split_dataset_ratio: float = 0., seed: Union[int, np.random.RandomState, None] = 42, num_proc: int = 1, load_from_cache_file: bool = True, shuffle: bool = False, streaming: bool = False, interleave_prob: Optional[List[float]] = None, stopping_strategy: Literal['first_exhausted', 'all_exhausted'] = 'first_exhausted', shuffle_buffer_size: int = 1000, use_hf: Optional[bool] = None, hub_token: Optional[str] = None, strict: bool = False, download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists', columns: Optional[Dict[str, str]] = None, # columns_mapping remove_unused_columns: bool = True, disable_auto_column_mapping: bool = False, # self-cognition model_name: Optional[Union[Tuple[str, str], List[str]]] = None, # zh, en model_author: Optional[Union[Tuple[str, str], List[str]]] = None, ) -> Tuple[DATASET_TYPE, Optional[DATASET_TYPE]]: """Load and preprocess datasets. This function provides a unified interface to load datasets from various sources (HuggingFace, ModelScope, or local paths), with support for splitting, shuffling, streaming, and interleaving multiple datasets. It also handles self-cognition dataset preprocessing for model training. Args: datasets: Single dataset name or list of dataset names to load. Can use special syntax for advanced configurations (e.g., 'dataset_name#1000' for sampling). split_dataset_ratio: Ratio for splitting dataset into train/validation sets. Value between 0 and 1. If 0, no validation split is created. Default: 0. seed: Random seed for reproducibility. Can be an integer or numpy RandomState object. If None, results will be non-deterministic. Default: 42. num_proc: Number of processes to use for dataset preprocessing. Set to None for streaming mode. Default: 1. load_from_cache_file: Whether to load preprocessed data from cache if available. Default: True. shuffle: Whether to shuffle the dataset(s) after loading. Default: False. streaming: Enable streaming mode for large datasets that don't fit in memory. When True, num_proc is automatically set to None. Default: False. interleave_prob: Probability weights for interleaving multiple datasets. Must have same length as datasets list. If None, datasets are concatenated instead. Default: None. stopping_strategy: Strategy when interleaving datasets of different lengths: - 'first_exhausted': Stop when shortest dataset is exhausted - 'all_exhausted': Continue until all datasets are exhausted Default: 'first_exhausted'. shuffle_buffer_size: Buffer size for shuffling in streaming mode. Larger values provide better randomization but use more memory. Default: 1000. use_hf: Force using HuggingFace Hub (True) or ModelScope (False). If None, it is controlled by the environment variable `USE_HF`, which defaults to '0'. Default: None. hub_token: Authentication token for accessing private datasets on the hub. Default: None. strict: If True, raise exceptions when encountering malformed data rows. If False, skip invalid rows with warnings. Default: False. download_mode: How to handle existing cached datasets: - 'reuse_dataset_if_exists': Use cached version if available - 'force_redownload': Always download fresh copy Default: 'reuse_dataset_if_exists'. columns: Manual column name mapping for datasets. Dictionary mapping source column names to target column names (e.g., {'text': 'content'}). Default: None. remove_unused_columns: Whether to remove columns not used in preprocessing. Helps reduce memory usage. Default: True. disable_auto_column_mapping: By default, column names in the dataset are automatically mapped. This parameter disables that behavior (the `columns` parameter remains effective), defaulting to `False`. model_name: Model name for self-cognition task preprocessing. Can be a tuple of (Chinese_name, English_name) or list of names. Default: None. model_author: Model author for self-cognition task preprocessing. Can be a tuple of (Chinese_author, English_author) or list of authors. Default: None. Returns: A tuple of (train_dataset, val_dataset): - train_dataset: The training dataset - val_dataset: The validation dataset if split_dataset_ratio > 0, otherwise None Examples: >>> # Load single dataset >>> train_ds, val_ds = load_dataset('AI-ModelScope/alpaca-gpt4-data-zh', split_dataset_ratio=0.1) >>> # Load multiple datasets >>> train_ds, _ = load_dataset( ... ['AI-ModelScope/alpaca-gpt4-data-zh#500', 'swift/self-cognition#500'], ... model_name=('我的模型', 'MyModel'), ... model_author=('作者', 'Author') ... ) """ init_self_cognition_preprocessor(DATASET_MAPPING.get('self-cognition'), model_name, model_author) if isinstance(datasets, str): datasets = [datasets] if not isinstance(seed, np.random.RandomState): seed = np.random.RandomState(seed) if streaming: num_proc = None train_datasets = [] val_datasets = [] use_hf_default = use_hf if use_hf_default is None: use_hf_default = True if use_hf_hub() else False for dataset in datasets: dataset_syntax = DatasetSyntax.parse(dataset) use_hf = dataset_syntax.use_hf or use_hf_default # compat dataset_name if dataset_syntax.dataset in DATASET_MAPPING: dataset_meta = DATASET_MAPPING[dataset_syntax.dataset] if dataset_syntax.use_hf is None and dataset_meta.dataset_path is not None: dataset_syntax.dataset = dataset_meta.dataset_path dataset_syntax.dataset_type = 'path' else: dataset_syntax.dataset = dataset_meta.hf_dataset_id if use_hf else dataset_meta.ms_dataset_id else: dataset_meta = dataset_syntax.get_dataset_meta(use_hf) loader = dataset_meta.loader( num_proc=num_proc, load_from_cache_file=load_from_cache_file, streaming=streaming, hub_token=hub_token, strict=strict, download_mode=download_mode, columns=columns, # columns_mapping remove_unused_columns=remove_unused_columns, disable_auto_column_mapping=disable_auto_column_mapping, ) train_dataset = loader.load(dataset_syntax, dataset_meta, use_hf=use_hf) train_dataset, val_dataset = loader.post_process( train_dataset, dataset_sample=dataset_syntax.dataset_sample, split_dataset_ratio=split_dataset_ratio, streaming=streaming, shuffle=shuffle, random_state=seed, ) if train_dataset is not None: # Inject dataset_syntax.dataset as routing tag for multi-teacher ds_name = dataset_syntax.dataset train_dataset = _inject_dataset_routing_tag(train_dataset, ds_name) train_datasets.append(train_dataset) if val_dataset is not None: ds_name = dataset_syntax.dataset val_dataset = _inject_dataset_routing_tag(val_dataset, ds_name) val_datasets.append(val_dataset) if interleave_prob is None: train_datasets = loader.concat_datasets(train_datasets) val_datasets = loader.concat_datasets(val_datasets) else: train_datasets = loader.interleave_datasets( train_datasets, interleave_prob, seed=get_seed(seed), stopping_strategy=stopping_strategy) val_datasets = loader.interleave_datasets( val_datasets, interleave_prob, seed=get_seed(seed), stopping_strategy=stopping_strategy) if shuffle: if train_datasets: train_datasets = loader.shuffle_dataset( train_datasets, seed=get_seed(seed), buffer_size=shuffle_buffer_size) if val_datasets: val_datasets = loader.shuffle_dataset(val_datasets, seed=get_seed(seed), buffer_size=shuffle_buffer_size) return train_datasets, val_datasets