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