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

belle_sft.py

@@ -0,0 +1,10 @@
+from datasets import load_dataset
+from utils import pre_tarin_process
+
+if __name__ == "__main__":
+    dataset = load_dataset("BelleGroup/train_3.5M_CN")
+    # pre_tarin_process(
+    #     dataset_dict=dataset,
+    #     output_subdir="belle_sft",
+    #     max_chunk_size=5,
+    # )

chinese-c4.py

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

chinese-cosmopedia.py

@@ -1,5 +1,5 @@
 from datasets import load_dataset
-from utils import process_dataset
+from utils import pre_tarin_process
 
 if __name__ == "__main__":
     max_chunk_num = 10
@@ -10,7 +10,7 @@ if __name__ == "__main__":
         data_files={"train": [f"data/0000{i}.parquet" for i in range(5)]}
     )
     
-    process_dataset(
+    pre_tarin_process(
         dataset_dict=dataset,
         output_subdir="chinese-wiki",
         max_chunk_num=max_chunk_num,

english-fineweb.py

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

english-wiki.py

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

pre_train_file.py

@@ -8,10 +8,10 @@ def processor(intput_str: str):
 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"),
+        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 = []

utils.py

@@ -57,7 +57,7 @@ def dump_pkl_files(
             tensor = torch.cat(arrows)
             pkl.dump(tensor, f)
             
-def process_dataset(
+def pre_tarin_process(
     dataset_dict: DatasetDict,
     output_subdir: str,
     max_chunk_num: int = None,
@@ -90,3 +90,10 @@ def process_dataset(
                 f.write(json.dumps(json_line, ensure_ascii=False) + "\n")
 
         print(f"Saved text chunk {i} to {output_path}")        
+
+
+def sft_process(
+    datset_dict: DatasetDict,
+    output_subdir: str,
+):
+    pass