refactor: 重构数据管道模块

pre_train/chinese-c4.py

@@ -1,9 +1,11 @@
 from datasets import load_dataset
-from modules.utils import process_dataset
+from modules.datapipeline import DataPipeline
 
 if __name__ == "__main__":
     dataset = load_dataset("shjwudp/chinese-c4")
-    process_dataset(
+    
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=dataset,
         output_subdir="chinese-c4-pretrain"
     )

pre_train/chinese-cosmopedia.py

@@ -1,5 +1,5 @@
 from datasets import load_dataset
-from modules.utils import process_dataset
+from modules.datapipeline import DataPipeline
 
 if __name__ == "__main__":
     chunk_size = 1000000
@@ -9,7 +9,8 @@ if __name__ == "__main__":
         data_files={"train": [f"data/000{i:02d}.parquet" for i in range(25)]}
     )
     
-    process_dataset(
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=dataset,
         output_subdir="chinese-wiki-pretrain",
         chunk_size=chunk_size,

pre_train/english-fineweb.py

@@ -1,9 +1,11 @@
 from datasets import load_dataset
-from modules.utils import process_dataset
+from modules.datapipeline import DataPipeline
 
 if __name__ == "__main__":
     dataset = load_dataset("HuggingFaceFW/fineweb", "sample-10BT")
-    process_dataset(
+    
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=dataset,
         output_subdir="english-fineweb-pretrain",
     )

pre_train/english-wiki.py

@@ -1,10 +1,12 @@
 from datasets import load_dataset
-from modules.utils import process_dataset
+from modules.datapipeline import DataPipeline
 
 if __name__ == "__main__":
     dataset = load_dataset("Blaze7451/enwiki_structured_content")
-    process_dataset(
+    
+    pipeline = DataPipeline()
+    pipeline.process_dataset(
         dataset_dict=dataset,
         output_subdir="english-wiki-pretrain",
-        max_chunk_size=5,
+        max_chunk_num=5,
     )