From b69ba2c5ff3b35808827d18dc885554198284424 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Wed, 16 Jul 2025 09:56:31 +0800
Subject: [PATCH] =?UTF-8?q?refactor(dataset):=20=E9=87=8D=E6=9E=84?=
 =?UTF-8?q?=E6=95=B0=E6=8D=AE=E9=9B=86=E5=A4=84=E7=90=86=E6=B5=81=E7=A8=8B?=
 =?UTF-8?q?=E5=B9=B6=E6=B7=BB=E5=8A=A0=E6=96=B0=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 belle_sft.py                   | 10 ++++++++++
 chinese-c4.py                  |  4 ++--
 chinese-cosmopedia.py          |  4 ++--
 english-fineweb.py             |  4 ++--
 english-wiki.py                |  4 ++--
 to_ids.py => pre_train_file.py |  8 ++++----
 sft_file.py                    |  0
 utils.py                       |  9 ++++++++-
 8 files changed, 30 insertions(+), 13 deletions(-)
 create mode 100644 belle_sft.py
 rename to_ids.py => pre_train_file.py (68%)
 create mode 100644 sft_file.py

diff --git a/belle_sft.py b/belle_sft.py
new file mode 100644
index 0000000..c4e1258
--- /dev/null
+++ b/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,
+    # )
\ No newline at end of file
diff --git a/chinese-c4.py b/chinese-c4.py
index 9344039..154f342 100644
--- a/chinese-c4.py
+++ b/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"
     )
\ No newline at end of file
diff --git a/chinese-cosmopedia.py b/chinese-cosmopedia.py
index 94d7ba0..d22cfa0 100644
--- a/chinese-cosmopedia.py
+++ b/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,
diff --git a/english-fineweb.py b/english-fineweb.py
index e71a42d..855741c 100644
--- a/english-fineweb.py
+++ b/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",
     )
\ No newline at end of file
diff --git a/english-wiki.py b/english-wiki.py
index e15f265..c7e180d 100644
--- a/english-wiki.py
+++ b/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,
diff --git a/to_ids.py b/pre_train_file.py
similarity index 68%
rename from to_ids.py
rename to pre_train_file.py
index ee0adf7..1822141 100644
--- a/to_ids.py
+++ b/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 = []
diff --git a/sft_file.py b/sft_file.py
new file mode 100644
index 0000000..e69de29
diff --git a/utils.py b/utils.py
index d62d4ac..4abf9ea 100644
--- a/utils.py
+++ b/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
\ No newline at end of file