From 64222aeab19b980ac0da1dd9973a422a83b6ea63 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Tue, 22 Jul 2025 13:30:53 +0800
Subject: [PATCH] =?UTF-8?q?refactor(dump=5Fsft=5Ffile.py=EF=BC=8C=20utils.?=
 =?UTF-8?q?py):=20=20=E4=BC=98=E5=8C=96=E5=86=85=E5=AD=98=E5=8D=A0?=
 =?UTF-8?q?=E7=94=A8?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 dump_sft_file.py |  4 ++--
 utils.py         | 16 +++++++++-------
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/dump_sft_file.py b/dump_sft_file.py
index 6a2ca7a..1c7a9ca 100644
--- a/dump_sft_file.py
+++ b/dump_sft_file.py
@@ -13,8 +13,8 @@ def get_processor(tokenizer: BpeTokenizer):
         
         tokens = prefix_ids + suffix_ids
         tokens = torch.tensor(tokens, dtype=torch.int32)
-        masks = torch.zeros_like(tokens)
-        masks[:len(prefix_ids)] = 1
+        masks = torch.zeros_like(tokens, dtype=torch.bool)
+        masks[:len(prefix_ids)] = True
         
         return {"sequence": tokens, "mask": masks}
     
diff --git a/utils.py b/utils.py
index 5244f13..c76247b 100644
--- a/utils.py
+++ b/utils.py
@@ -78,9 +78,7 @@ def dump_pkl_files(
         file_name = os.path.basename(file_path)
         os.makedirs(os.path.dirname(out_file_path), exist_ok=True)
         
-        arrows: Dict[str, List[Tensor]] = {}
-        for key in output_keys:
-            arrows[key] = []
+        arrows: List[Dict[str, Tensor]] = []
         
         with open(file_path, "r") as f:    
             lines = f.readlines()
@@ -88,16 +86,20 @@ def dump_pkl_files(
         for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
             line_dict = json.loads(line)
             arrow = process_func(line_dict)
-            for key in output_keys:
-                arrows[key].extend(arrow[key])
+            arrows.append(arrow)
+        
+        package: Dict[str, List[Tensor]] = {}
+        for key in output_keys:
+            list_tensor = [arrow[key] for arrow in arrows]
+            package[key] = list_tensor
             
         output_package: Dict[str, Tensor] = {}
         
         for key in output_keys:
             if packing_size > 0:
                 print(f"Packaging key: '{key}'")
-                arrows[key] = pack_sequences(arrows[key], packing_size, pad_value)
-            sequence = torch.cat(arrows[key])
+                package[key] = pack_sequences(package[key], packing_size, pad_value)
+            sequence = torch.cat(package[key])
             output_package[key] = sequence
          
         with open(out_file_path, "w") as f: