wehub-resource-sync
203 lines
8.0 KiB
Python
203 lines
8.0 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import numpy as np
|
|
import os
|
|
import shutil
|
|
from abc import ABC, abstractmethod
|
|
from copy import deepcopy
|
|
from dataclasses import dataclass, field
|
|
from datasets import Dataset as HfDataset
|
|
from datasets import concatenate_datasets, interleave_datasets
|
|
from modelscope.hub.api import ModelScopeConfig
|
|
from modelscope.utils.config_ds import MS_CACHE_HOME
|
|
from tempfile import TemporaryDirectory
|
|
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union
|
|
|
|
from swift.utils import download_ms_file, get_logger, get_seed, safe_ddp_context
|
|
from .preprocessor import DATASET_TYPE, AutoPreprocessor
|
|
from .utils import sample_dataset
|
|
|
|
PreprocessFunc = Callable[..., DATASET_TYPE]
|
|
logger = get_logger()
|
|
|
|
if TYPE_CHECKING:
|
|
from .dataset_syntax import DatasetSyntax
|
|
|
|
|
|
@dataclass
|
|
class SubsetDataset:
|
|
# `Name` is used for matching subsets of the dataset, and `subset` refers to the subset_name on the hub.
|
|
name: Optional[str] = None
|
|
# If set to None, then subset is set to subset_name.
|
|
subset: str = 'default'
|
|
|
|
# Higher priority. If set to None, the attributes of the DatasetMeta will be used.
|
|
split: Optional[List[str]] = None
|
|
preprocess_func: Optional[PreprocessFunc] = None
|
|
|
|
# If the dataset specifies "all," weak subsets will be skipped.
|
|
is_weak_subset: bool = False
|
|
|
|
def __post_init__(self):
|
|
if self.name is None:
|
|
self.name = self.subset
|
|
|
|
def set_default(self, dataset_meta: 'DatasetMeta') -> 'SubsetDataset':
|
|
subset_dataset = deepcopy(self)
|
|
for k in ['split', 'preprocess_func']:
|
|
v = getattr(subset_dataset, k)
|
|
if v is None:
|
|
setattr(subset_dataset, k, deepcopy(getattr(dataset_meta, k)))
|
|
return subset_dataset
|
|
|
|
|
|
class BaseDatasetLoader(ABC):
|
|
|
|
@abstractmethod
|
|
def load(
|
|
self,
|
|
dataset_syntax: Optional['DatasetSyntax'] = None,
|
|
dataset_meta: Optional['DatasetMeta'] = None,
|
|
*,
|
|
use_hf: Optional[bool] = None,
|
|
) -> HfDataset:
|
|
pass
|
|
|
|
@staticmethod
|
|
def download_ms_dataset(ms_dataset_id: str, files: List[str], force_download: bool = False) -> str:
|
|
"""Download dataset from repo manually
|
|
Args:
|
|
ms_dataset_id: The dataset id of ModelScope
|
|
files: Which files to download
|
|
force_download: Force download or not
|
|
Returns:
|
|
The dataset dir
|
|
"""
|
|
assert isinstance(files, list)
|
|
url = f'http://www.modelscope.cn/api/v1/datasets/{ms_dataset_id}/repo?Revision=master&FilePath={{fpath}}'
|
|
cache_dir = os.path.join(MS_CACHE_HOME, 'datasets', ms_dataset_id, 'master')
|
|
local_dir = os.path.join(cache_dir, 'raw')
|
|
tmp_dir = os.path.join(cache_dir, 'tmp')
|
|
os.makedirs(local_dir, exist_ok=True)
|
|
os.makedirs(tmp_dir, exist_ok=True)
|
|
cookies = ModelScopeConfig.get_cookies()
|
|
with TemporaryDirectory(dir=tmp_dir) as temp_dir:
|
|
for remote_fpath in files:
|
|
url = url.format(fpath=remote_fpath)
|
|
temp_fpath = os.path.join(temp_dir, remote_fpath)
|
|
local_fpath = os.path.join(local_dir, remote_fpath)
|
|
if not force_download and os.path.exists(local_fpath):
|
|
continue
|
|
download_ms_file(url, temp_fpath, cookies)
|
|
shutil.copy2(temp_fpath, local_fpath)
|
|
|
|
return local_dir
|
|
|
|
@staticmethod
|
|
def concat_datasets(datasets: List[HfDataset]) -> Optional[HfDataset]:
|
|
if len(datasets) == 0:
|
|
return
|
|
if len(datasets) == 1:
|
|
return datasets[0]
|
|
return concatenate_datasets(datasets)
|
|
|
|
@staticmethod
|
|
def interleave_datasets(datasets, *args, **kwargs):
|
|
if len(datasets) == 0:
|
|
return
|
|
if len(datasets) == 1:
|
|
return datasets[0]
|
|
return interleave_datasets(datasets, *args, **kwargs)
|
|
|
|
@staticmethod
|
|
def shuffle_dataset(dataset, seed: int, buffer_size: int = 1000):
|
|
if isinstance(dataset, HfDataset):
|
|
with safe_ddp_context(None, True):
|
|
return dataset.shuffle(seed=seed)
|
|
else:
|
|
return dataset.shuffle(seed=seed, buffer_size=buffer_size)
|
|
|
|
@staticmethod
|
|
def post_process(
|
|
train_dataset: DATASET_TYPE,
|
|
*,
|
|
dataset_sample: Optional[int] = None,
|
|
split_dataset_ratio: float = 0.,
|
|
streaming: bool = False,
|
|
shuffle: bool = True,
|
|
random_state: Optional[np.random.RandomState] = None,
|
|
) -> Tuple[DATASET_TYPE, Optional[DATASET_TYPE]]:
|
|
"""Split into train/val datasets and perform dataset sampling."""
|
|
assert dataset_sample is None or dataset_sample > 0
|
|
assert 0 <= split_dataset_ratio <= 1
|
|
if streaming:
|
|
if dataset_sample is None:
|
|
if split_dataset_ratio == 0:
|
|
val_dataset = None
|
|
elif split_dataset_ratio == 1:
|
|
train_dataset, val_dataset = None, train_dataset
|
|
else:
|
|
raise ValueError('The IterableDataset does not support splitting the training set '
|
|
'and validation set when dataset_sample is None.')
|
|
else:
|
|
# not shuffle
|
|
train_dataset = train_dataset.take(dataset_sample)
|
|
val_sample = int(dataset_sample * split_dataset_ratio)
|
|
val_dataset = None if val_sample == 0 else train_dataset.take(val_sample)
|
|
if val_sample:
|
|
train_dataset = train_dataset.skip(val_sample)
|
|
else:
|
|
if dataset_sample is None:
|
|
dataset_sample = len(train_dataset)
|
|
if split_dataset_ratio == 0:
|
|
train_dataset = sample_dataset(train_dataset, dataset_sample, shuffle, random_state)
|
|
val_dataset = None
|
|
elif split_dataset_ratio == 1:
|
|
train_dataset, val_dataset = None, train_dataset
|
|
val_sample = dataset_sample
|
|
# Avoid duplication in the val_dataset.
|
|
assert val_sample <= len(val_dataset), f'val_sample: {val_sample}, len(val_dataset): {len(val_dataset)}'
|
|
val_dataset = sample_dataset(val_dataset, val_sample, shuffle, random_state)
|
|
else:
|
|
# Avoid duplication in the val_dataset.
|
|
train_len = min(len(train_dataset), dataset_sample)
|
|
val_sample = max(int(train_len * split_dataset_ratio), 1)
|
|
train_sample = dataset_sample - val_sample
|
|
assert train_sample > 0
|
|
with safe_ddp_context(None, True):
|
|
train_dataset, val_dataset = train_dataset.train_test_split(
|
|
test_size=val_sample, shuffle=shuffle, seed=get_seed(random_state)).values()
|
|
train_dataset = sample_dataset(train_dataset, train_sample, shuffle, random_state)
|
|
return train_dataset, val_dataset
|
|
|
|
|
|
@dataclass
|
|
class DatasetMeta:
|
|
ms_dataset_id: Optional[str] = None
|
|
hf_dataset_id: Optional[str] = None
|
|
dataset_path: Optional[str] = None # or dataset_dir
|
|
dataset_name: Optional[str] = None
|
|
ms_revision: Optional[str] = None
|
|
hf_revision: Optional[str] = None
|
|
|
|
subsets: List[Union[SubsetDataset, str]] = field(default_factory=lambda: ['default'])
|
|
# Applicable to all subsets.
|
|
split: List[str] = field(default_factory=lambda: ['train'])
|
|
# First perform column mapping, then proceed with the preprocess_func.
|
|
preprocess_func: PreprocessFunc = field(default_factory=lambda: AutoPreprocessor())
|
|
loader: Optional[BaseDatasetLoader] = None
|
|
|
|
tags: List[str] = field(default_factory=list)
|
|
help: Optional[str] = None
|
|
huge_dataset: bool = False
|
|
|
|
def __post_init__(self):
|
|
from .loader import DatasetLoader
|
|
if self.loader is None:
|
|
self.loader = DatasetLoader
|
|
for i, subset in enumerate(self.subsets):
|
|
if isinstance(subset, str):
|
|
self.subsets[i] = SubsetDataset(subset=subset)
|
|
|
|
|
|
DATASET_MAPPING: Dict[Tuple[str, str, str], DatasetMeta] = {}
|