modelscope 1.36.3__tar.gz → 1.37.0__tar.gz

{modelscope-1.36.3/modelscope.egg-info → modelscope-1.37.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modelscope
-Version: 1.36.3
+Version: 1.37.0
 Summary: ModelScope: bring the notion of Model-as-a-Service to life.
 Author: ModelScope team
 Author-email: contact@modelscope.cn

{modelscope-1.36.3 → modelscope-1.37.0}/modelscope/hub/api.py

@@ -2316,6 +2316,18 @@ class HubApi:
             commit_message=commit_message,
         )
 
+        # Guard: skip sending empty commits (no effective file changes)
+        if not payload['actions']:
+            logger.info(
+                'Commit skipped: no effective actions in payload '
+                '(all files already exist).')
+            return CommitInfo(
+                commit_url='',
+                commit_message=commit_message,
+                commit_description=commit_description or '',
+                oid='no-op',
+            )
+
         response = self.session.post(
             url,
             headers=self.builder_headers(self.headers),
@@ -2753,8 +2765,9 @@ class HubApi:
         )
         commit_description = commit_description or 'Uploading files'
 
-        # Exclude internal cache/checkpoint files from upload
-        _internal_ignore = [UPLOAD_HASH_CACHE_FILE, _LEGACY_PROGRESS_FILE]
+        # Exclude internal cache/checkpoint files from upload at any directory depth
+        _internal_files = [UPLOAD_HASH_CACHE_FILE, _LEGACY_PROGRESS_FILE]
+        _internal_ignore = [p for f in _internal_files for p in (f, f'*/{f}')]
         if ignore_patterns is None:
             ignore_patterns = _internal_ignore
         elif isinstance(ignore_patterns, str):
@@ -2957,18 +2970,27 @@ class HubApi:
                     except Exception as e:
                         logger.error(
                             f'Batch {batch_idx + 1}/{num_batches} commit failed: {e}')
-                        for r in results:
-                            tracker.mark_failed(
-                                r['file_path_in_repo'], r['file_mtime'],
-                                r['file_size_on_disk'],
-                                error_type='commit_failed')
-                        # Recover uploaded files to retry queue
-                        for r in results:
-                            total_failed_files.append(
-                                ((r['file_path_in_repo'], r['file_path']), e))
-                        logger.warning(
-                            f'Batch {batch_idx + 1}/{num_batches}: '
-                            f'{len(results)} uploaded file(s) recovered to retry queue.')
+                        category = classify_error(e)
+                        if not category.is_retryable:
+                            # Permanent error: mark files as failed, do not retry
+                            for r in results:
+                                tracker.mark_failed(
+                                    r['file_path_in_repo'], r['file_mtime'],
+                                    r['file_size_on_disk'],
+                                    error_type='commit_' + category.value)
+                            logger.error(
+                                f'Batch {batch_idx + 1}/{num_batches}: '
+                                f'permanent failure ({category.value}), '
+                                f'{len(results)} file(s) will not be retried.')
+                        else:
+                            # Transient error: recover to retry queue
+                            for r in results:
+                                total_failed_files.append(
+                                    ((r['file_path_in_repo'], r['file_path']), e))
+                            logger.warning(
+                                f'Batch {batch_idx + 1}/{num_batches}: '
+                                f'{len(results)} file(s) recovered to retry queue '
+                                f'(error_category={category.value}).')
         finally:
             tracker.save()
 
@@ -3030,10 +3052,19 @@ class HubApi:
                         except Exception as e:
                             logger.error(
                                 f'  Retry round {retry_round + 1} commit failed: {e}')
-                            for result in retry_successes:
-                                retry_failures.append(
-                                    ((result['file_path_in_repo'],
-                                      result.get('file_path', '')), e))
+                            category = classify_error(e)
+                            if not category.is_retryable:
+                                for result in retry_successes:
+                                    tracker.mark_failed(
+                                        result['file_path_in_repo'],
+                                        result['file_mtime'],
+                                        result['file_size_on_disk'],
+                                        error_type='commit_' + category.value)
+                            else:
+                                for result in retry_successes:
+                                    retry_failures.append(
+                                        ((result['file_path_in_repo'],
+                                          result.get('file_path', '')), e))
                 total_failed_files = retry_failures
 
         # Final tracker save
@@ -3253,7 +3284,7 @@ class HubApi:
                         repo_type=repo_type,
                         revision=revision)
                     commit_infos.append(commit_info)
