refactor(data): 重构数据处理流程以适应新需求

2025-07-22 10:03:06 +08:00
parent b588c30f09
commit 5ea1c6355d
3 changed files with 7 additions and 4 deletions
@@ -79,12 +79,15 @@ def dump_pkl_files(
        os.makedirs(os.path.dirname(out_file_path), exist_ok=True)
        
        arrows: Dict[str, List[Tensor]] = {}
+        for key in output_keys:
+            arrows[key] = []
        
        with open(file_path, "r") as f:    
            lines = f.readlines()
            
        for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
-            arrow = process_func(line)
+            line_dict = json.loads(line)
+            arrow = process_func(line_dict)
            for key in output_keys:
                arrows[key].extend(arrow[key])