reafactor: 使用h5 文件格式存储

modules/utils.py

@@ -1,11 +1,12 @@
-from typing import Dict, List, Callable, Union
+from typing import Dict, List, Callable, Tuple, Union
 from datasets import DatasetDict
-from .tokenizer import BpeTokenizer
+import numpy as np
+from modules.tokenizer import BpeTokenizer
 from tqdm import tqdm
 from torch import Tensor
 
-import pickle as pkl
 import torch
+import h5py
 import json
 import os
 import re
@@ -25,6 +26,42 @@ def fetch_folders(root_dir, filter_func=None):
                 folders.append(folder_path)
     return folders
 
+def save_h5(file_path: str, tensor_group: Dict[str, List[Tensor]]):
+    os.makedirs(os.path.dirname(file_path), exist_ok=True)
+    with h5py.File(file_path, 'w') as f:
+        for key, tensors in tensor_group.items():
+            grp = f.create_group(key)
+            grp.attrs['num_tensors'] = len(tensors)
+            
+            for idx, tensor in enumerate(tensors):
+                arr = tensor.cpu().numpy()
+                dset = grp.create_dataset(
+                    f'data_{idx}',
+                    data=arr,
+                    compression='gzip',
+                    compression_opts=4,
+                    shuffle=True
+                )
+                dset.attrs['numel'] = tensor.numel()
+
+def load_h5(file_path: str) -> Tuple[Dict[str, List[Tensor]], int]:
+    tensor_group: Dict[str, List[Tensor]] = {}
+    total_samples = 0
+
+    with h5py.File(file_path, 'r') as f:
+        for key in f.keys():
+            grp = f[key]
+            dsets = []
+            for dset_name in grp.keys():
+                dset = grp[dset_name]
+                dsets.append(torch.from_numpy(dset[:]).share_memory_())
+                total_samples += dset.attrs.get('numel', np.prod(dset.shape))
+            tensor_group[key] = dsets
+
+    num_keys = max(len(tensor_group), 1)
+    sample_per_key = total_samples // num_keys
+
+    return tensor_group, sample_per_key
 
 def comprehensive_normalization(text):
     replacements = {
@@ -62,7 +99,7 @@ def pack_sequences(sequences: List[Tensor], pack_size: int, pad_value: int) -> L
     return packages
 
 
-def dump_pkl_files(
+def dump_files(
     files: List[str], 
     base_out_dir: str,
     process_func: Callable[[dict], dict],
@@ -101,8 +138,7 @@ def dump_pkl_files(
             sequence = torch.cat(package[key])
             output_package[key] = sequence
          
-        with open(out_file_path, "wb") as f:
-            pkl.dump(output_package, f)
+        save_h5(out_file_path, output_package)
 
 
 def get_pt_processor(tokenizer: BpeTokenizer):
@@ -151,7 +187,7 @@ def cache_files(tokenizer, files, base_out_dir, cache_type, packing_size: int =
     else:
         raise ValueError("Invalid cache type")
     
-    dump_pkl_files(files, base_out_dir, processor, keys, packing_size, pad_value)
+    dump_files(files, base_out_dir, processor, keys, packing_size, pad_value)
                           
             
 def process_dataset(