From faf9f81e7cea8e616591e8807a8f44b45e587944 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 16 Jul 2025 12:17:58 +0800 Subject: [PATCH] =?UTF-8?q?refactor(utils):=20=E5=B0=86=20pre=5Ftarin=5Fpr?= =?UTF-8?q?ocess=20=E5=92=8C=20sft=5Fprocess=20=E5=87=BD=E6=95=B0=E5=90=88?= =?UTF-8?q?=E5=B9=B6=E4=B8=BA=20process=5Fdataset?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- belle_sft.py | 4 ++-- chinese-c4.py | 4 ++-- chinese-cosmopedia.py | 4 ++-- english-fineweb.py | 4 ++-- english-wiki.py | 4 ++-- utils.py | 51 ++++++++++--------------------------------- 6 files changed, 21 insertions(+), 50 deletions(-) diff --git a/belle_sft.py b/belle_sft.py index c4e1258..4d6a521 100644 --- a/belle_sft.py +++ b/belle_sft.py @@ -1,9 +1,9 @@ from datasets import load_dataset -from utils import pre_tarin_process +from utils import process_dataset if __name__ == "__main__": dataset = load_dataset("BelleGroup/train_3.5M_CN") - # pre_tarin_process( + # process_dataset( # dataset_dict=dataset, # output_subdir="belle_sft", # max_chunk_size=5, diff --git a/chinese-c4.py b/chinese-c4.py index 154f342..9344039 100644 --- a/chinese-c4.py +++ b/chinese-c4.py @@ -1,9 +1,9 @@ from datasets import load_dataset -from utils import pre_tarin_process +from utils import process_dataset if __name__ == "__main__": dataset = load_dataset("shjwudp/chinese-c4") - pre_tarin_process( + process_dataset( dataset_dict=dataset, output_subdir="chinese-c4" ) \ No newline at end of file diff --git a/chinese-cosmopedia.py b/chinese-cosmopedia.py index d22cfa0..94d7ba0 100644 --- a/chinese-cosmopedia.py +++ b/chinese-cosmopedia.py @@ -1,5 +1,5 @@ from datasets import load_dataset -from utils import pre_tarin_process +from utils import process_dataset if __name__ == "__main__": max_chunk_num = 10 @@ -10,7 +10,7 @@ if __name__ == "__main__": data_files={"train": [f"data/0000{i}.parquet" for i in range(5)]} ) - pre_tarin_process( + process_dataset( dataset_dict=dataset, output_subdir="chinese-wiki", max_chunk_num=max_chunk_num, diff --git a/english-fineweb.py b/english-fineweb.py index 855741c..e71a42d 100644 --- a/english-fineweb.py +++ b/english-fineweb.py @@ -1,9 +1,9 @@ from datasets import load_dataset -from utils import pre_tarin_process +from utils import process_dataset if __name__ == "__main__": dataset = load_dataset("HuggingFaceFW/fineweb", "sample-10BT") - pre_tarin_process( + process_dataset( dataset_dict=dataset, output_subdir="english-fineweb", ) \ No newline at end of file diff --git a/english-wiki.py b/english-wiki.py index c7e180d..e15f265 100644 --- a/english-wiki.py +++ b/english-wiki.py @@ -1,9 +1,9 @@ from datasets import load_dataset -from utils import pre_tarin_process +from utils import process_dataset if __name__ == "__main__": dataset = load_dataset("Blaze7451/enwiki_structured_content") - pre_tarin_process( + process_dataset( dataset_dict=dataset, output_subdir="english-wiki", max_chunk_size=5, diff --git a/utils.py b/utils.py index 05057ac..8dc740d 100644 --- a/utils.py +++ b/utils.py @@ -57,13 +57,14 @@ def dump_pkl_files( tensor = torch.cat(arrows) pkl.dump(tensor, f) -def pre_tarin_process( +def process_dataset( dataset_dict: DatasetDict, output_subdir: str, max_chunk_num: int = None, chunk_size: int = 1000000, split_name: str = "train", column_name: str = "text", + process_func: Callable[[dict], dict] = None, normalization_func=comprehensive_normalization, ): train_dataset = dataset_dict[split_name] @@ -83,43 +84,13 @@ def pre_tarin_process( 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[column_name] - if normalization_func: - text = normalization_func(text) - json_line = {column_name : text} - f.write(json.dumps(json_line, ensure_ascii=False) + "\n") + if process_func is not None: + 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") - print(f"Saved text chunk {i} to {output_path}") - - -def sft_process( - dataset_dict: DatasetDict, - output_subdir: str, - max_chunk_num: int = None, - chunk_size: int = 1000000, - split_name: str = "train", - processsor: Callable[[str], str] = None, -): - 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 + print(f"Saved text chunk {i} to {output_path}")