refactor(data_processing): 重构数据处理脚本并添加 fetch_folders 函数

dump_pt_file.py

@@ -1,16 +1,10 @@
-from modules.utils import dump_pkl_files, fetch_files, get_pt_processor
+from modules.utils import dump_pkl_files, fetch_files, fetch_folders, get_pt_processor
 from modules.tokenizer import BpeTokenizer
-import os
 
 
 if __name__ == "__main__":
     tokenizer = BpeTokenizer("tokenizer.json")
-    base_dir = [
-        os.path.join("dataset", "chinese-c4"),
-        os.path.join("dataset", "english-fineweb"),
-        os.path.join("dataset", "english-wiki"),
-        os.path.join("dataset", "chinese-wiki"),
-    ]
+    base_dir = fetch_folders("dataset")
     
     base_out_dir = "pkl_output"
     files = []

dump_sft_file.py

@@ -1,19 +1,14 @@
-from modules.utils import dump_pkl_files, fetch_files, get_sft_processor
+from modules.utils import dump_pkl_files, fetch_files, fetch_folders, get_sft_processor
 from modules.tokenizer import BpeTokenizer
-import os
-
 
 if __name__ == "__main__":
     tokenizer = BpeTokenizer("tokenizer.json")
-    base_dir = [
-        os.path.join("dataset", "Ling-Coder-SFT"),
-        os.path.join("dataset", "chinese-instruct")
-    ]
+    base_dir = fetch_folders("dataset")
+    
     base_out_dir = "pkl_output"
     files = []
     for dir_path in base_dir:
         files.extend(fetch_files(dir_path))
         
     processor = get_sft_processor(tokenizer)
-    
     dump_pkl_files(files, base_out_dir,processor, ["sequence", "mask"])

modules/utils.py

@@ -16,6 +16,16 @@ def fetch_files(directory):
             for root, _, files in os.walk(directory) for f in files]
     
 
+def fetch_folders(root_dir, filter_func=None):
+    folders = []
+    for root, dirs, _ in os.walk(root_dir):
+        for dir_name in dirs:
+            folder_path = os.path.join(root, dir_name)
+            if filter_func is None or filter_func(folder_path):
+                folders.append(folder_path)
+    return folders
+
+
 def comprehensive_normalization(text):
     replacements = {
         '\u2018': "'", '\u2019': "'", '\u0060': "'",