refactor(utils): 重构数据处理脚本

chinese-c4.py

@@ -1,7 +1,9 @@
+from datasets import load_dataset
 from utils import process_dataset
 
 if __name__ == "__main__":
+    dataset = load_dataset("shjwudp/chinese-c4")
     process_dataset(
-        dataset_name="shjwudp/chinese-c4",
+        dataset_dict=dataset,
         output_subdir="chinese-c4"
     )

chinese-cosmopedia.py

@@ -1,14 +1,18 @@
+from datasets import load_dataset
 from utils import process_dataset
 
 if __name__ == "__main__":
     max_chunk_num = 10
     chunk_size = 1000000
-    item_size = max_chunk_num * chunk_size
+    
+    dataset = load_dataset(
+        "opencsg/chinese-cosmopedia", 
+        data_files=[f"0000{i}.parquet" for i in range(5)]
+    )
     
     process_dataset(
-        dataset_name="opencsg/chinese-cosmopedia",
+        dataset_dict=dataset,
         output_subdir="chinese-wiki",
-        data_files=[f"0000{i}.parquet" for i in range(5)],
         max_chunk_num=max_chunk_num,
         chunk_size=chunk_size,
     )

english-fineweb.py

@@ -1,8 +1,9 @@
+from datasets import load_dataset
 from utils import process_dataset
 
 if __name__ == "__main__":
+    dataset = load_dataset("HuggingFaceFW/fineweb", "sample-10BT")
     process_dataset(
-        dataset_name="HuggingFaceFW/fineweb",
+        dataset_dict=dataset,
         output_subdir="english-fineweb",
-        dataset_config="sample-10BT",
     )

english-wiki.py

@@ -1,8 +1,10 @@
+from datasets import load_dataset
 from utils import process_dataset
 
 if __name__ == "__main__":
+    dataset = load_dataset("Blaze7451/enwiki_structured_content")
     process_dataset(
-        dataset_name="Blaze7451/enwiki_structured_content",
+        dataset_dict=dataset,
         output_subdir="english-wiki",
         max_chunk_size=5,
     )

utils.py

@@ -1,5 +1,5 @@
 from typing import List, Callable
-from datasets import load_dataset
+from datasets import DatasetDict
 from tokenizer import BpeTokenizer
 from tqdm import tqdm
 from torch import Tensor
@@ -58,25 +58,15 @@ def dump_pkl_files(
             pkl.dump(tensor, f)
             
 def process_dataset(
-    dataset_name: str,
+    dataset_dict: DatasetDict,
     output_subdir: str,
-    dataset_config: str = None,
     max_chunk_num: int = None,
-    split: str = None,
+    chunk_size: int = 1000000,
-    data_files: List[str] = None,
     split_name: str = "train",
     column_name: str = "text",
-    chunk_size: int = 1000000,
+    normalization_func=comprehensive_normalization,
-    normalization_func=comprehensive_normalization
 ):
 
-    dataset_dict = load_dataset(
-        path=dataset_name, 
-        name=dataset_config,
-        split=split,
-        data_files=data_files
-    )
-    
     train_dataset = dataset_dict[split_name]
     total_samples = len(train_dataset)
     num_chunks = (total_samples // chunk_size) + 1