diff --git a/modelscope/hub/api.py b/modelscope/hub/api.py index 2f14976d..a17a0538 100644 --- a/modelscope/hub/api.py +++ b/modelscope/hub/api.py @@ -44,8 +44,8 @@ from modelscope.hub.repository import Repository from modelscope.utils.constant import (DEFAULT_DATASET_REVISION, DEFAULT_MODEL_REVISION, DEFAULT_REPOSITORY_REVISION, - MASTER_MODEL_BRANCH, DatasetFormations, - DatasetMetaFormats, + MASTER_MODEL_BRANCH, META_FILES_FORMAT, + DatasetFormations, DatasetMetaFormats, DatasetVisibilityMap, DownloadChannel, DownloadMode, ModelFile, VirgoDatasetConfig) @@ -643,42 +643,57 @@ class HubApi: return local_paths, dataset_formation @staticmethod - def fetch_csv_from_url(url, out_path, chunk_size=100000, mode=DownloadMode.REUSE_DATASET_IF_EXISTS): - from io import StringIO + def fetch_meta_files_from_url(url, out_path, chunk_size=1024, mode=DownloadMode.REUSE_DATASET_IF_EXISTS): + """ + Fetch the meta-data files from the url, e.g. csv/jsonl files. + """ import hashlib + import json + from tqdm import tqdm out_path = os.path.join(out_path, hashlib.md5(url.encode(encoding='UTF-8')).hexdigest()) if mode == DownloadMode.FORCE_REDOWNLOAD and os.path.exists(out_path): os.remove(out_path) if os.path.exists(out_path): - logger.info(f'Reusing cached meta-csv file: {out_path}') + logger.info(f'Reusing cached meta-data file: {out_path}') return out_path cookies = ModelScopeConfig.get_cookies() # Make the request and get the response content as TextIO - logger.info('Loading meta-csv file ...') + logger.info('Loading meta-data file ...') + response = requests.get(url, cookies=cookies, stream=True) + total_size = int(response.headers.get('content-length', 0)) + progress = tqdm(total=total_size, dynamic_ncols=True) - response = requests.get(url, cookies=cookies) - data = StringIO(response.text) + def get_chunk(resp): + chunk_data = [] + for data in resp.iter_lines(): + data = data.decode('utf-8') + chunk_data.append(data) + if len(chunk_data) >= chunk_size: + yield chunk_data + chunk_data = [] + yield chunk_data - # Use read_csv with the TextIO object - csv_file_reader = pd.read_csv(data, iterator=True, dtype=str, delimiter=None) - - loop = True iter_num = 0 - while loop: - try: - chunk = csv_file_reader.get_chunk(size=chunk_size) - logger.info(f'Receiving chunk {iter_num}, shape: {chunk.shape}') - if iter_num == 0: - with_header = True + with open(out_path, 'a') as f: + for chunk in get_chunk(response): + progress.update(len(chunk)) + if url.endswith('jsonl'): + chunk = [json.loads(line) for line in chunk if line.strip()] + if len(chunk) == 0: + continue + if iter_num == 0: + with_header = True + else: + with_header = False + chunk_df = pd.DataFrame(chunk) + chunk_df.to_csv(f, index=False, header=with_header) + iter_num += 1 else: - with_header = False - - chunk.to_csv(out_path, mode='a', index=False, header=with_header) - iter_num += 1 - except StopIteration: - loop = False - logger.info('stop chunk iteration') + # csv or others + for line in chunk: + f.write(line + '\n') + progress.close() return out_path @@ -688,7 +703,7 @@ class HubApi: dataset_name: str, namespace: str, revision: Optional[str] = DEFAULT_DATASET_REVISION): - if file_name.endswith('.csv'): + if file_name and os.path.splitext(file_name)[-1] in META_FILES_FORMAT: file_name = f'{self.endpoint}/api/v1/datasets/{namespace}/{dataset_name}/repo?' \ f'Revision={revision}&FilePath={file_name}' return file_name diff --git a/modelscope/msdatasets/data_files/data_files_manager.py b/modelscope/msdatasets/data_files/data_files_manager.py index d5f43533..206a4d74 100644 --- a/modelscope/msdatasets/data_files/data_files_manager.py +++ b/modelscope/msdatasets/data_files/data_files_manager.py @@ -13,8 +13,8 @@ from modelscope.msdatasets.download.dataset_builder import ( from modelscope.msdatasets.download.download_config import DataDownloadConfig from modelscope.msdatasets.download.download_manager import ( DataDownloadManager, DataStreamingDownloadManager) -from modelscope.utils.constant import (DatasetPathName, DownloadMode, - MetaDataFields) +from modelscope.utils.constant import (META_FILES_FORMAT, DatasetPathName, + DownloadMode, MetaDataFields) class DataFilesManager(object): @@ -85,7 +85,8 @@ class DataFilesManager(object): builder = TaskSpecificDatasetBuilder( dataset_context_config=self.dataset_context_config) - elif meta_data_file.endswith('.csv'): + elif meta_data_file and os.path.splitext( + meta_data_file)[-1] in META_FILES_FORMAT: builder = CsvDatasetBuilder( dataset_context_config=self.dataset_context_config) else: diff --git a/modelscope/msdatasets/download/dataset_builder.py b/modelscope/msdatasets/download/dataset_builder.py index 796e3d83..6a8251f1 100644 --- a/modelscope/msdatasets/download/dataset_builder.py +++ b/modelscope/msdatasets/download/dataset_builder.py @@ -206,7 +206,7 @@ class CsvDatasetBuilder(csv.Csv): os.makedirs(target_cache_dir, exist_ok=True) self.local_meta_csv_paths = { - k: HubApi.fetch_csv_from_url(v, target_cache_dir) + k: HubApi.fetch_meta_files_from_url(v, target_cache_dir) for k, v in self.meta_data_files.items() } @@ -478,7 +478,7 @@ class IterableDatasetBuilder(csv.Csv): def _get_meta_csv_df(self, meta_file_url: str) -> None: if not self.meta_csv_df: - meta_csv_file_path = HubApi.fetch_csv_from_url( + meta_csv_file_path = HubApi.fetch_meta_files_from_url( meta_file_url, self.meta_cache_dir) self.meta_csv_df = pd.read_csv( meta_csv_file_path, diff --git a/modelscope/msdatasets/utils/dataset_utils.py b/modelscope/msdatasets/utils/dataset_utils.py index dc7f9b38..b40915eb 100644 --- a/modelscope/msdatasets/utils/dataset_utils.py +++ b/modelscope/msdatasets/utils/dataset_utils.py @@ -207,7 +207,7 @@ def get_dataset_files(subset_split_into: dict, if args_dict and args_dict.get(MetaDataFields.ARGS_BIG_DATA): meta_csv_file_url = meta_map[split] - meta_csv_file_path = HubApi.fetch_csv_from_url( + meta_csv_file_path = HubApi.fetch_meta_files_from_url( meta_csv_file_url, meta_cache_dir) csv_delimiter = context_config.config_kwargs.get('delimiter', ',') diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index 1f44fc01..32b1e2da 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -517,6 +517,8 @@ EXTENSIONS_TO_LOAD = { 'txt': 'text' } +META_FILES_FORMAT = ('.csv', '.jsonl') + class DatasetPathName: META_NAME = 'meta' diff --git a/tests/msdatasets/test_load_meta_jsonl.py b/tests/msdatasets/test_load_meta_jsonl.py new file mode 100644 index 00000000..64198566 --- /dev/null +++ b/tests/msdatasets/test_load_meta_jsonl.py @@ -0,0 +1,30 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import unittest + +from modelscope.msdatasets import MsDataset +from modelscope.utils import logger as logging +from modelscope.utils.constant import DownloadMode +from modelscope.utils.test_utils import test_level + +logger = logging.get_logger() + + +class TestLoadMetaJsonl(unittest.TestCase): + + def setUp(self): + self.dataset_id = 'modelscope/ms_ds_meta_jsonlines' + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_load_jsonl_in_meta(self): + ds = MsDataset.load( + self.dataset_id, + split='test', + download_mode=DownloadMode.FORCE_REDOWNLOAD) + ds_one = next(iter(ds)) + logger.info(next(iter(ds))) + assert ds_one['text'] + + +if __name__ == '__main__': + unittest.main()