From e77fb56c452b48e485b673746c79961bb5338ad3 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Tue, 15 Jul 2025 13:10:38 +0800
Subject: [PATCH] =?UTF-8?q?refactor(utils):=20=E9=87=8D=E6=9E=84=E6=95=B0?=
 =?UTF-8?q?=E6=8D=AE=E5=A4=84=E7=90=86=E8=84=9A=E6=9C=AC?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 chinese-c4.py         |  4 +++-
 chinese-cosmopedia.py | 10 +++++++---
 english-fineweb.py    |  5 +++--
 english-wiki.py       |  4 +++-
 utils.py              | 18 ++++--------------
 5 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/chinese-c4.py b/chinese-c4.py
index 008ad32..9344039 100644
--- a/chinese-c4.py
+++ b/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"
     )
\ No newline at end of file
diff --git a/chinese-cosmopedia.py b/chinese-cosmopedia.py
index 95bb2f3..f451dc2 100644
--- a/chinese-cosmopedia.py
+++ b/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,
     )
\ No newline at end of file
diff --git a/english-fineweb.py b/english-fineweb.py
index 4d53dab..e71a42d 100644
--- a/english-fineweb.py
+++ b/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",
     )
\ No newline at end of file
diff --git a/english-wiki.py b/english-wiki.py
index cdeca88..e15f265 100644
--- a/english-wiki.py
+++ b/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,
     )
\ No newline at end of file
diff --git a/utils.py b/utils.py
index 96b63b4..83ea164 100644
--- a/utils.py
+++ b/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,
-    data_files: List[str] = None,
+    chunk_size: int = 1000000,
     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