refactor(dataset): 重构数据集处理流程并添加新功能

2025-07-16 09:56:31 +08:00
parent 6edc31ce2e
commit b69ba2c5ff
8 changed files with 30 additions and 13 deletions
@@ -0,0 +1,21 @@
+from utils import dump_pkl_files, fetch_files
+from tokenizer import BpeTokenizer
+import os
+
+def processor(intput_str: str):
+    return f"{intput_str} <eos>"
+
+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_out_dir = "pkl_output"
+    files = []
+    for dir_path in base_dir:
+        files.extend(fetch_files(dir_path))
+    
+    dump_pkl_files(tokenizer, files, base_out_dir, processor)