refactor(dataset): 重构数据集处理逻辑

This commit is contained in:
2025-07-14 08:45:01 +08:00
parent d4df706bc5
commit 7591a231df
4 changed files with 57 additions and 80 deletions
+8 -28
View File
@@ -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}")
process_dataset(
dataset_name="HuggingFaceFW/fineweb",
output_subdir="english-fineweb",
dataset_config="sample-10BT",
normalization_func=comprehensive_normalization
)