From f64a1a15ee12db1d327efb242f47faa6c7ec4f70 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 16 Jul 2025 11:35:08 +0800 Subject: [PATCH] =?UTF-8?q?feat(utils):=20=E6=96=B0=E5=A2=9E=20sft=5Fproce?= =?UTF-8?q?ss=20=E5=87=BD=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- utils.py | 30 ++++++++++++++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/utils.py b/utils.py index 4abf9ea..05057ac 100644 --- a/utils.py +++ b/utils.py @@ -93,7 +93,33 @@ def pre_tarin_process( def sft_process( - datset_dict: DatasetDict, + dataset_dict: DatasetDict, output_subdir: str, + max_chunk_num: int = None, + chunk_size: int = 1000000, + split_name: str = "train", + processsor: Callable[[str], str] = None, ): - pass \ No newline at end of file + train_dataset = dataset_dict[split_name] + total_samples = len(train_dataset) + num_chunks = (total_samples // chunk_size) + 1 + lim_chunks = min(max_chunk_num, num_chunks) if max_chunk_num 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(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)) + + 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 processsor is not None: + example = processsor(example) + f.write(json.dumps(example, ensure_ascii=False) + "\n") + + print(f"Saved text chunk {i} to {output_path}") \ No newline at end of file