From 976bdf4ebc33e5439c7d6618751badfa866ec4ed Mon Sep 17 00:00:00 2001 From: "xingjun.wxj" Date: Tue, 18 Apr 2023 17:01:43 +0800 Subject: [PATCH] Add virgo hub 1. Add Hubs.virgo in MsDataset.load() 2. Add VirgoDownloader pipeline 3. Add auth for virgo hub. 4. Add other utils. 5. cherry-pick from https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/12352198?file=8c317278a40b19fe76825fb749e768fdbfd6a711 Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/12351633 --- modelscope/hub/api.py | 45 +++++- modelscope/msdatasets/auth/auth_config.py | 9 ++ .../context/dataset_context_config.py | 10 ++ .../msdatasets/data_loader/data_loader.py | 150 ++++++++++++++++-- .../data_loader/data_loader_manager.py | 4 +- modelscope/msdatasets/dataset_cls/__init__.py | 2 +- modelscope/msdatasets/dataset_cls/dataset.py | 81 +++++++++- .../msdatasets/meta/data_meta_manager.py | 8 + modelscope/msdatasets/ms_dataset.py | 30 +++- modelscope/utils/constant.py | 21 +++ modelscope/utils/service_utils.py | 5 + modelscope/utils/url_utils.py | 36 +++++ tests/msdatasets/test_virgo_dataset.py | 76 +++++++++ 13 files changed, 450 insertions(+), 27 deletions(-) create mode 100644 modelscope/utils/url_utils.py create mode 100644 tests/msdatasets/test_virgo_dataset.py diff --git a/modelscope/hub/api.py b/modelscope/hub/api.py index ad8d0c5d..0991dbb9 100644 --- a/modelscope/hub/api.py +++ b/modelscope/hub/api.py @@ -15,6 +15,7 @@ from http.cookiejar import CookieJar from os.path import expanduser from typing import Dict, List, Optional, Tuple, Union +import requests from requests import Session from requests.adapters import HTTPAdapter, Retry @@ -45,7 +46,7 @@ from modelscope.utils.constant import (DEFAULT_DATASET_REVISION, MASTER_MODEL_BRANCH, DatasetFormations, DatasetMetaFormats, DatasetVisibilityMap, DownloadChannel, - ModelFile) + ModelFile, VirgoDatasetConfig) from modelscope.utils.logger import get_logger from .utils.utils import (get_endpoint, get_release_datetime, model_id_to_group_owner_name) @@ -581,6 +582,17 @@ class HubApi: file_list = file_list['Files'] return file_list + @staticmethod + def dump_datatype_file(dataset_type: int, meta_cache_dir: str): + """ + Dump the data_type as a local file, in order to get the dataset formation without calling the datahub. + More details, please refer to the class `modelscope.utils.constant.DatasetFormations`. + """ + dataset_type_file_path = os.path.join(meta_cache_dir, + f'{str(dataset_type)}{DatasetFormations.formation_mark_ext.value}') + with open(dataset_type_file_path, 'w') as fp: + fp.write('*** Automatically-generated file, do not modify ***') + def get_dataset_meta_files_local_paths(self, dataset_name: str, namespace: str, revision: str, @@ -591,10 +603,7 @@ class HubApi: cookies = ModelScopeConfig.get_cookies() # Dump the data_type as a local file - dataset_type_file_path = os.path.join(meta_cache_dir, - f'{str(dataset_type)}{DatasetFormations.formation_mark_ext.value}') - with open(dataset_type_file_path, 'w') as fp: - fp.write('*** Automatically-generated file, do not modify ***') + HubApi.dump_datatype_file(dataset_type=dataset_type, meta_cache_dir=meta_cache_dir) for file_info in file_list: file_path = file_info['Path'] @@ -661,7 +670,6 @@ class HubApi: cookies = self._check_cookie(use_cookies=True) else: cookies = ModelScopeConfig.get_cookies() - r = self.session.get(url=datahub_url, cookies=cookies, headers=self.headers) r = self.session.get( url=datahub_url, cookies=cookies, headers=self.headers) @@ -669,6 +677,31 @@ class HubApi: raise_on_error(resp) return resp['Data'] + def get_virgo_meta(self, dataset_id: str, version: int = 1) -> dict: + """ + Get virgo dataset meta info. + """ + virgo_endpoint = os.environ.get(VirgoDatasetConfig.env_virgo_endpoint, '') + if not virgo_endpoint: + raise RuntimeError(f'Virgo endpoint is not set in env: {VirgoDatasetConfig.env_virgo_endpoint}') + + virgo_dataset_url = f'{virgo_endpoint}/data/set/download' + cookies = requests.utils.dict_from_cookiejar(ModelScopeConfig.get_cookies()) + + dataset_info = dict( + dataSetId=dataset_id, + dataSetVersion=version + ) + data = dict( + data=dataset_info, + ) + r = self.session.post(url=virgo_dataset_url, json=data, cookies=cookies, headers=self.headers, timeout=900) + resp = r.json() + if resp['code'] != 0: + raise RuntimeError(f'Failed to get virgo dataset: {resp}') + + return resp['data'] + def get_dataset_access_config_for_unzipped(self, dataset_name: str, namespace: str, diff --git a/modelscope/msdatasets/auth/auth_config.py b/modelscope/msdatasets/auth/auth_config.py index 576a6efd..e09db93c 100644 --- a/modelscope/msdatasets/auth/auth_config.py +++ b/modelscope/msdatasets/auth/auth_config.py @@ -23,6 +23,15 @@ class OssAuthConfig(BaseAuthConfig): cookies=cookies, git_token=git_token, user_info=user_info) +class VirgoAuthConfig(BaseAuthConfig): + """The authorization config for virgo dataset.""" + + def __init__(self, cookies: CookieJar, git_token: str, + user_info: Tuple[str, str]): + super().__init__( + cookies=cookies, git_token=git_token, user_info=user_info) + + class MaxComputeAuthConfig(BaseAuthConfig): # TODO: MaxCompute dataset to be supported. def __init__(self, cookies: CookieJar, git_token: str, diff --git a/modelscope/msdatasets/context/dataset_context_config.py b/modelscope/msdatasets/context/dataset_context_config.py index 26b05f7d..929e68d0 100644 --- a/modelscope/msdatasets/context/dataset_context_config.py +++ b/modelscope/msdatasets/context/dataset_context_config.py @@ -24,6 +24,7 @@ class DatasetContextConfig: self._config_kwargs = kwargs self._dataset_version_cache_root_dir = None self._auth_config = None + self._ext_config: dict = {} # The lock file path for meta-files and data-files self._global_meta_lock_file_path = None @@ -42,6 +43,7 @@ class DatasetContextConfig: self.data_files = data_files self.cache_root_dir = cache_root_dir self.use_streaming = use_streaming + self.download_virgo_files: bool = False @property def config_kwargs(self) -> dict: @@ -98,3 +100,11 @@ class DatasetContextConfig: @auth_config.setter def auth_config(self, val: BaseAuthConfig): self._auth_config = val + + @property + def ext_config(self) -> dict: + return self._ext_config + + @ext_config.setter + def ext_config(self, val: dict): + self._ext_config = val diff --git a/modelscope/msdatasets/data_loader/data_loader.py b/modelscope/msdatasets/data_loader/data_loader.py index 1ef92372..38b83c46 100644 --- a/modelscope/msdatasets/data_loader/data_loader.py +++ b/modelscope/msdatasets/data_loader/data_loader.py @@ -1,11 +1,12 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +import os from abc import ABC, abstractmethod from typing import Optional, Union from datasets import (Dataset, DatasetBuilder, DatasetDict, IterableDataset, IterableDatasetDict) -from datasets import load_dataset as hf_data_loader +from datasets import load_dataset as hf_load_dataset from modelscope.hub.api import ModelScopeConfig from modelscope.msdatasets.auth.auth_config import OssAuthConfig @@ -15,11 +16,16 @@ from modelscope.msdatasets.data_files.data_files_manager import \ DataFilesManager from modelscope.msdatasets.dataset_cls.dataset import ExternalDataset from modelscope.msdatasets.meta.data_meta_manager import DataMetaManager -from modelscope.utils.constant import DatasetFormations +from modelscope.utils.constant import (DatasetFormations, DatasetPathName, + DownloadMode, VirgoDatasetConfig) +from modelscope.utils.logger import get_logger +from modelscope.utils.url_utils import valid_url + +logger = get_logger() -class BaseDataLoader(ABC): - """Base dataset loader to load data.""" +class BaseDownloader(ABC): + """Base dataset downloader to load data.""" def __init__(self, dataset_context_config: DatasetContextConfig): self.dataset_context_config = dataset_context_config @@ -28,35 +34,35 @@ class BaseDataLoader(ABC): def process(self): """The entity processing pipeline for fetching the data. """ raise NotImplementedError( - f'No default implementation provided for {BaseDataLoader.__name__}.process.' + f'No default implementation provided for {BaseDownloader.__name__}.process.' ) @abstractmethod def _authorize(self): raise NotImplementedError( - f'No default implementation provided for {BaseDataLoader.__name__}._authorize.' + f'No default implementation provided for {BaseDownloader.__name__}._authorize.' ) @abstractmethod def _build(self): raise NotImplementedError( - f'No default implementation provided for {BaseDataLoader.__name__}._build.' + f'No default implementation provided for {BaseDownloader.__name__}._build.' ) @abstractmethod def _prepare_and_download(self): raise NotImplementedError( - f'No default implementation provided for {BaseDataLoader.__name__}._prepare_and_download.' + f'No default implementation provided for {BaseDownloader.__name__}._prepare_and_download.' ) @abstractmethod def _post_process(self): raise NotImplementedError( - f'No default implementation provided for {BaseDataLoader.__name__}._post_process.' + f'No default implementation provided for {BaseDownloader.__name__}._post_process.' ) -class OssDataLoader(BaseDataLoader): +class OssDownloader(BaseDownloader): def __init__(self, dataset_context_config: DatasetContextConfig): super().__init__(dataset_context_config) @@ -127,7 +133,7 @@ class OssDataLoader(BaseDataLoader): raise f'meta-file: {dataset_name}.py not found on the modelscope hub.' if dataset_py_script and dataset_formation == DatasetFormations.hf_compatible: - self.dataset = hf_data_loader( + self.dataset = hf_load_dataset( dataset_py_script, name=subset_name, revision=version, @@ -147,8 +153,126 @@ class OssDataLoader(BaseDataLoader): self.dataset.custom_map = self.dataset_context_config.data_meta_config.meta_type_map -class MaxComputeDataLoader(BaseDataLoader): - """Data loader for MaxCompute data source.""" +class VirgoDownloader(BaseDownloader): + """Data downloader for Virgo data source.""" + + def __init__(self, dataset_context_config: DatasetContextConfig): + super().__init__(dataset_context_config) + self.dataset = None + + def process(self): + """ + Sequential data fetching virgo dataset process: authorize -> build -> prepare_and_download -> post_process + """ + self._authorize() + self._build() + self._prepare_and_download() + self._post_process() + + def _authorize(self): + """Authorization of virgo dataset.""" + from modelscope.msdatasets.auth.auth_config import VirgoAuthConfig + + cookies = ModelScopeConfig.get_cookies() + user_info = ModelScopeConfig.get_user_info() + + if not self.dataset_context_config.auth_config: + auth_config = VirgoAuthConfig( + cookies=cookies, git_token='', user_info=user_info) + else: + auth_config = self.dataset_context_config.auth_config + auth_config.cookies = cookies + auth_config.git_token = '' + auth_config.user_info = user_info + + self.dataset_context_config.auth_config = auth_config + + def _build(self): + """ + Fetch virgo meta and build virgo dataset. + """ + from modelscope.msdatasets.dataset_cls import VirgoDataset + + meta_manager = DataMetaManager(self.dataset_context_config) + meta_manager.fetch_virgo_meta() + self.dataset_context_config = meta_manager.dataset_context_config + self.dataset = VirgoDataset(**self.dataset_context_config.ext_config) + + virgo_cache_dir = os.path.join( + self.dataset_context_config.cache_root_dir, + self.dataset_context_config.namespace, + self.dataset_context_config.dataset_name, + self.dataset_context_config.version) + os.makedirs( + os.path.join(virgo_cache_dir, DatasetPathName.META_NAME), + exist_ok=True) + meta_content_cache_file = os.path.join(virgo_cache_dir, + DatasetPathName.META_NAME, + 'meta_content.csv') + meta_content_df = self.dataset.meta + meta_content_df.to_csv(meta_content_cache_file, index=False) + self.dataset.meta_content_cache_file = meta_content_cache_file + self.dataset.virgo_cache_dir = virgo_cache_dir + logger.info(f'Virgo meta content saved to {meta_content_cache_file}') + + def _prepare_and_download(self): + """ + Fetch data-files from oss-urls in the virgo meta content. + """ + + if self.dataset_context_config.download_virgo_files: + import requests + import json + import shutil + from urllib.parse import urlparse + + def download_file(meta_info_val): + file_url = '' + file_path = '' + try: + file_url = json.loads(meta_info_val)['url'] + is_url = valid_url(file_url) + if is_url: + url_parse_res = urlparse(file_url) + file_name = os.path.basename(url_parse_res.path) + else: + raise ValueError(f'Unsupported url: {file_url}') + file_path = os.path.join(data_files_dir, file_name) + except Exception as e: + logger.warning(e) + + if file_path and not os.path.exists(file_path): + logger.info( + f'Downloading file from {file_url} to {file_path}') + with open(file_path, 'wb') as f: + f.write(requests.get(file_url).content) + + return file_path + + self.dataset.download_virgo_files = True + download_mode = self.dataset_context_config.download_mode + data_files_dir = os.path.join(self.dataset.virgo_cache_dir, + DatasetPathName.DATA_FILES_NAME) + if download_mode == DownloadMode.REUSE_DATASET_IF_EXISTS: + self.dataset.meta[VirgoDatasetConfig. + col_cache_file] = self.dataset.meta.apply( + lambda row: download_file(row.meta_info), + axis=1) + elif download_mode == DownloadMode.FORCE_REDOWNLOAD: + shutil.rmtree(data_files_dir, ignore_errors=True) + self.dataset.meta[VirgoDatasetConfig. + col_cache_file] = self.dataset.meta.apply( + lambda row: download_file(row.meta_info), + axis=1) + else: + raise ValueError(f'Unsupported download mode: {download_mode}') + + def _post_process(self): + ... + + +class MaxComputeDownloader(BaseDownloader): + """Data downloader for MaxCompute data source.""" # TODO: MaxCompute data source to be supported . def __init__(self, dataset_context_config: DatasetContextConfig): diff --git a/modelscope/msdatasets/data_loader/data_loader_manager.py b/modelscope/msdatasets/data_loader/data_loader_manager.py index 3c8a638a..5be32de1 100644 --- a/modelscope/msdatasets/data_loader/data_loader_manager.py +++ b/modelscope/msdatasets/data_loader/data_loader_manager.py @@ -9,7 +9,7 @@ from datasets import load_dataset as hf_data_loader from modelscope.hub.api import HubApi from modelscope.msdatasets.context.dataset_context_config import \ DatasetContextConfig -from modelscope.msdatasets.data_loader.data_loader import OssDataLoader +from modelscope.msdatasets.data_loader.data_loader import OssDownloader from modelscope.utils.constant import EXTENSIONS_TO_LOAD from modelscope.utils.logger import get_logger @@ -127,7 +127,7 @@ class RemoteDataLoaderManager(DataLoaderManager): return dataset_ret # To use the modelscope data loader elif data_loader_type == RemoteDataLoaderType.MS_DATA_LOADER: - oss_data_loader = OssDataLoader( + oss_data_loader = OssDownloader( dataset_context_config=self.dataset_context_config) oss_data_loader.process() # download statistics diff --git a/modelscope/msdatasets/dataset_cls/__init__.py b/modelscope/msdatasets/dataset_cls/__init__.py index a5b2e73d..ced94f2f 100644 --- a/modelscope/msdatasets/dataset_cls/__init__.py +++ b/modelscope/msdatasets/dataset_cls/__init__.py @@ -1,3 +1,3 @@ # Copyright (c) Alibaba, Inc. and its affiliates. -from .dataset import ExternalDataset, NativeIterableDataset +from .dataset import ExternalDataset, NativeIterableDataset, VirgoDataset diff --git a/modelscope/msdatasets/dataset_cls/dataset.py b/modelscope/msdatasets/dataset_cls/dataset.py index 4acf51b1..e5e9664a 100644 --- a/modelscope/msdatasets/dataset_cls/dataset.py +++ b/modelscope/msdatasets/dataset_cls/dataset.py @@ -4,11 +4,13 @@ import copy import os import datasets +import pandas as pd from datasets import IterableDataset from PIL import Image -from modelscope.utils.constant import EXTENSIONS_TO_LOAD +from modelscope.utils.constant import EXTENSIONS_TO_LOAD, VirgoDatasetConfig from modelscope.utils.logger import get_logger +from modelscope.utils.url_utils import fetch_csv_with_url, valid_url logger = get_logger() @@ -108,3 +110,80 @@ class NativeIterableDataset(IterableDataset): def __len__(self): return 1 + + +class VirgoDataset(object): + """Dataset class for Virgo. + + Attributes: + _meta_content (str): Virgo meta data content, could be a url that contains csv file. + _data_type (int): Virgo dataset type, 0-Standard virgo dataset; Others-User define dataset (to be supported) + + Examples: + >>> from modelscope.msdatasets.dataset_cls import VirgoDataset + >>> input_kwargs = {'metaContent': 'http://xxx-xxx/xxx.csv', 'samplingType': 0} + >>> virgo_dataset = VirgoDataset(**input_kwargs) + >>> print(virgo_dataset[1]) + >>> print(len(virgo_dataset)) + >>> for line in virgo_dataset: + >>> print(line) + + Note: If you set `download_virgo_files` to True by using + MsDataset.load(dataset_name='your-virgo-dataset-id', hub=Hubs.virgo, download_virgo_files=True), + you can get the cache file path of the virgo dataset, the column name is `cache_file`. + >>> if virgo_dataset.download_virgo_files: + >>> print(virgo_dataset[1].get('cache_file')) + """ + + def __init__(self, **kwargs): + + self._meta_content: str = '' + self._data_type: int = 0 + self._meta: pd.DataFrame = pd.DataFrame() + + if VirgoDatasetConfig.meta_content in kwargs: + self._meta_content = kwargs.pop(VirgoDatasetConfig.meta_content) + if VirgoDatasetConfig.sampling_type in kwargs: + self._data_type = kwargs.pop(VirgoDatasetConfig.sampling_type) + + self._check_variables() + self._parse_meta() + + self.meta_content_cache_file = '' + self.virgo_cache_dir = '' + self.download_virgo_files: bool = False + + def __getitem__(self, index): + return self._meta.iloc[index].to_dict() + + def __len__(self): + return len(self._meta) + + def __iter__(self): + for _, row in self._meta.iterrows(): + yield row.to_dict() + + @property + def meta(self) -> pd.DataFrame: + """ + Virgo meta data. Contains columns: id, meta_info, analysis_result, external_info and + cache_file (if download_virgo_files is True). + """ + return self._meta + + def _parse_meta(self): + # Fetch csv content + meta_content_df = fetch_csv_with_url(self._meta_content) + self._meta = meta_content_df + + def _check_variables(self): + """Check member variables in this class. + 1. Condition-1: self._meta_content cannot be empty + 2. Condition-2: self._meta_content must be url when self._data_type is 0 + """ + if not self._meta_content: + raise 'Them meta content cannot be empty.' + if self._data_type != 0: + raise 'Supported samplingType should be 0, others are not supported yet.' + if not valid_url(self._meta_content): + raise 'The meta content must be url when data type is 0.' diff --git a/modelscope/msdatasets/meta/data_meta_manager.py b/modelscope/msdatasets/meta/data_meta_manager.py index d90b8d5e..6a34d5d6 100644 --- a/modelscope/msdatasets/meta/data_meta_manager.py +++ b/modelscope/msdatasets/meta/data_meta_manager.py @@ -140,6 +140,14 @@ class DataMetaManager(object): self.dataset_context_config.data_meta_config = data_meta_config + def fetch_virgo_meta(self) -> None: + virgo_dataset_id = self.dataset_context_config.dataset_name + version = int(self.dataset_context_config.version) + + meta_content = self.api.get_virgo_meta( + dataset_id=virgo_dataset_id, version=version) + self.dataset_context_config.ext_config = meta_content + def _fetch_meta_from_cache(self, meta_cache_dir): local_paths = defaultdict(list) dataset_type = None diff --git a/modelscope/msdatasets/ms_dataset.py b/modelscope/msdatasets/ms_dataset.py index 0a88eb91..6729e4e6 100644 --- a/modelscope/msdatasets/ms_dataset.py +++ b/modelscope/msdatasets/ms_dataset.py @@ -13,13 +13,15 @@ from datasets.utils.file_utils import is_relative_path from modelscope.hub.repository import DatasetRepository from modelscope.msdatasets.context.dataset_context_config import \ DatasetContextConfig +from modelscope.msdatasets.data_loader.data_loader import VirgoDownloader from modelscope.msdatasets.data_loader.data_loader_manager import ( LocalDataLoaderManager, LocalDataLoaderType, RemoteDataLoaderManager, RemoteDataLoaderType) +from modelscope.msdatasets.dataset_cls import (ExternalDataset, + NativeIterableDataset, + VirgoDataset) from modelscope.msdatasets.dataset_cls.custom_datasets.builder import \ build_custom_dataset -from modelscope.msdatasets.dataset_cls.dataset import (ExternalDataset, - NativeIterableDataset) from modelscope.msdatasets.utils.delete_utils import DatasetDeleteManager from modelscope.msdatasets.utils.upload_utils import DatasetUploadManager from modelscope.preprocessors import build_preprocessor @@ -28,7 +30,7 @@ from modelscope.utils.config_ds import MS_DATASETS_CACHE from modelscope.utils.constant import (DEFAULT_DATASET_NAMESPACE, DEFAULT_DATASET_REVISION, ConfigFields, DownloadMode, Hubs, ModeKeys, Tasks, - UploadMode) + UploadMode, VirgoDatasetConfig) from modelscope.utils.import_utils import is_tf_available, is_torch_available from modelscope.utils.logger import get_logger @@ -170,7 +172,7 @@ class MsDataset: use_streaming: Optional[bool] = False, custom_cfg: Optional[Config] = Config(), **config_kwargs, - ) -> Union[dict, 'MsDataset', NativeIterableDataset]: + ) -> Union[dict, 'MsDataset', NativeIterableDataset, VirgoDataset]: """Load a MsDataset from the ModelScope Hub, Hugging Face Hub, urls, or a local dataset. Args: @@ -287,6 +289,26 @@ class MsDataset: custom_cfg=custom_cfg, **config_kwargs) dataset_inst.is_custom = True return dataset_inst + elif hub == Hubs.virgo: + # Rewrite the namespace, version and cache_dir for virgo dataset. + if namespace == DEFAULT_DATASET_NAMESPACE: + dataset_context_config.namespace = VirgoDatasetConfig.default_virgo_namespace + if version == DEFAULT_DATASET_REVISION: + dataset_context_config.version = VirgoDatasetConfig.default_dataset_version + if cache_dir == MS_DATASETS_CACHE: + from modelscope.utils.config_ds import CACHE_HOME + cache_dir = os.path.join(CACHE_HOME, 'virgo', 'hub', + 'datasets') + dataset_context_config.cache_root_dir = cache_dir + if 'download_virgo_files' in config_kwargs: + dataset_context_config.download_virgo_files = config_kwargs.pop( + 'download_virgo_files') + + virgo_downloader = VirgoDownloader(dataset_context_config) + virgo_downloader.process() + + return virgo_downloader.dataset + else: raise 'Please adjust input args to specify a loading mode, we support following scenes: ' \ 'loading from local disk, huggingface hub and modelscope hub.' diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index 2382825a..a19eef20 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -327,6 +327,7 @@ class Hubs(enum.Enum): """ modelscope = 'modelscope' huggingface = 'huggingface' + virgo = 'virgo' class DownloadMode(enum.Enum): @@ -539,3 +540,23 @@ class DistributedParallelType(object): class DatasetTensorflowConfig: BATCH_SIZE = 'batch_size' DEFAULT_BATCH_SIZE_VALUE = 5 + + +class VirgoDatasetConfig: + + default_virgo_namespace = 'default_namespace' + + default_dataset_version = '1' + + env_virgo_endpoint = 'VIRGO_ENDPOINT' + + # Columns for meta request + meta_content = 'metaContent' + sampling_type = 'samplingType' + + # Columns for meta content + col_id = 'id' + col_meta_info = 'meta_info' + col_analysis_result = 'analysis_result' + col_external_info = 'external_info' + col_cache_file = 'cache_file' diff --git a/modelscope/utils/service_utils.py b/modelscope/utils/service_utils.py index 6e7c0fc1..e20988f7 100644 --- a/modelscope/utils/service_utils.py +++ b/modelscope/utils/service_utils.py @@ -8,6 +8,7 @@ import requests from modelscope.outputs import TASK_OUTPUTS, OutputKeys from modelscope.pipeline_inputs import TASK_INPUTS, InputType +from modelscope.utils.url_utils import valid_url # service data decoder func decodes data from network and convert it to pipeline's input @@ -82,12 +83,16 @@ def get_mimetype(filename): def decode_base64_to_binary(encoding): + if valid_url(encoding): + return encoding, '' extension = get_extension(encoding) data = encoding.split(',')[1] return base64.b64decode(data), extension def decode_base64_to_image(encoding): + if valid_url(encoding): + return encoding from PIL import Image content = encoding.split(';')[1] image_encoded = content.split(',')[1] diff --git a/modelscope/utils/url_utils.py b/modelscope/utils/url_utils.py new file mode 100644 index 00000000..59cc2efd --- /dev/null +++ b/modelscope/utils/url_utils.py @@ -0,0 +1,36 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +from urllib.parse import urlparse + +import pandas as pd + +from modelscope.utils.logger import get_logger + +logger = get_logger() + + +def valid_url(url) -> bool: + try: + result = urlparse(url) + return all([result.scheme, result.netloc]) + except ValueError as e: + logger.warning(e) + return False + + +def fetch_csv_with_url(csv_url: str) -> pd.DataFrame: + """Fetch the csv content from url. + + Args: + csv_url (str): The input url of csv data. + + Returns: + A pandas DataFrame object which contains the csv content. + """ + try: + df = pd.read_csv(csv_url) + except Exception as e: + logger.error(f'Failed to fetch csv from url: {csv_url}') + raise e + + return df diff --git a/tests/msdatasets/test_virgo_dataset.py b/tests/msdatasets/test_virgo_dataset.py new file mode 100644 index 00000000..03e79593 --- /dev/null +++ b/tests/msdatasets/test_virgo_dataset.py @@ -0,0 +1,76 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import os +import unittest + +from modelscope.hub.api import HubApi +from modelscope.msdatasets import MsDataset +from modelscope.msdatasets.dataset_cls import VirgoDataset +from modelscope.utils.constant import DownloadMode, Hubs, VirgoDatasetConfig +from modelscope.utils.logger import get_logger + +logger = get_logger() + +# Please use your own access token for buc account. +YOUR_ACCESS_TOKEN = 'your_access_token' +# Please use your own virgo dataset id and ensure you have access to it. +VIRGO_DATASET_ID = 'your_virgo_dataset_id' + + +class TestVirgoDataset(unittest.TestCase): + + def setUp(self): + self.api = HubApi() + self.api.login(YOUR_ACCESS_TOKEN) + + @unittest.skip('to be used for local test only') + def test_download_virgo_dataset_meta(self): + ds = MsDataset.load(dataset_name=VIRGO_DATASET_ID, hub=Hubs.virgo) + ds_one = next(iter(ds)) + logger.info(ds_one) + + self.assertTrue(ds_one) + self.assertIsInstance(ds_one, VirgoDataset) + self.assertIn(VirgoDatasetConfig.col_id, ds_one) + self.assertIn(VirgoDatasetConfig.col_meta_info, ds_one) + self.assertIn(VirgoDatasetConfig.col_analysis_result, ds_one) + self.assertIn(VirgoDatasetConfig.col_external_info, ds_one) + + @unittest.skip('to be used for local test only') + def test_download_virgo_dataset_files(self): + ds = MsDataset.load( + dataset_name=VIRGO_DATASET_ID, + hub=Hubs.virgo, + download_virgo_files=True) + + ds_one = next(iter(ds)) + logger.info(ds_one) + + self.assertTrue(ds_one) + self.assertIsInstance(ds_one, VirgoDataset) + self.assertTrue(ds_one.download_virgo_files) + self.assertIn(VirgoDatasetConfig.col_cache_file, ds_one) + cache_file_path = ds_one[VirgoDatasetConfig.col_cache_file] + self.assertTrue(os.path.exists(cache_file_path)) + + @unittest.skip('to be used for local test only') + def test_force_download_virgo_dataset_files(self): + ds = MsDataset.load( + dataset_name=VIRGO_DATASET_ID, + hub=Hubs.virgo, + download_mode=DownloadMode.FORCE_REDOWNLOAD, + download_virgo_files=True) + + ds_one = next(iter(ds)) + logger.info(ds_one) + + self.assertTrue(ds_one) + self.assertIsInstance(ds_one, VirgoDataset) + self.assertTrue(ds_one.download_virgo_files) + self.assertIn(VirgoDatasetConfig.col_cache_file, ds_one) + cache_file_path = ds_one[VirgoDatasetConfig.col_cache_file] + self.assertTrue(os.path.exists(cache_file_path)) + + +if __name__ == '__main__': + unittest.main()