refactor(sft_file, utils): 重构数据处理流程并支持新数据集

sft_file.py

@@ -5,7 +5,8 @@ import os
 if __name__ == "__main__":
     tokenizer = BpeTokenizer("tokenizer.json")
     base_dir = [
-        os.path.join("dataset", "belle-sft"),
+        # os.path.join("dataset", "belle-sft"),
+        os.path.join("dataset", "chinese-instruct")
     ]
     base_out_dir = "pkl_output"
     files = []

utils.py

@@ -27,21 +27,18 @@ def comprehensive_normalization(text):
     pattern = re.compile('|'.join(re.escape(k) for k in replacements))
     return pattern.sub(lambda m: replacements[m.group()], text)      
 
-   
 def dump_pkl_files(
     tokenizer: BpeTokenizer,
     files: List[str], 
     base_out_dir: str,
-    encoder: Callable[[str], str]=None,
-    key: str='text',
-    packing_size: int=None
+    process_func: Callable[[dict], str],
+    packing_size: int = -1
 ):
     def process_line(line: str) -> Tensor:
-            line = json.loads(line)[key]
-            processed_line = encoder(line) if encoder else line
-            ids = tokenizer.encode(processed_line)
-            arrow = torch.tensor(ids, dtype=torch.int32)
-            return arrow
+        dict_line = json.loads(line)
+        tokens = process_func(dict_line)
+        ids = tokenizer.encode(tokens)
+        return torch.tensor(ids, dtype=torch.int32)
         
     for file_path in files:
         out_file_name = os.path.basename(file_path).replace(".jsonl", ".pkl")
@@ -55,8 +52,7 @@ def dump_pkl_files(
         for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
             arrow = process_line(line)
             arrows.append(arrow)
-        
-        if packing_size is None:
+            if packing_size > 0:
                 with open(out_file_path, "wb") as f:
                     package_tensor = torch.cat(arrows)
                     pkl.dump(package_tensor, f)
@@ -90,7 +86,6 @@ def dump_pkl_files(
                     pkl.dump(package_tensor, f)
                     
             
-            
 def process_dataset(
     dataset_dict: DatasetDict,
     output_subdir: str,