diff --git a/modelscope/msdatasets/ms_dataset.py b/modelscope/msdatasets/ms_dataset.py index 7d99a7cb..cb92a5e6 100644 --- a/modelscope/msdatasets/ms_dataset.py +++ b/modelscope/msdatasets/ms_dataset.py @@ -21,8 +21,9 @@ from modelscope.msdatasets.dataset_cls import (ExternalDataset, from modelscope.msdatasets.dataset_cls.custom_datasets.builder import \ build_custom_dataset from modelscope.msdatasets.utils.delete_utils import DatasetDeleteManager -from modelscope.msdatasets.utils.hf_datasets_util import \ - load_dataset as hf_load_dataset_wrapper +# from modelscope.msdatasets.utils.hf_datasets_util import \ +# load_dataset as hf_load_dataset_wrapper +from modelscope.msdatasets.utils.hf_datasets_util import load_dataset_with_ctx from modelscope.msdatasets.utils.upload_utils import DatasetUploadManager from modelscope.preprocessors import build_preprocessor from modelscope.utils.config import Config, ConfigDict @@ -293,21 +294,25 @@ class MsDataset: # Load from the ModelScope Hub for type=4 (general) if str(dataset_type) == str(DatasetFormations.general.value): - return hf_load_dataset_wrapper( - path=namespace + '/' + dataset_name, - name=subset_name, - data_dir=data_dir, - data_files=data_files, - split=split, - cache_dir=cache_dir, - features=None, - download_config=None, - download_mode=download_mode.value, - revision=version, - token=token, - streaming=use_streaming, - dataset_info_only=dataset_info_only, - **config_kwargs) + + with load_dataset_with_ctx( + path=namespace + '/' + dataset_name, + name=subset_name, + data_dir=data_dir, + data_files=data_files, + split=split, + cache_dir=cache_dir, + features=None, + download_config=None, + download_mode=download_mode.value, + revision=version, + token=token, + streaming=use_streaming, + dataset_info_only=dataset_info_only, + **config_kwargs) as dataset_res: + + return dataset_res + else: remote_dataloader_manager = RemoteDataLoaderManager( diff --git a/modelscope/msdatasets/utils/hf_datasets_util.py b/modelscope/msdatasets/utils/hf_datasets_util.py index 8b067fda..fca641fc 100644 --- a/modelscope/msdatasets/utils/hf_datasets_util.py +++ b/modelscope/msdatasets/utils/hf_datasets_util.py @@ -2,6 +2,7 @@ # Copyright (c) Alibaba, Inc. and its affiliates. # Copyright 2020 The HuggingFace Datasets Authors and the TensorFlow Datasets Authors. import importlib +import contextlib import os import warnings from functools import partial @@ -52,7 +53,7 @@ from datasets.utils.track import tracked_str from fsspec import filesystem from fsspec.core import _un_chain from fsspec.utils import stringify_path -from huggingface_hub import (DatasetCard, DatasetCardData, HfFileSystem) +from huggingface_hub import (DatasetCard, DatasetCardData) from huggingface_hub.hf_api import DatasetInfo as HfDatasetInfo from huggingface_hub.hf_api import HfApi, RepoFile, RepoFolder from packaging import version @@ -66,14 +67,8 @@ from modelscope.utils.logger import get_logger logger = get_logger() -config.HF_ENDPOINT = get_endpoint() - -file_utils.get_from_cache = get_from_cache_ms - - -def _download(self, url_or_filename: str, - download_config: DownloadConfig) -> str: +def _download_ms(self, url_or_filename: str, download_config: DownloadConfig) -> str: url_or_filename = str(url_or_filename) # for temp val revision = None @@ -94,9 +89,6 @@ def _download(self, url_or_filename: str, return out -DownloadManager._download = _download - - def _dataset_info( self, repo_id: str, @@ -193,9 +185,6 @@ def _dataset_info( return HfDatasetInfo(**data) -HfApi.dataset_info = _dataset_info - - def _list_repo_tree( self, repo_id: str, @@ -244,9 +233,6 @@ def _list_repo_tree( **path_info) -HfApi.list_repo_tree = _list_repo_tree - - def _get_paths_info( self, repo_id: str, @@ -282,9 +268,6 @@ def _get_paths_info( ] -HfApi.get_paths_info = _get_paths_info - - def get_fs_token_paths( urlpath, storage_options=None, @@ -420,9 +403,6 @@ def _resolve_pattern( return out -data_files.resolve_pattern = _resolve_pattern - - def _get_data_patterns( base_path: str, download_config: Optional[DownloadConfig] = None) -> Dict[str, @@ -668,9 +648,6 @@ def get_module_without_script(self) -> DatasetModule: ) -HubDatasetModuleFactoryWithoutScript.get_module = get_module_without_script - - def _download_additional_modules( name: str, dataset_name: str, @@ -863,9 +840,6 @@ def get_module_with_script(self) -> DatasetModule: return DatasetModule(module_path, hash, builder_kwargs) -HubDatasetModuleFactoryWithScript.get_module = get_module_with_script - - class DatasetsWrapperHF: @staticmethod @@ -1336,4 +1310,40 @@ class DatasetsWrapperHF: f'any data file in the same directory.') -load_dataset = DatasetsWrapperHF.load_dataset +@contextlib.contextmanager +def load_dataset_with_ctx(*args, **kwargs): + hf_endpoint_origin = config.HF_ENDPOINT + get_from_cache_origin = file_utils.get_from_cache + _download_origin = DownloadManager._download + dataset_info_origin = HfApi.dataset_info + list_repo_tree_origin = HfApi.list_repo_tree + get_paths_info_origin = HfApi.get_paths_info + resolve_pattern_origin = data_files.resolve_pattern + get_module_without_script_origin = HubDatasetModuleFactoryWithoutScript.get_module + get_module_with_script_origin = HubDatasetModuleFactoryWithScript.get_module + + config.HF_ENDPOINT = get_endpoint() + file_utils.get_from_cache = get_from_cache_ms + DownloadManager._download = _download_ms + HfApi.dataset_info = _dataset_info + HfApi.list_repo_tree = _list_repo_tree + HfApi.get_paths_info = _get_paths_info + data_files.resolve_pattern = _resolve_pattern + HubDatasetModuleFactoryWithoutScript.get_module = get_module_without_script + HubDatasetModuleFactoryWithScript.get_module = get_module_with_script + + try: + dataset_res = DatasetsWrapperHF.load_dataset(*args, **kwargs) + yield dataset_res + finally: + config.HF_ENDPOINT = hf_endpoint_origin + file_utils.get_from_cache = get_from_cache_origin + DownloadManager._download = _download_origin + HfApi.dataset_info = dataset_info_origin + HfApi.list_repo_tree = list_repo_tree_origin + HfApi.get_paths_info = get_paths_info_origin + data_files.resolve_pattern = resolve_pattern_origin + HubDatasetModuleFactoryWithoutScript.get_module = get_module_without_script_origin + HubDatasetModuleFactoryWithScript.get_module = get_module_with_script_origin + + logger.info('Context manager of ms-dataset exited.') diff --git a/modelscope/msdatasets/utils/hf_file_utils.py b/modelscope/msdatasets/utils/hf_file_utils.py index fea2506a..f1a4f1f7 100644 --- a/modelscope/msdatasets/utils/hf_file_utils.py +++ b/modelscope/msdatasets/utils/hf_file_utils.py @@ -20,7 +20,7 @@ from filelock import FileLock from modelscope.utils.config_ds import MS_DATASETS_CACHE from modelscope.utils.logger import get_logger -from modelscope.hub.api import HubApi, ModelScopeConfig +from modelscope.hub.api import ModelScopeConfig logger = get_logger() diff --git a/requirements/framework.txt b/requirements/framework.txt index d4987429..8dfbd912 100644 --- a/requirements/framework.txt +++ b/requirements/framework.txt @@ -1,6 +1,6 @@ addict attrs -datasets>=2.14.5 +datasets>=2.16.0 einops filelock>=3.3.0 gast>=0.2.2