wehub-resource-sync
385 lines
18 KiB
Python
385 lines
18 KiB
Python
# 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
|