add ideadata2.0

Author: Zimix0

ideadata2.0/__init__.py

@@ -0,0 +1,2 @@
+from .load import load_dataset
+__all__ = ['load_dataset']

ideadata2.0/load.py

@@ -0,0 +1,131 @@
+import datasets
+from pathlib import Path
+from concurrent.futures import ProcessPoolExecutor
+
+
+_SPLIT_DATA_PATH = '/cognitive_comp/common_data/big_corpus/ideaData/ideaData2.0'
+# 缓存文件
+_CACHE_SPLIT_DATA_PATH = '/cognitive_comp/common_data/ideadata2.0/'
+# feats = datasets.Features({"text": datasets.Value('string')})
+
+
+class CommonTextCorpusGenerate(object):
+    """
+    处理通用文本数据，输入格式为jsonline，例如每一行都是：{"text":"这是段文本的内容"}
+    处理成memory map的数据类型，便于后续的加载
+    """
+    def __init__(self,
+                 data_files=_SPLIT_DATA_PATH,
+                 save_path=_CACHE_SPLIT_DATA_PATH,
+                 train_test_validation='9900,70,30',
+                 num_proc=1,
+                 if_shuffle=False,
+                 cache=False):
+        self.data_files = Path(data_files)
+        if save_path:
+            self.save_path = Path(save_path)
+        else:
+            self.save_path = self.file_check(
+                Path(self.data_files.parent, self.data_files.name+'_FSDataset'),'save')
+        self.num_proc = num_proc
+        self.cache = cache
+        self.split_idx = self.split_train_test_validation_index(train_test_validation)
+        self.shuffle = if_shuffle
+        if cache:
+            self.cache_path = self.file_check(
+                Path(self.save_path.parent, 'FSDataCache', self.data_files.name), 'cache')
+        else:
+            self.cache_path = None
+
+    @staticmethod
+    def file_check(path, path_type):
+        print(path)
+        if not path.exists():
+            path.mkdir(parents=True)
+        print(f"Since no {path_type} directory is specified, the program will automatically create it in {path} directory.")
+        return str(path)
+
+    @staticmethod
+    def split_train_test_validation_index(train_test_validation):
+        split_idx_ = [int(i) for i in train_test_validation.split(',')]
+        idx_dict = {
+            'train_rate': split_idx_[0]/sum(split_idx_),
+            'test_rate': split_idx_[1]/sum(split_idx_[1:])
+        }
+        return idx_dict
+
+    def process(self, index, path):
+        print('saving dataset shard {}'.format(index))
+
+        ds = (datasets.load_dataset('json', data_files=str(path),
+                                    cache_dir=self.cache_path,
+                                    features=None))
+        # 局部shuffle
+        # TODO 全局shuffle
+        if self.shuffle:
+            ds = ds.shuffle()
+        # 添加分句操作 没加
+        # 这里是拆分了数据集，
+        ds = ds['train'].train_test_split(train_size=self.split_idx['train_rate'])
+        ds_ = ds['test'].train_test_split(train_size=self.split_idx['test_rate'])
+        ds = datasets.DatasetDict({
+            'train': ds['train'],
+            'test': ds_['train'],
+            'validation': ds_['test']
+        })
+        
+        ds.save_to_disk(Path(self.save_path, path.name))
+        return 'saving dataset shard {} done'.format(index)
+
+    def generate_cache_arrow(self) -> None:
+        '''
+        生成HF支持的缓存文件，加速后续的加载
+        '''
+        data_dict_paths = self.data_files.rglob('*')
+        p = ProcessPoolExecutor(max_workers=self.num_proc)
+        res = list()
+
+        for index, path in enumerate(data_dict_paths):
+            # # for test 
+            # if index >10:
+            #     break
+            res.append(p.submit(self.process, index, path))
+
+        p.shutdown(wait=True)
+        for future in res:
+            print(future.result(), flush=True)
+
+
+def load_dataset(num_proc=4, **kargs):
+    cache_dict_paths = Path(_CACHE_SPLIT_DATA_PATH).glob('*')
+    ds = []
+    res = []
+    p = ProcessPoolExecutor(max_workers=num_proc)
+    for path in cache_dict_paths:
+        res.append(p.submit(datasets.load_from_disk,
+                            str(path), **kargs))
+
+    p.shutdown(wait=True)
+    for future in res:
+        ds.append(future.result())
+        # print(future.result())
+    train = []
+    test = []
+    validation = []
+    for ds_ in ds:
+        train.append(ds_['train'])
+        test.append(ds_['test'])
+        validation.append(ds_['validation'])
+    # ds = datasets.concatenate_datasets(ds)
+    # print(ds)
+    return datasets.DatasetDict({
+        'train': datasets.concatenate_datasets(train),
+        'test': datasets.concatenate_datasets(test),
+        'validation': datasets.concatenate_datasets(validation)
+    })
+
+
+
+if __name__ == '__main__':
+    dataset = CommonTextCorpusGenerate(_SPLIT_DATA_PATH, num_proc=16)
+    dataset.generate_cache_arrow()