diff --git a/chinese-instruct.py b/chinese-instruct.py
index 98cb7f1..d3048db 100644
--- a/chinese-instruct.py
+++ b/chinese-instruct.py
@@ -3,25 +3,23 @@ from datasets import load_dataset, concatenate_datasets
 from utils import process_dataset
 
 
-def build_prompt(query:str, response:str) -> str:
+def replace_seg(query:str, response:str) -> str:
     replacements = {
-        "\\[": "$$",
-        "\\]": "$$",
-        "\\(": "$",
-        "\\)": "$"
+        "\\[": "$$", "\\]": "$$",
+        "\\(": "$", "\\)": "$"
     }
-
     for old, new in replacements.items():
         query = query.replace(old, new)
         response = response.replace(old, new)
-         
-    return f"<|user|> {query} <|system|> <bos>{response}<eos>\n"
 
+    return {"query": query, "response": response}
 
 def process_func(input_dict: dict):
     query = input_dict["prompt"] if input_dict["prompt"] else ""
-    response = input_dict["response"] if input_dict["response"] else ""
-    return {"text": build_prompt(query, response)}
+    resp = input_dict["response"] if input_dict["response"] else ""
+    query, resp = replace_seg(query, resp)
+    
+    return {"query": query, "response": resp }
 
 
 if __name__ == "__main__":
diff --git a/dump_sft_file.py b/dump_sft_file.py
index 6038797..70b2670 100644
--- a/dump_sft_file.py
+++ b/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)
\ No newline at end of file
+    dump_pkl_files(files, base_out_dir,)
\ No newline at end of file