146 lines
8.8 KiB
Python
146 lines
8.8 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Literal, Optional, Union
|
|
|
|
from swift.dataset import register_dataset_info
|
|
from swift.utils import get_logger, json_parse_to_dict
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
@dataclass
|
|
class DataArguments:
|
|
"""Holds arguments related to dataset handling and processing.
|
|
|
|
Args:
|
|
dataset (List[str]): A list of dataset IDs or paths. Defaults to [].
|
|
Format for each dataset: 'dataset_id_or_path:subset#count'. Both subset and count are optional.
|
|
- Subsets: Only effective for dataset IDs or folders. Use '/' to select multiple subsets (e.g.,
|
|
'dataset_id:subset1/subset2') or 'all' to select all registered subsets. If only one subset is
|
|
registered, it will be used by default; otherwise, 'default' is the default.
|
|
- Sampling Count: By default, the full dataset is used. Use '#count' to sample. If count <
|
|
total samples, it performs random sampling without replacement. If count > total, it repeats
|
|
the full dataset `count // total` times and then randomly samples an additional `count % total`
|
|
samples. Note: Streaming datasets or setting `--dataset_shuffle false` will result in sequential
|
|
sampling.
|
|
- Local datasets: Supports formats like jsonl, csv, json, and folders.
|
|
val_dataset (List[str]): A list of validation dataset IDs or paths. Defaults to [].
|
|
cached_dataset (List[str]): Use cached datasets to avoid GPU time being occupied by tokenization during
|
|
training/inference on large datasets. This parameter is used to set the folder path(s) of
|
|
cached training datasets, and defaults to `[]`.
|
|
This is generated by the `swift export --to_cached_dataset true ...` command.
|
|
ms-swift only stores an extra 'length' field and filters out erroneous samples
|
|
to reduce storage. Actual preprocessing happens concurrently with training.
|
|
cached_val_dataset (List[str]): Folder path(s) for cached validation datasets, default is [].
|
|
split_dataset_ratio (float): The ratio to split from the training set for validation if `val_dataset` is not
|
|
provided. Defaults to 0.0. Note: The default was 0.01 in `ms-swift<3.6`.
|
|
data_seed (int): The random seed for dataset shuffling. Defaults to 42.
|
|
dataset_num_proc (int): The number of processes to use for dataset preprocessing. Defaults to 1.
|
|
load_from_cache_file (bool): Whether to load the dataset from cache files. Recommended to set to `True` during
|
|
actual runs and `False` during debugging. Defaults to False.
|
|
Note: The default was `True` in `ms-swift<3.9`.
|
|
dataset_shuffle (bool): Whether to shuffle the training dataset. Defaults to True.
|
|
Note: For CPT/SFT, shuffling occurs at both the dataset level (controlled by this flag) and the dataloader
|
|
level.
|
|
val_dataset_shuffle (bool): Whether to shuffle the validation dataset. Defaults to False.
|
|
streaming (bool): Enables streaming to read and process the dataset on-the-fly. `--max_steps` must be set as the
|
|
dataset length is unknown. This allows preprocessing to overlap with training but can become a bottleneck
|
|
with a large `world_size` as preprocessing only runs on rank 0. Defaults to False.
|
|
interleave_prob (Optional[List[float]]): If set, combines datasets using `interleave_datasets` with the
|
|
provided probabilities instead of `concatenate_datasets`. Typically used for streaming. Defaults to None.
|
|
stopping_strategy (str): The stopping strategy for `interleave_datasets`. Can be "first_exhausted" or
|
|
"all_exhausted". Defaults to "first_exhausted".
|
|
shuffle_buffer_size (int): The buffer size for shuffling in streaming mode. Only effective if `dataset_shuffle`
|
|
is `True`. Defaults to 1000.
|
|
download_mode (str): The dataset download mode. Options are 'reuse_dataset_if_exists' and 'force_redownload'.
|
|
Defaults to 'reuse_dataset_if_exists'.
|
|
columns (Optional[str]): A JSON string for column mapping to fit the format required by `AutoPreprocessor`.
|
|
Example: '{"text1": "query", "text2": "response"}'. Defaults to None.
|
|
strict (bool): If `True`, raises an error on any problematic data row. If `False`, discards the problematic
|
|
sample and continues. Typically used for debugging. Defaults to False.
|
|
remove_unused_columns (bool): Whether to remove columns not used by the model. If `False`, extra columns are
|
|
passed to the trainer's `compute_loss` function, which is useful for custom loss calculations.
|
|
Defaults to True. Note: The default is `False` for GPRO.
|
|
disable_auto_column_mapping (bool): 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 (Optional[List[str]]): For self-cognition tasks, replaces the `{{NAME}}` placeholder in the
|
|
`swift/self-cognition` dataset. Pass Chinese and English names.
|
|
Example: `--model_name 小黄 'Xiao Huang'`. Defaults to None.
|
|
model_author (Optional[List[str]]): For self-cognition tasks, replaces the `{{AUTHOR}}` placeholder in the
|
|
`swift/self-cognition` dataset. Pass author's Chinese and English names.
|
|
Example: `--model_author '魔搭' 'ModelScope'`. Defaults to None.
|
|
custom_dataset_info (List[str]): Path to a custom dataset registration JSON file. Defaults to [].
|
|
"""
|
|
# dataset_id or dataset_dir or dataset_path
|
|
dataset: List[str] = field(default_factory=list)
|
|
val_dataset: List[str] = field(default_factory=list)
|
|
cached_dataset: List[str] = field(default_factory=list)
|
|
cached_val_dataset: List[str] = field(default_factory=list)
|
|
split_dataset_ratio: float = 0.
|
|
|
|
data_seed: int = 42
|
|
dataset_num_proc: int = 1
|
|
load_from_cache_file: bool = False
|
|
dataset_shuffle: bool = True
|
|
val_dataset_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
|
|
|
|
download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists'
|
|
columns: Optional[Union[dict, str]] = None
|
|
strict: bool = False
|
|
remove_unused_columns: bool = True
|
|
disable_auto_column_mapping: bool = False
|
|
# Chinese name and English name
|
|
model_name: Optional[List[str]] = field(default=None, metadata={'help': "e.g. ['小黄', 'Xiao Huang']"})
|
|
model_author: Optional[List[str]] = field(default=None, metadata={'help': "e.g. ['魔搭', 'ModelScope']"})
|
|
|
|
custom_dataset_info: List[str] = field(default_factory=list) # .json
|
|
|
|
def _init_custom_dataset_info(self):
|
|
"""register custom dataset_info.json to datasets"""
|
|
if isinstance(self.custom_dataset_info, str):
|
|
self.custom_dataset_info = [self.custom_dataset_info]
|
|
for path in self.custom_dataset_info:
|
|
register_dataset_info(path)
|
|
|
|
def __post_init__(self):
|
|
self.columns = json_parse_to_dict(self.columns)
|
|
if len(self.val_dataset) > 0 or self.streaming and self.split_dataset_ratio > 0:
|
|
self.split_dataset_ratio = 0.
|
|
if len(self.val_dataset) > 0:
|
|
msg = 'len(args.val_dataset) > 0'
|
|
else:
|
|
msg = 'args.streaming is True'
|
|
logger.info(f'Because {msg}, setting split_dataset_ratio: {self.split_dataset_ratio}')
|
|
self._init_custom_dataset_info()
|
|
if isinstance(self.cached_dataset, str):
|
|
self.cached_dataset = [self.cached_dataset]
|
|
self._init_val_dataset_exists()
|
|
|
|
def _init_val_dataset_exists(self):
|
|
self._val_dataset_exists = bool(self.dataset and self.split_dataset_ratio > 0 or self.val_dataset
|
|
or self.cached_val_dataset)
|
|
|
|
def get_dataset_kwargs(self):
|
|
return {
|
|
'seed': self.data_seed,
|
|
'num_proc': self.dataset_num_proc,
|
|
'load_from_cache_file': self.load_from_cache_file,
|
|
'streaming': self.streaming,
|
|
'interleave_prob': self.interleave_prob,
|
|
'stopping_strategy': self.stopping_strategy,
|
|
'shuffle_buffer_size': self.shuffle_buffer_size,
|
|
'use_hf': self.use_hf,
|
|
'hub_token': self.hub_token,
|
|
'download_mode': self.download_mode,
|
|
'columns': self.columns,
|
|
'strict': self.strict,
|
|
'model_name': self.model_name,
|
|
'model_author': self.model_author,
|
|
'remove_unused_columns': self.remove_unused_columns,
|
|
'disable_auto_column_mapping': self.disable_auto_column_mapping,
|
|
}
|