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

chinese-instruct.py

@@ -12,7 +12,7 @@ def replace_seg(query:str, response:str) -> str:
         query = query.replace(old, new)
         response = response.replace(old, new)
 
-    return {"query": query, "response": response}
+    return query, response
 
 def process_func(input_dict: dict):
     query = input_dict["prompt"] if input_dict["prompt"] else ""

dump_sft_file.py

@@ -16,7 +16,7 @@ def get_processor(tokenizer: BpeTokenizer):
         masks = torch.zeros_like(tokens)
         masks[:len(prefix_ids)] = 1
         
-        return tokens, masks
+        return {"sequence": tokens, "mask": masks}
     
     return processor
 
@@ -34,4 +34,4 @@ if __name__ == "__main__":
         
     processor = get_processor(tokenizer)
     
-    dump_pkl_files(files, base_out_dir,)
+    dump_pkl_files(files, base_out_dir,processor, ["sequence", "mask"])

utils.py

@@ -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])