refactor(dataset): 重构数据集处理逻辑

chinese-c4.py

@@ -1,28 +1,7 @@
-from datasets import load_dataset
-import json
-import os
+from utils import process_dataset
 
 if __name__ == "__main__":
-    dataset_dict = load_dataset("shjwudp/chinese-c4")
-    train_dataset = dataset_dict["train"]
-
-    chunk_size = 1000000
-    total_samples = len(train_dataset)
-    num_chunks = (total_samples // chunk_size) + 1
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-    output_dir = os.path(script_dir, "dataset", "chinese-c4")
-    os.makedirs(output_dir, exist_ok=True)
-
-    for i in range(num_chunks):
-        start_idx = i * chunk_size
-        end_idx = min((i + 1) * chunk_size, total_samples)
-        chunk = train_dataset.select(range(start_idx, end_idx))
-
-        output_path = f"{output_dir}/chinese-c4_text_chunk_{i}.jsonl"
-        with open(output_path, "w", encoding="utf-8") as f:
-            for example in chunk:
-                # 每行写入一个 {"text": "xxx"} 对象
-                json_line = {"text": example["text"]}
-                f.write(json.dumps(json_line, ensure_ascii=False) + "\n")
-
-        print(f"Saved text chunk {i} to {output_path}")
+    process_dataset(
+        dataset_name="shjwudp/chinese-c4",
+        output_subdir="chinese-c4"
+    )

chinese-wiki.py

@@ -1,28 +1,7 @@
-from datasets import load_dataset
-import json
-import os
-
+from utils import process_dataset
 
 if __name__ == "__main__":
-    dataset_dict = load_dataset("Blaze7451/Wiki-zh-20250601")
-    train_dataset = dataset_dict["train"]
-    chunk_size = 1000000
-    total_samples = len(train_dataset)
-    num_chunks = (total_samples // chunk_size) + 1
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-    output_dir = os.path.join(script_dir, "dataset", "chinese-wiki")
-    os.makedirs(output_dir, exist_ok=True)
-    
-    
-    for i in range(num_chunks):
-        start_idx = i * chunk_size
-        end_idx = min((i + 1) * chunk_size, total_samples)
-        chunk = train_dataset.select(range(start_idx, end_idx))
-
-        output_path = f"{output_dir}/chinese-wiki_text_chunk_{i}.jsonl"
-        with open(output_path, "w", encoding="utf-8") as f:
-            for example in chunk:
-                json_line = {"text": example["text"]}
-                f.write(json.dumps(json_line, ensure_ascii=False) + "\n")
-
-        print(f"Saved text chunk {i} to {output_path}")
+    process_dataset(
+        dataset_name="Blaze7451/Wiki-zh-20250601",
+        output_subdir="chinese-wiki"
+    )

english-fineweb.py

@@ -1,6 +1,4 @@
-from datasets import load_dataset
-import json
-import os
+from utils import process_dataset
 import re
 
 def comprehensive_normalization(text):
@@ -14,29 +12,11 @@ def comprehensive_normalization(text):
     pattern = re.compile('|'.join(re.escape(k) for k in replacements))
     return pattern.sub(lambda m: replacements[m.group()], text)
 
+
 if __name__ == "__main__":
-    dataset_dict = load_dataset("HuggingFaceFW/fineweb","sample-10BT")
-    train_dataset = dataset_dict["train"]
-
-    chunk_size = 1000000
-    total_samples = len(train_dataset)
-    num_chunks = (total_samples // chunk_size) + 1
-
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-    output_dir = os.path(script_dir, "dataset", "english-fineweb")
-    os.makedirs(output_dir, exist_ok=True)
-
-    for i in range(num_chunks):
-        if i == 10: 
-            break
-        start_idx = i * chunk_size
-        end_idx = min((i + 1) * chunk_size, total_samples)
-        chunk = train_dataset.select(range(start_idx, end_idx))
-        
-        output_path = f"{output_dir}/english-fineweb_text_chunk_{i}.jsonl"
-        with open(output_path, "w", encoding="utf-8") as f:
-            for example in chunk:
-                json_line = {"text": comprehensive_normalization(example["text"])}
-                f.write(json.dumps(json_line, ensure_ascii=False) + "\n")
-        
-        print(f"Saved text chunk {i} to {output_path}")
+    process_dataset(
+        dataset_name="HuggingFaceFW/fineweb",
+        output_subdir="english-fineweb",
+        dataset_config="sample-10BT",
+        normalization_func=comprehensive_normalization
+    )

utils.py

@@ -0,0 +1,39 @@
+from datasets import load_dataset
+import json
+import os
+
+def process_dataset(
+    dataset_name: str,
+    output_subdir: str,
+    dataset_config: str = None,
+    split_name: str = "train",
+    chunk_size: int = 1000000,
+    normalization_func=None
+):
+
+    dataset_dict = load_dataset(dataset_name, dataset_config) if dataset_config else load_dataset(dataset_name)
+    train_dataset = dataset_dict[split_name]
+
+    total_samples = len(train_dataset)
+    num_chunks = (total_samples // chunk_size) + 1
+
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    output_dir = os.path.join(script_dir, "dataset", output_subdir)
+    os.makedirs(output_dir, exist_ok=True)
+
+    for i in range(num_chunks):
+        start_idx = i * chunk_size
+        end_idx = min((i + 1) * chunk_size, total_samples)
+        chunk = train_dataset.select(range(start_idx, end_idx))
+
+        output_path = os.path.join(output_dir, f"{output_subdir}_text_chunk_{i}.jsonl")
+        with open(output_path, "w", encoding="utf-8") as f:
+            for example in chunk:
+                text = example["text"]
+                if normalization_func:
+                    text = normalization_func(text)
+                json_line = {"text": text}
+                f.write(json.dumps(json_line, ensure_ascii=False) + "\n")
+
+        print(f"Saved text chunk {i} to {output_path}")
+