diff --git a/chinese-cosmopedia.py b/chinese-cosmopedia.py
index f451dc2..94d7ba0 100644
--- a/chinese-cosmopedia.py
+++ b/chinese-cosmopedia.py
@@ -7,7 +7,7 @@ if __name__ == "__main__":
     
     dataset = load_dataset(
         "opencsg/chinese-cosmopedia", 
-        data_files=[f"0000{i}.parquet" for i in range(5)]
+        data_files={"train": [f"data/0000{i}.parquet" for i in range(5)]}
     )
     
     process_dataset(
diff --git a/utils.py b/utils.py
index 83ea164..d62d4ac 100644
--- a/utils.py
+++ b/utils.py
@@ -66,7 +66,6 @@ def process_dataset(
     column_name: str = "text",
     normalization_func=comprehensive_normalization,
 ):
-
     train_dataset = dataset_dict[split_name]
     total_samples = len(train_dataset)
     num_chunks = (total_samples // chunk_size) + 1