diff --git a/chinese-c4.py b/chinese-c4.py index 172270c..008ad32 100644 --- a/chinese-c4.py +++ b/chinese-c4.py @@ -1,28 +1,7 @@ -from datasets import load_dataset -import json -import os +from utils import process_dataset if __name__ == "__main__": - dataset_dict = load_dataset("shjwudp/chinese-c4") - train_dataset = dataset_dict["train"] - - chunk_size = 1000000 - total_samples = len(train_dataset) - num_chunks = (total_samples // chunk_size) + 1 - script_dir = os.path.dirname(os.path.abspath(__file__)) - output_dir = os.path(script_dir, "dataset", "chinese-c4") - os.makedirs(output_dir, exist_ok=True) - - for i in range(num_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 = f"{output_dir}/chinese-c4_text_chunk_{i}.jsonl" - with open(output_path, "w", encoding="utf-8") as f: - for example in chunk: - # 每行写入一个 {"text": "xxx"} 对象 - json_line = {"text": example["text"]} - f.write(json.dumps(json_line, ensure_ascii=False) + "\n") - - print(f"Saved text chunk {i} to {output_path}") + process_dataset( + dataset_name="shjwudp/chinese-c4", + output_subdir="chinese-c4" + ) \ No newline at end of file diff --git a/chinese-wiki.py b/chinese-wiki.py index 815a811..e0afb99 100644 --- a/chinese-wiki.py +++ b/chinese-wiki.py @@ -1,28 +1,7 @@ -from datasets import load_dataset -import json -import os - +from utils import process_dataset if __name__ == "__main__": - dataset_dict = load_dataset("Blaze7451/Wiki-zh-20250601") - train_dataset = dataset_dict["train"] - chunk_size = 1000000 - total_samples = len(train_dataset) - num_chunks = (total_samples // chunk_size) + 1 - script_dir = os.path.dirname(os.path.abspath(__file__)) - output_dir = os.path.join(script_dir, "dataset", "chinese-wiki") - os.makedirs(output_dir, exist_ok=True) - - - for i in range(num_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 = f"{output_dir}/chinese-wiki_text_chunk_{i}.jsonl" - with open(output_path, "w", encoding="utf-8") as f: - for example in chunk: - json_line = {"text": example["text"]} - f.write(json.dumps(json_line, ensure_ascii=False) + "\n") - - print(f"Saved text chunk {i} to {output_path}") \ No newline at end of file + process_dataset( + dataset_name="Blaze7451/Wiki-zh-20250601", + output_subdir="chinese-wiki" + ) \ No newline at end of file diff --git a/english-fineweb.py b/english-fineweb.py index 871d606..f429e6b 100644 --- a/english-fineweb.py +++ b/english-fineweb.py @@ -1,6 +1,4 @@ -from datasets import load_dataset -import json -import os +from utils import process_dataset import re def comprehensive_normalization(text): @@ -14,29 +12,11 @@ def comprehensive_normalization(text): pattern = re.compile('|'.join(re.escape(k) for k in replacements)) return pattern.sub(lambda m: replacements[m.group()], text) + if __name__ == "__main__": - dataset_dict = load_dataset("HuggingFaceFW/fineweb","sample-10BT") - train_dataset = dataset_dict["train"] - - chunk_size = 1000000 - total_samples = len(train_dataset) - num_chunks = (total_samples // chunk_size) + 1 - - script_dir = os.path.dirname(os.path.abspath(__file__)) - output_dir = os.path(script_dir, "dataset", "english-fineweb") - os.makedirs(output_dir, exist_ok=True) - - for i in range(num_chunks): - if i == 10: - break - 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 = f"{output_dir}/english-fineweb_text_chunk_{i}.jsonl" - with open(output_path, "w", encoding="utf-8") as f: - for example in chunk: - json_line = {"text": comprehensive_normalization(example["text"])} - f.write(json.dumps(json_line, ensure_ascii=False) + "\n") - - print(f"Saved text chunk {i} to {output_path}") \ No newline at end of file + process_dataset( + dataset_name="HuggingFaceFW/fineweb", + output_subdir="english-fineweb", + dataset_config="sample-10BT", + normalization_func=comprehensive_normalization + ) \ No newline at end of file diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..e3b3ddd --- /dev/null +++ b/utils.py @@ -0,0 +1,39 @@ +from datasets import load_dataset +import json +import os + +def process_dataset( + dataset_name: str, + output_subdir: str, + dataset_config: str = None, + split_name: str = "train", + chunk_size: int = 1000000, + normalization_func=None +): + + dataset_dict = load_dataset(dataset_name, dataset_config) if dataset_config else load_dataset(dataset_name) + train_dataset = dataset_dict[split_name] + + total_samples = len(train_dataset) + num_chunks = (total_samples // chunk_size) + 1 + + 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): + 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: + text = example["text"] + if normalization_func: + text = normalization_func(text) + json_line = {"text": text} + f.write(json.dumps(json_line, ensure_ascii=False) + "\n") + + print(f"Saved text chunk {i} to {output_path}") +