Support jsonl format in meta data

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/13071970

* support jsonl in meta


* add UT and refine fetch_meta_files_from_url

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

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:

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,

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', ',')

modelscope/utils/constant.py

@@ -517,6 +517,8 @@ EXTENSIONS_TO_LOAD = {
     'txt': 'text'
 }
 
+META_FILES_FORMAT = ('.csv', '.jsonl')
+
 
 class DatasetPathName:
     META_NAME = 'meta'

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()