From b588c30f095d52aae38de3f5bdfa05f8664aba68 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Mon, 21 Jul 2025 22:18:22 +0800 Subject: [PATCH] =?UTF-8?q?refactor(utils):=20=E4=BC=98=E5=8C=96=20process?= =?UTF-8?q?=5Fdataset=20=E5=87=BD=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- utils.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/utils.py b/utils.py index 97418e9..0f7b11d 100644 --- a/utils.py +++ b/utils.py @@ -109,7 +109,7 @@ def process_dataset( chunk_size: int = 1000000, split_name: str = "train", column_name: str = "text", - process_func: Callable[[Union[dict, List[dict]]], dict] = None, + process_func: Union[Callable[[dict], dict], Callable[[List[dict]], List[dict]]] = None, normalization_func=comprehensive_normalization, ): train_dataset = dataset_dict[split_name] @@ -129,13 +129,19 @@ 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: - if process_func is not None: + if process_func: processed_example = process_func(example) else: text = example[column_name] if normalization_func: text = normalization_func(text) processed_example = {column_name: text} - f.write(json.dumps(processed_example, ensure_ascii=False) + "\n") + + if isinstance(processed_example, dict): + f.write(json.dumps(processed_example, ensure_ascii=False) + "\n") + elif isinstance(processed_example, list): + for item in processed_example: + f.write(json.dumps(item, ensure_ascii=False) + "\n") + print(f"Saved text chunk {i} to {output_path}")