wehub-resource-sync
154 lines
5.6 KiB
Python
154 lines
5.6 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import inspect
|
|
import numpy as np
|
|
import os
|
|
import tempfile
|
|
from datasets import Dataset as HfDataset
|
|
from modelscope.hub.utils.utils import get_cache_dir
|
|
from torch.utils.data import Dataset
|
|
from typing import Any, Callable, Dict, Optional, Union
|
|
|
|
from swift.template import MaxLengthError, Template
|
|
from swift.utils import get_logger
|
|
from .preprocessor import RowPreprocessor
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
def sample_dataset(
|
|
dataset: HfDataset,
|
|
dataset_sample: Optional[int],
|
|
shuffle: bool = True,
|
|
random_state: Optional[np.random.RandomState] = None,
|
|
shuffle_all: bool = False, # For compatibility, this defaults to False.
|
|
) -> HfDataset:
|
|
"""Sample dataset by a dataset_sample number
|
|
Args:
|
|
dataset: The dataset instance, iterable dataset is not supported
|
|
dataset_sample: The sample number
|
|
shuffle: Whether to perform random sampling on non-streaming datasets
|
|
random_state: The random state
|
|
Returns:
|
|
The sampled dataset
|
|
"""
|
|
if dataset_sample is None:
|
|
return dataset
|
|
|
|
n_repeat_sample = dataset_sample // len(dataset)
|
|
n_remain_sample = dataset_sample % len(dataset)
|
|
if n_repeat_sample >= 1 and n_remain_sample >= 1:
|
|
logger.warning(f'dataset_sample:{dataset_sample} is greater than len(dataset):{len(dataset)}, '
|
|
'repeated sampling will be performed.')
|
|
idx = np.tile(range(len(dataset)), n_repeat_sample)
|
|
if random_state is None:
|
|
random_state = np.random.RandomState()
|
|
if n_remain_sample >= 1:
|
|
if shuffle:
|
|
idx_remain = random_state.permutation(len(dataset))[:n_remain_sample]
|
|
else:
|
|
idx_remain = np.arange(n_remain_sample)
|
|
idx = np.concatenate([idx, idx_remain])
|
|
if n_repeat_sample >= 1 and shuffle and shuffle_all:
|
|
random_state.shuffle(idx)
|
|
dataset = dataset.select(idx)
|
|
return dataset
|
|
|
|
|
|
class LazyLLMDataset(Dataset):
|
|
"""This class if used to lazy tokenize the dataset, and skips bad ones when training"""
|
|
|
|
def __init__(self,
|
|
dataset: HfDataset,
|
|
encode_func: Callable[[Dict[str, Any]], Dict[str, Any]],
|
|
*,
|
|
n_try_fetch: int = 10,
|
|
strict: bool = False,
|
|
random_state: Optional[Union[np.random.RandomState, int]] = None,
|
|
traceback_limit: int = 10) -> None:
|
|
self.dataset = dataset
|
|
self.encode_func = encode_func
|
|
|
|
n_try_fetch = 1 if strict else min(n_try_fetch, len(self.dataset))
|
|
assert n_try_fetch >= 1
|
|
self.strict = strict
|
|
self.n_try_fetch = n_try_fetch
|
|
|
|
if not isinstance(random_state, np.random.RandomState):
|
|
random_state = np.random.RandomState(random_state)
|
|
self.random_state = random_state
|
|
|
|
self.traceback_limit = traceback_limit
|
|
self._traceback_counter = 0
|
|
self._idx = 0
|
|
self._idx_list = self.random_state.permutation(len(self.dataset)).tolist()
|
|
|
|
def __getitem__(self, idx: int) -> Dict[str, Any]:
|
|
if isinstance(idx, str):
|
|
return self.dataset[idx]
|
|
for i in range(self.n_try_fetch):
|
|
if i > 0:
|
|
idx = self._idx_list[self._idx]
|
|
self._idx = (self._idx + 1) % len(self.dataset)
|
|
data = self.dataset[idx]
|
|
try:
|
|
return self.encode_func(data, return_length=True)
|
|
except Exception as e:
|
|
if self.strict:
|
|
logger.warning('To avoid errors, you can pass `strict=False`.')
|
|
raise
|
|
if isinstance(e, MaxLengthError):
|
|
continue
|
|
if self.traceback_limit is not None and self._traceback_counter < self.traceback_limit:
|
|
import traceback
|
|
logger.info(traceback.format_exc())
|
|
logger.warning('👆👆👆There are errors in the template.encode, '
|
|
'and another piece of data will be randomly selected.')
|
|
self._traceback_counter += 1
|
|
|
|
raise ValueError('Failed to retrieve the dataset. You can avoid this issue by increasing `max_length` or '
|
|
'modifying the `truncation_strategy`.')
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.dataset)
|
|
|
|
|
|
class EncodePreprocessor(RowPreprocessor):
|
|
|
|
def __init__(self, template: 'Template'):
|
|
super().__init__()
|
|
self.template = template
|
|
|
|
def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
|
return self.template.encode(row, return_length=True)
|
|
|
|
|
|
class AddLengthPreprocessor(EncodePreprocessor):
|
|
|
|
def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
|
encoded = super().preprocess(row)
|
|
row['lengths'] = encoded['lengths']
|
|
return row
|
|
|
|
|
|
TEMP_DIR_POOL = {}
|
|
|
|
|
|
def get_temporary_cache_files_directory(prefix=None):
|
|
if prefix is None:
|
|
import datasets.config
|
|
prefix = datasets.config.TEMP_CACHE_DIR_PREFIX
|
|
if prefix in TEMP_DIR_POOL:
|
|
TEMP_DIR = TEMP_DIR_POOL[prefix]
|
|
else:
|
|
tmp_dir = os.path.join(get_cache_dir(), 'tmp')
|
|
os.makedirs(tmp_dir, exist_ok=True)
|
|
kwargs = {}
|
|
parameters = inspect.signature(tempfile.TemporaryDirectory.__init__).parameters
|
|
if 'ignore_cleanup_errors' in parameters:
|
|
kwargs['ignore_cleanup_errors'] = True
|
|
TEMP_DIR = tempfile.TemporaryDirectory(prefix=prefix, dir=tmp_dir, **kwargs)
|
|
logger.info(f'create tmp_dir: {TEMP_DIR.name}')
|
|
TEMP_DIR_POOL[prefix] = TEMP_DIR
|
|
|
|
return TEMP_DIR.name
|