# Copyright (c) Alibaba, Inc. and its affiliates. import inspect import os from dataclasses import dataclass from typing import List, Optional, Union from evalscope.constants import HubType from evalscope.utils.logger import get_logger logger = get_logger() @dataclass(frozen=True) class DatasetHub: """Small hub handle shared by dataset loaders and benchmark-specific media resolvers.""" data_id_or_path: str data_source: Optional[str] = HubType.MODELSCOPE revision: Optional[str] = None trust_remote: bool = True force_redownload: bool = False cache_dir: Optional[str] = None def load(self, split: str, subset: str = 'default', **kwargs): return load_dataset_from_hub( data_id_or_path=self.data_id_or_path, split=split, subset=subset, data_source=self.data_source, version=self.revision, trust_remote=self.trust_remote, force_redownload=self.force_redownload, **kwargs, ) def download_file(self, file_path: str) -> str: return download_dataset_file( data_id_or_path=self.data_id_or_path, file_path=file_path, data_source=self.data_source, revision=self.revision, force_redownload=self.force_redownload, cache_dir=self.cache_dir, ) def download_snapshot( self, allow_file_pattern: Optional[Union[str, List[str]]] = None, ignore_file_pattern: Optional[Union[str, List[str]]] = None, ) -> str: return download_dataset_snapshot( data_id_or_path=self.data_id_or_path, data_source=self.data_source, revision=self.revision, force_redownload=self.force_redownload, cache_dir=self.cache_dir, allow_file_pattern=allow_file_pattern, ignore_file_pattern=ignore_file_pattern, ) def _resolve_data_source(data_id_or_path: str, data_source: Optional[str]) -> str: if data_source == HubType.LOCAL or os.path.exists(data_id_or_path): return HubType.LOCAL return data_source or HubType.MODELSCOPE def load_dataset_from_hub( data_id_or_path: str, split: str, subset: str = 'default', data_source: Optional[str] = HubType.MODELSCOPE, version: Optional[str] = None, trust_remote: bool = True, force_redownload: bool = False, **kwargs, ): """Load a dataset split from ModelScope, Hugging Face, or a local dataset path.""" import datasets from datasets import DownloadMode as HFDownloadMode from modelscope import MsDataset from modelscope.utils.constant import DownloadMode as MSDownloadMode data_source = _resolve_data_source(data_id_or_path, data_source) hf_download_mode = None if not force_redownload else HFDownloadMode.FORCE_REDOWNLOAD ms_download_mode = None if not force_redownload else MSDownloadMode.FORCE_REDOWNLOAD if data_source == HubType.MODELSCOPE: load_kwargs = dict( dataset_name=data_id_or_path, split=split, subset_name=subset if subset != 'default' else None, trust_remote_code=trust_remote, **kwargs, ) if version: load_kwargs['version'] = version if ms_download_mode: load_kwargs['download_mode'] = ms_download_mode dataset = MsDataset.load(**load_kwargs) if not isinstance(dataset, datasets.Dataset): dataset = dataset.to_hf_dataset() return dataset if data_source in [HubType.HUGGINGFACE, HubType.LOCAL]: # Hugging Face datasets may fail on local mirrors that contain a stale dataset_infos.json. dataset_infos_path = os.path.join(data_id_or_path, 'dataset_infos.json') if os.path.exists(dataset_infos_path): logger.info(f'Removing dataset_infos.json file at {dataset_infos_path} to avoid datasets errors.') os.remove(dataset_infos_path) load_kwargs = dict( path=data_id_or_path, name=subset if subset != 'default' else None, split=split, revision=version, download_mode=hf_download_mode, **kwargs, ) if 'trust_remote_code' in inspect.signature(datasets.load_dataset).parameters: load_kwargs['trust_remote_code'] = trust_remote return datasets.load_dataset(**load_kwargs) raise ValueError(f'Unsupported dataset hub: {data_source}') def download_dataset_file( data_id_or_path: str, file_path: str, data_source: Optional[str] = HubType.MODELSCOPE, revision: Optional[str] = None, force_redownload: bool = False, cache_dir: Optional[str] = None, ) -> str: """Download or resolve a single file from a dataset hub.""" data_source = _resolve_data_source(data_id_or_path, data_source) if data_source == HubType.LOCAL: root_dir = os.path.realpath(data_id_or_path) resolved_path = os.path.realpath(os.path.join(root_dir, file_path)) if os.path.commonpath([root_dir, resolved_path]) != root_dir: raise ValueError(f'Invalid dataset file path: {file_path}') if not os.path.exists(resolved_path): raise FileNotFoundError(f'Dataset file {file_path} was not found in {root_dir}.') return resolved_path if data_source == HubType.HUGGINGFACE: from huggingface_hub import hf_hub_download return hf_hub_download( repo_id=data_id_or_path, filename=file_path, repo_type='dataset', revision=revision, cache_dir=cache_dir, force_download=force_redownload, ) if data_source == HubType.MODELSCOPE: from modelscope import dataset_snapshot_download download_kwargs = {'allow_file_pattern': file_path} if revision: download_kwargs['revision'] = revision if cache_dir: download_kwargs['cache_dir'] = cache_dir snapshot_dir = dataset_snapshot_download(data_id_or_path, **download_kwargs) resolved_path = os.path.join(snapshot_dir, file_path) if not os.path.exists(resolved_path): raise FileNotFoundError(f'Dataset file {file_path} was not found in {snapshot_dir}.') return resolved_path raise ValueError(f'Unsupported dataset hub: {data_source}') def download_dataset_snapshot( data_id_or_path: str, data_source: Optional[str] = HubType.MODELSCOPE, revision: Optional[str] = None, force_redownload: bool = False, cache_dir: Optional[str] = None, allow_file_pattern: Optional[Union[str, List[str]]] = None, ignore_file_pattern: Optional[Union[str, List[str]]] = None, ) -> str: """Download or resolve a dataset snapshot root from a supported hub.""" data_source = _resolve_data_source(data_id_or_path, data_source) if data_source == HubType.LOCAL: root_dir = os.path.realpath(data_id_or_path) if not os.path.isdir(root_dir): raise FileNotFoundError(f'Local dataset directory was not found: {data_id_or_path}') return root_dir if data_source == HubType.HUGGINGFACE: from huggingface_hub import snapshot_download return snapshot_download( repo_id=data_id_or_path, repo_type='dataset', revision=revision, cache_dir=cache_dir, force_download=force_redownload, allow_patterns=allow_file_pattern, ignore_patterns=ignore_file_pattern, ) if data_source == HubType.MODELSCOPE: from modelscope import dataset_snapshot_download download_kwargs = {} if revision: download_kwargs['revision'] = revision if cache_dir: download_kwargs['cache_dir'] = cache_dir if allow_file_pattern is not None: download_kwargs['allow_file_pattern'] = allow_file_pattern if ignore_file_pattern is not None: download_kwargs['ignore_file_pattern'] = ignore_file_pattern return dataset_snapshot_download(data_id_or_path, **download_kwargs) raise ValueError(f'Unsupported dataset hub: {data_source}')