diff --git a/chinese-cosmopedia.py b/chinese-cosmopedia.py
index 6f55162..95bb2f3 100644
--- a/chinese-cosmopedia.py
+++ b/chinese-cosmopedia.py
@@ -6,9 +6,9 @@ if __name__ == "__main__":
     item_size = max_chunk_num * chunk_size
     
     process_dataset(
-        dataset_name="chinese-cosmopedia",
+        dataset_name="opencsg/chinese-cosmopedia",
         output_subdir="chinese-wiki",
-        split=f"train[:{item_size}]",
+        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/utils.py b/utils.py
index a2441b4..96b63b4 100644
--- a/utils.py
+++ b/utils.py
@@ -63,6 +63,7 @@ def process_dataset(
     dataset_config: str = None,
     max_chunk_num: int = None,
     split: str = None,
+    data_files: List[str] = None,
     split_name: str = "train",
     column_name: str = "text",
     chunk_size: int = 1000000,
@@ -70,9 +71,10 @@ def process_dataset(
 ):
 
     dataset_dict = load_dataset(
-        data_dir=dataset_name, 
-        data_files=dataset_config,
-        split=split
+        path=dataset_name, 
+        name=dataset_config,
+        split=split,
+        data_files=data_files
     )
     
     train_dataset = dataset_dict[split_name]