-                    # Mark committed
+                    # Mark committed only after successful commit
                     self._track_committed_batch(tracker, batch)
                     logger.info(
                         f'[ReAct] {round_name}: '
@@ -3261,11 +3292,18 @@ class HubApi:
                 except Exception as e:
                     logger.error(
                         f'[ReAct] {round_name} commit failed: {e}')
-                    # Recover uploaded files back to failures
-                    for r in batch:
-                        round_failures.append(
-                            ((r['file_path_in_repo'],
-                              r['file_path']), e))
+                    category = classify_error(e)
+                    if not category.is_retryable:
+                        for r in batch:
+                            tracker.mark_failed(
+                                r['file_path_in_repo'], r['file_mtime'],
+                                r['file_size_on_disk'],
+                                error_type='commit_' + category.value)
+                    else:
+                        for r in batch:
+                            round_failures.append(
+                                ((r['file_path_in_repo'],
+                                  r['file_path']), e))
 
             # OBSERVE: classify new failures, enforce per-file retry limit
             new_retryable = []

{modelscope-1.36.3 → modelscope-1.37.0}/modelscope/models/base/base_model.py

@@ -14,7 +14,8 @@ from modelscope.utils.config import Config, ConfigDict
 from modelscope.utils.constant import DEFAULT_MODEL_REVISION, Invoke, ModelFile
 from modelscope.utils.device import verify_device
 from modelscope.utils.logger import get_logger
-from modelscope.utils.plugins import (register_modelhub_repo,
+from modelscope.utils.plugins import (filter_plugin_in_whitelist,
+                                      register_modelhub_repo,
                                       register_plugins_repo)
 
 logger = get_logger()
@@ -190,7 +191,7 @@ class Model(ABC):
 
         # Security check: Only allow execution of remote code or plugins if trust_remote_code is True
         plugins = cfg.safe_get('plugins')
-        if plugins and not trust_remote_code:
+        if filter_plugin_in_whitelist(plugins) and not trust_remote_code:
             raise RuntimeError(
                 'Detected plugins field in the model configuration file, but '
                 'trust_remote_code=True was not explicitly set.\n'

{modelscope-1.36.3 → modelscope-1.37.0}/modelscope/msdatasets/ms_dataset.py

@@ -246,7 +246,7 @@ class MsDataset:
                 'you can trust the external codes.')
 
         # Raise csv field size limit to avoid errors with large cells
-        if config_kwargs.get('engine') == 'python':
+        if config_kwargs.pop('engine', None) == 'python':
             import csv as csv_module
             import sys
             try:

{modelscope-1.36.3 → modelscope-1.37.0}/modelscope/msdatasets/utils/_module_factories.py

@@ -9,12 +9,12 @@ by :func:`~hf_datasets_util.load_dataset_with_ctx`.
 import importlib
 import inspect
 import os
+import re
 from functools import partial
 from pathlib import Path
-from typing import Dict, List, Optional, Sequence, Tuple, Union
+from typing import Dict, List, Optional, Tuple
 
-from datasets import (BuilderConfig, DownloadConfig, DownloadMode, Features,
-                      Version, config, data_files)
+from datasets import (BuilderConfig, DownloadConfig, config)
 from datasets.data_files import (
     FILES_TO_IGNORE, DataFilesDict, EmptyDatasetError,
     _get_data_files_patterns, _is_inside_unrequested_special_dir,
@@ -24,17 +24,16 @@ from datasets.download.streaming_download_manager import (
     _prepare_path_and_storage_options, xbasename, xjoin)
 from datasets.exceptions import DataFilesNotFoundError
 from datasets.info import DatasetInfosDict
-from datasets.load import (BuilderConfigsParameters, DatasetModule,
+from datasets.load import (BuilderConfigsParameters,
+                           DatasetModule,
                            create_builder_configs_from_metadata_configs,
-                           get_dataset_builder_class, import_main_class,
+                           import_main_class,
                            infer_module_for_data_files)
 from datasets.naming import camelcase_to_snakecase
 from datasets.packaged_modules import (_MODULE_TO_EXTENSIONS,
                                        _PACKAGED_DATASETS_MODULES)
-from datasets.utils.file_utils import (cached_path, is_local_path,
-                                       relative_to_absolute_path)
+from datasets.utils.file_utils import (cached_path, is_local_path)
 from datasets.utils.metadata import MetadataConfigs
-from datasets.utils.track import tracked_str
 from fsspec import filesystem
 from fsspec.core import _un_chain
 from fsspec.utils import stringify_path
@@ -42,14 +41,16 @@ from huggingface_hub import DatasetCard, DatasetCardData
 from packaging import version
 
 from modelscope import HubApi
-from modelscope.msdatasets.utils._compat import (
-    _HAS_SCRIPT_LOADING, _create_importable_file, _get_importable_file_path,
-    _load_importable_file, files_to_hash, get_imports, init_dynamic_modules,
-    resolve_trust_remote_code)
+from modelscope.msdatasets.utils._compat import (_create_importable_file,
+                                                 _get_importable_file_path,
+                                                 _load_importable_file,
+                                                 files_to_hash,
+                                                 get_imports,
+                                                 init_dynamic_modules,
+                                                 resolve_trust_remote_code)
 from modelscope.utils.constant import (DEFAULT_DATASET_REVISION,
                                        REPO_TYPE_DATASET)
 from modelscope.utils.file_utils import is_relative_path
-from modelscope.utils.import_utils import has_attr_in_class
 from modelscope.utils.logger import get_logger
 
 # ALL_ALLOWED_EXTENSIONS moved to datasets.packaged_modules in datasets 4.0
@@ -60,6 +61,73 @@ except ImportError:
 
 logger = get_logger()
 
+
+def _extract_split_names(split):
+    """Extract base split names from a split specification string.
+
+    Handles simple names ("tool"), sliced splits ("train[:100]"),
+    and combined splits ("train+test").
+
+    Args:
+        split: A split specification string, or None.
+
+    Returns:
+        A set of split name strings, or None if *split* is None or
+        cannot be parsed.
+    """
+    if split is None:
+        return None
+    split_str = str(split)
+    parts = split_str.split('+')
+    names = set()
+    for part in parts:
+        # Remove slice notation like "[:100]" or "[50%:]"
+        name = re.split(r'\[', part.strip())[0]
+        if name:
+            names.add(name)
+    return names if names else None
+
+
+def _filter_data_files_by_split(data_files, download_config):
+    """Filter data_files entries to only include the requested split(s).
+
+    Args:
+        data_files: The raw data_files value from metadata_configs.
+            Expected format: list of dicts with 'split' and 'path' keys.
+        download_config: The DownloadConfig instance (may be None).
+
+    Returns:
+        Filtered data_files if a split filter is active; otherwise
+        the original *data_files* unchanged.
+    """
+    # 1. Safely retrieve the split value
+    split_str = None
+    if download_config is not None:
+        storage_opts = getattr(download_config, 'storage_options', None)
+        if isinstance(storage_opts, dict):
+            split_str = storage_opts.get('split')
+
+    if split_str is None:
+        return data_files
+
+    # 2. Parse split names
+    split_names = _extract_split_names(split_str)
+    if not split_names:
+        return data_files
+
+    # 3. Only filter list[dict] format
+    if not isinstance(data_files, list):
+        return data_files
+    if not all(isinstance(item, dict) and 'split' in item for item in data_files):
+        return data_files
+
+    # 4. Filter
+    filtered = [df for df in data_files if df.get('split') in split_names]
+
+    # 5. Fallback: if filtered is empty, return original data
+    return filtered if filtered else data_files
+
+
 # ---------------------------------------------------------------------------
 # Shared HubApi instance (avoids creating a new requests.Session per call)
 # ---------------------------------------------------------------------------
@@ -575,6 +643,8 @@ def get_module_without_script(self) -> DatasetModule:
         else:
             subset_data_files = next(iter(
                 metadata_configs.values()))['data_files']
+        subset_data_files = _filter_data_files_by_split(
+            subset_data_files, self.download_config)
         patterns = sanitize_patterns(subset_data_files)
     else:
         patterns = _get_data_patterns(

{modelscope-1.36.3 → modelscope-1.37.0}/modelscope/msdatasets/utils/hf_datasets_util.py

@@ -12,6 +12,7 @@ Sub-modules:
     _module_factories – dataset module factory functions & data-file resolution
 """
 import contextlib
+import inspect
 import os
 import warnings
 from dataclasses import fields
@@ -25,7 +26,7 @@ from datasets import (Dataset, DatasetBuilder, DatasetDict,
                       DownloadConfig, DownloadManager, DownloadMode, Features,
                       IterableDataset, IterableDatasetDict, Split,
                       VerificationMode, Version, config, data_files, LargeList,
-                      Sequence as SequenceHf)
+                      Sequence as SequenceHf, SplitDict)
 
 try:
     from datasets import List as DatasetList
@@ -407,10 +408,26 @@ def _get_paths_info(
 
 
 # ===================================================================
-# HfFileSystem patch (_hf_fs_open)
+# HfFileSystem patches (_hf_fs_open, _hf_fs_init)
 # ===================================================================
 
 _hf_fs_open_original = None
+_hf_fs_init_original = None
+
+
+def _hf_fs_init_with_cookie(self, *args, endpoint=None, token=None, **kwargs):
+    """HfFileSystem.__init__ wrapper that injects ModelScope cookie auth.
+
+    ModelScope's /resolve/ endpoint authenticates via `m_session_id` cookie
+    rather than the `Authorization: Bearer` header used by HuggingFace Hub.
+    This wrapper ensures the cookie is included in all subsequent HTTP
+    requests made by the HfFileSystem instance.
+    """
+    _hf_fs_init_original(self, *args, endpoint=endpoint, token=token, **kwargs)
+    if token and isinstance(token, str):
+        if not hasattr(self._api, 'headers') or self._api.headers is None:
+            self._api.headers = {}
+        self._api.headers['Cookie'] = f'm_session_id={token}'
 
 
 def _hf_fs_open(self, path, mode='rb', **kwargs):
@@ -456,6 +473,154 @@ def _hf_fs_open(self, path, mode='rb', **kwargs):
     return _hf_fs_open_original(self, path, mode=mode, **kwargs)
 
 
+class _DryRunDownloadManager:
+    """Minimal download-manager stub for split discovery without I/O.
+
+    Returns placeholder paths for all download/extract calls, allowing
+    _split_generators() to execute and return SplitGenerator objects
+    (which carry split names) without triggering actual network or disk I/O.
+    """
+
+    _PLACEHOLDER = os.devnull
+
+    def download(self, url_or_urls):
+        return self._map(url_or_urls)
+
+    def download_and_extract(self, url_or_urls):
+        return self._map(url_or_urls)
+
+    def extract(self, path_or_paths):
+        return self._map(path_or_paths)
+
+    def download_custom(self, url_or_urls, custom_download):
+        return self._map(url_or_urls)
+
+    def iter_archive(self, path):
+        return iter([])
+
+    def iter_files(self, paths):
+        return iter([])
+
+    @property
+    def manual_dir(self):
+        return None
+
+    @property
+    def is_streaming(self):
+        return False
+
+    def _map(self, input_):
+        if isinstance(input_, dict):
+            return {k: self._PLACEHOLDER for k in input_}
+        if isinstance(input_, (list, tuple, set)):
+            return type(input_)(self._PLACEHOLDER for _ in input_)
+        return self._PLACEHOLDER
+
+
+def _discover_splits_from_builder(builder_instance):
+    """Discover available split names by dry-running _split_generators().
+
+    For script-based datasets that lack README split metadata, this calls
+    the builder's _split_generators() with a stub download manager to
+    extract split names without performing any actual downloads.
+
+    Returns:
+        A set of split name strings, or an empty set if discovery fails.
+    """
+    if not hasattr(builder_instance, '_split_generators'):
+        return set()
+    try:
+        generators = builder_instance._split_generators(_DryRunDownloadManager())
+        return {str(sg.name) for sg in generators}
+    except Exception as e:
+        logger.debug(f'Failed to discover splits from builder: {e}')
+        return set()
+
+
+def _validate_split_exists(builder_instance, split):
+    """Fail-fast check: raise ValueError before downloading if the
+    requested split does not exist in the dataset metadata.
+
+    Args:
+        builder_instance: The DatasetBuilder instance with info/config.
+        split: The user-requested split specification (may be None).
+
+    Raises:
+        ValueError: If any requested split name is not found among
+            the available splits declared in the dataset metadata.
+    """
+    if split is None:
+        return
+
+    from modelscope.msdatasets.utils._module_factories import _extract_split_names
+    split_names = _extract_split_names(split)
+    if not split_names:
+        return
+
+    # Source 1: info.splits (original metadata from README)
+    available = set()
+    info = getattr(builder_instance, 'info', None)
+    if info is not None and info.splits:
+        available = set(info.splits.keys())
+
+    # Source 2: config.data_files keys
+    if not available:
+        config = getattr(builder_instance, 'config', None)
+        data_files = getattr(config, 'data_files', None)
+        if isinstance(data_files, dict):
+            available = set(data_files.keys())
+
+    # Source 3: dry-run _split_generators() for script-based datasets
+    if not available:
+        available = _discover_splits_from_builder(builder_instance)
+
+    if not available:
+        logger.debug(
+            'Cannot determine available splits from dataset metadata; '
+            'split validation skipped. Invalid splits will be caught downstream.'
+        )
+        return
+
+    missing = split_names - available
+    if missing:
+        raise ValueError(
+            f'Split {sorted(missing)} not found in dataset. '
+            f'Available splits: {sorted(available)}'
+        )
+
+
+def _align_builder_splits_with_data_files(builder_instance, split):
+    """Align builder.info.splits with the actually requested split(s).
+
+    When data_files have been filtered to a subset of splits (see
+    _filter_data_files_by_split in _module_factories.py), the builder's
+    info.splits metadata may still list all original splits from the
+    README.  download_and_prepare() calls verify_splits() which would
+    then raise ExpectedMoreSplitsError.  This helper prunes info.splits
+    to only contain the splits that will actually be generated.
+    """
+    if split is None:
+        return
+    info = getattr(builder_instance, 'info', None)
+    if info is None or info.splits is None:
+        return
+
+    from modelscope.msdatasets.utils._module_factories import _extract_split_names
+    split_names = _extract_split_names(split)
+    if not split_names:
+        return
+
+    existing_keys = set(info.splits.keys())
+    if split_names >= existing_keys:
+        return  # All splits requested, no filtering needed
+
+    filtered = {k: v for k, v in info.splits.items() if k in split_names}
+    if not filtered:
+        return  # Safety: don't empty out splits
+
+    info.splits = SplitDict(filtered, dataset_name=info.splits.dataset_name)
+
+
 # ===================================================================
 # DatasetsWrapperHF
 # ===================================================================
@@ -544,34 +709,66 @@ class DatasetsWrapperHF:
             storage_options=storage_options,
             trust_remote_code=trust_remote_code,
             _require_default_config_name=name is None,
+            split=split,
             **config_kwargs,
         )
 
         if dataset_info_only:
             ret_dict = {}
+
+            # Case 1: Local .py script file
             if isinstance(path, str) and path.endswith('.py') and os.path.exists(path):
                 from datasets import get_dataset_config_names
                 subset_list = get_dataset_config_names(path)
                 ret_dict = {_subset: [] for _subset in subset_list}
                 return ret_dict
 
-            if builder_instance is None or not hasattr(builder_instance,
-                                                       'builder_configs'):
-                logger.error(f'No builder_configs found for {path} dataset.')
+            if builder_instance is None:
+                logger.error(f'No builder instance created for {path} dataset.')
                 return ret_dict
 
-            _tmp_builder_configs = builder_instance.builder_configs
-            for tmp_config_name, tmp_builder_config in _tmp_builder_configs.items():
-                tmp_config_name = str(tmp_config_name)
-                if hasattr(tmp_builder_config, 'data_files') and tmp_builder_config.data_files is not None:
-                    ret_dict[tmp_config_name] = [str(item) for item in list(tmp_builder_config.data_files.keys())]
+            # Case 2: Try builder_configs with data_files (packaged datasets)
+            _tmp_builder_configs = getattr(builder_instance, 'builder_configs', None)
+            if _tmp_builder_configs:
+                if hasattr(_tmp_builder_configs, 'items'):
+                    configs_iter = _tmp_builder_configs.items()
                 else:
-                    ret_dict[tmp_config_name] = []
+                    configs_iter = [(getattr(c, 'name', 'default'), c) for c in _tmp_builder_configs]
+                for tmp_config_name, tmp_builder_config in configs_iter:
+                    tmp_config_name = str(tmp_config_name)
+                    if hasattr(tmp_builder_config, 'data_files') and tmp_builder_config.data_files is not None:
+                        ret_dict[tmp_config_name] = [str(item) for item in list(tmp_builder_config.data_files.keys())]
+
+            # Case 3: Fallback for script datasets — use info.splits or dry-run discovery
+            if not ret_dict or all(not v for v in ret_dict.values()):
+                config_name = getattr(builder_instance, 'config_name', 'default') or 'default'
+                splits = []
+
+                # Try info.splits (from README metadata)
+                info = getattr(builder_instance, 'info', None)
+                if info is not None and info.splits:
+                    splits = sorted(info.splits.keys())
+
+                # Fallback: dry-run _split_generators()
+                if not splits:
+                    discovered = _discover_splits_from_builder(builder_instance)
+                    if discovered:
+                        splits = sorted(discovered)
+
+                if splits:
+                    ret_dict = {config_name: splits}
+                elif not ret_dict:
+                    ret_dict = {config_name: []}
+
             return ret_dict
 
+        _validate_split_exists(builder_instance, split)
+
         if streaming:
             return builder_instance.as_streaming_dataset(split=split)
 
+        _align_builder_splits_with_data_files(builder_instance, split)
+
         builder_instance.download_and_prepare(
             download_config=download_config,
             download_mode=download_mode,
@@ -624,6 +821,7 @@ class DatasetsWrapperHF:
         storage_options: Optional[Dict] = None,
         trust_remote_code: Optional[bool] = None,
         _require_default_config_name=True,
+        split: Optional[Union[str, Split]] = None,
         **config_kwargs,
     ) -> DatasetBuilder:
 
@@ -644,6 +842,12 @@ class DatasetsWrapperHF:
             download_config = download_config.copy(
             ) if download_config else DownloadConfig()
             download_config.storage_options.update(storage_options)
+        if split is not None:
+            download_config = download_config.copy(
+            ) if download_config else DownloadConfig()
+            if download_config.storage_options is None:
+                download_config.storage_options = {}
+            download_config.storage_options['split'] = split
 
         dataset_module = DatasetsWrapperHF.dataset_module_factory(
             path,
@@ -683,6 +887,24 @@ class DatasetsWrapperHF:
         builder_cls = get_dataset_builder_class(
             dataset_module, dataset_name=dataset_name)
 
+        _config_cls = builder_cls.BUILDER_CONFIG_CLASS
+        if hasattr(_config_cls, '__dataclass_fields__'):
+            _valid_fields = set(_config_cls.__dataclass_fields__.keys())
+            # Also preserve parameters accepted by the builder's
+            # __init__ (e.g. writer_batch_size, base_path, repo_id)
+            # so they are not inadvertently stripped.
+            try:
+                _init_params = set(
+                    inspect.signature(builder_cls.__init__).parameters.keys()
+                )
+            except (ValueError, TypeError):
+                _init_params = set()
+            _valid_fields = _valid_fields | _init_params
+            config_kwargs = {
+                k: v for k, v in config_kwargs.items()
+                if k in _valid_fields
+            }
+
         builder_instance: DatasetBuilder = builder_cls(
             cache_dir=cache_dir,
             dataset_name=dataset_name,
@@ -946,7 +1168,7 @@ def load_dataset_with_ctx(*args, **kwargs):
     non-streaming mode) or kept alive (for streaming mode, where lazy
     iteration needs the patches to remain active).
     """
-    global _hf_fs_open_original
+    global _hf_fs_open_original, _hf_fs_init_original
 
     # Save originals
     hf_endpoint_origin = config.HF_ENDPOINT
@@ -962,6 +1184,7 @@ def load_dataset_with_ctx(*args, **kwargs):
         HubDatasetModuleFactoryWithScript.get_module if _HAS_SCRIPT_LOADING else None)
     generate_from_dict_origin = features.generate_from_dict
     hf_fs_open_origin = HfFileSystem._open
+    hf_fs_init_origin = HfFileSystem.__init__
 
     # Apply patches
     config.HF_ENDPOINT = get_endpoint()
@@ -980,6 +1203,8 @@ def load_dataset_with_ctx(*args, **kwargs):
     features.generate_from_dict = generate_from_dict_ms
     _hf_fs_open_original = hf_fs_open_origin
     HfFileSystem._open = _hf_fs_open
+    _hf_fs_init_original = hf_fs_init_origin
+    HfFileSystem.__init__ = _hf_fs_init_with_cookie
 
     streaming = kwargs.get('streaming', False)
 
@@ -990,10 +1215,12 @@ def load_dataset_with_ctx(*args, **kwargs):
         _repo_tree_cache.clear()
         HubApi._dataset_id_type_cache.clear()
 
-        HfFileSystem._open = hf_fs_open_origin
-        _hf_fs_open_original = None
-
         if not streaming:
+            HfFileSystem._open = hf_fs_open_origin
+            _hf_fs_open_original = None
+            HfFileSystem.__init__ = hf_fs_init_origin
+            _hf_fs_init_original = None
+
             config.HF_ENDPOINT = hf_endpoint_origin
             file_utils.get_from_cache = get_from_cache_origin
             features.generate_from_dict = generate_from_dict_origin