feat(chinese-instruct): 重新处理数据集并生成 SFT 训练文件

dump_sft_file.py

@@ -1,7 +1,26 @@
 from utils import dump_pkl_files, fetch_files
 from tokenizer import BpeTokenizer
+import torch
 import os
 
+def get_processor(tokenizer: BpeTokenizer):
+    def processor(input_dict: dict):
+        query, response = input_dict["query"], input_dict["response"]
+        prefix_seg = f"<|user|> {query} <|system|> <bos>"
+        suffix_seg = f"{response}<eos>\n"
+        prefix_ids = tokenizer.encode(prefix_seg)
+        suffix_ids = tokenizer.encode(suffix_seg)
+        
+        tokens = prefix_ids + suffix_ids
+        tokens = torch.tensor(tokens, dtype=torch.int32)
+        masks = torch.zeros_like(tokens)
+        masks[:len(prefix_ids)] = 1
+        
+        return tokens, masks
+    
+    return processor
+
+
 if __name__ == "__main__":
     tokenizer = BpeTokenizer("tokenizer.json")
     base_dir = [
@@ -12,5 +31,7 @@ if __name__ == "__main__":
     files = []
     for dir_path in base_dir:
         files.extend(fetch_files(dir_path))
+        
+    processor = get_processor(tokenizer)
     
-    dump_pkl_files(tokenizer, files, base_out_dir, packing_size=2048)
+    dump_pkl_files(files, base_out_dir,)