From 820ddf3fc8d60127483fbb1b5a3b72ab84e8a39b Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Mon, 14 Jul 2025 13:34:03 +0800
Subject: [PATCH] =?UTF-8?q?feat(utils):=20=E5=A2=9E=E5=8A=A0=E6=95=B0?=
 =?UTF-8?q?=E6=8D=AE=E9=9B=86=E5=A4=84=E7=90=86=E7=9A=84=E7=81=B5=E6=B4=BB?=
 =?UTF-8?q?=E6=80=A7=E5=92=8C=E5=8F=AF=E9=85=8D=E7=BD=AE=E6=80=A7?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 utils.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/utils.py b/utils.py
index 610ab00..656f3d1 100644
--- a/utils.py
+++ b/utils.py
@@ -19,7 +19,9 @@ def process_dataset(
     dataset_name: str,
     output_subdir: str,
     dataset_config: str = None,
+    max_chunk_size: int = None,
     split_name: str = "train",
+    column_name: str = "text",
     chunk_size: int = 1000000,
     normalization_func=comprehensive_normalization
 ):
@@ -29,12 +31,13 @@ def process_dataset(
 
     total_samples = len(train_dataset)
     num_chunks = (total_samples // chunk_size) + 1
+    lim_chunks = max_chunk_size if max_chunk_size else num_chunks
 
     script_dir = os.path.dirname(os.path.abspath(__file__))
     output_dir = os.path.join(script_dir, "dataset", output_subdir)
     os.makedirs(output_dir, exist_ok=True)
 
-    for i in range(num_chunks):
+    for i in range(lim_chunks):
         start_idx = i * chunk_size
         end_idx = min((i + 1) * chunk_size, total_samples)
         chunk = train_dataset.select(range(start_idx, end_idx))
@@ -42,10 +45,10 @@ def process_dataset(
         output_path = os.path.join(output_dir, f"{output_subdir}_text_chunk_{i}.jsonl")
         with open(output_path, "w", encoding="utf-8") as f:
             for example in chunk:
-                text = example["text"]
+                text = example[column_name]
                 if normalization_func:
                     text = normalization_func(text)
-                json_line = {"text": text}
+                json_line = {column_name : text}
                 f.write(json.dumps(json_line, ensure_ascii=False) + "\n")
 
         print(f"Saved text chunk {i} to {output_path}")