From e6f5b26a75367ef22945521f738e382303d5cd2c Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Tue, 15 Jul 2025 12:54:31 +0800
Subject: [PATCH] =?UTF-8?q?refactor(chinese-cosmopedia):=20=E6=9B=B4?=
 =?UTF-8?q?=E6=96=B0=E6=95=B0=E6=8D=AE=E9=9B=86=E5=8A=A0=E8=BD=BD=E6=96=B9?=
 =?UTF-8?q?=E5=BC=8F=E5=B9=B6=E8=B0=83=E6=95=B4=E7=9B=B8=E5=85=B3=E5=8F=82?=
 =?UTF-8?q?=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 chinese-cosmopedia.py | 4 ++--
 utils.py              | 8 +++++---
 2 files changed, 7 insertions(+), 5 deletions(-)

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]