ViperEkura/DataPipeline
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(utils): 增加数据集处理的灵活性和可配置性

Browse Source
This commit is contained in:
ViperEkura
2025-07-14 13:34:03 +08:00
parent 2b0f805ff3
commit 820ddf3fc8
1 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils.py
+6 -3
View File
@@ -19,7 +19,9 @@ def process_dataset(
dataset_name: str, dataset_name: str,
output_subdir: str, output_subdir: str,
dataset_config: str = None, dataset_config: str = None,
max_chunk_size: int = None,
split_name: str = "train", split_name: str = "train",
column_name: str = "text",
chunk_size: int = 1000000, chunk_size: int = 1000000,
normalization_func=comprehensive_normalization normalization_func=comprehensive_normalization
): ):
@@ -29,12 +31,13 @@ def process_dataset(
total_samples = len(train_dataset) total_samples = len(train_dataset)
num_chunks = (total_samples // chunk_size) + 1 num_chunks = (total_samples // chunk_size) + 1
lim_chunks = max_chunk_size if max_chunk_size else num_chunks
script_dir = os.path.dirname(os.path.abspath(__file__)) script_dir = os.path.dirname(os.path.abspath(__file__))
output_dir = os.path.join(script_dir, "dataset", output_subdir) output_dir = os.path.join(script_dir, "dataset", output_subdir)
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
for i in range(num_chunks): for i in range(lim_chunks):
start_idx = i * chunk_size start_idx = i * chunk_size
end_idx = min((i + 1) * chunk_size, total_samples) end_idx = min((i + 1) * chunk_size, total_samples)
chunk = train_dataset.select(range(start_idx, end_idx)) chunk = train_dataset.select(range(start_idx, end_idx))
@@ -42,10 +45,10 @@ def process_dataset(
output_path = os.path.join(output_dir, f"{output_subdir}_text_chunk_{i}.jsonl") output_path = os.path.join(output_dir, f"{output_subdir}_text_chunk_{i}.jsonl")
with open(output_path, "w", encoding="utf-8") as f: with open(output_path, "w", encoding="utf-8") as f:
for example in chunk: for example in chunk:
text = example["text"] text = example[column_name]
if normalization_func: if normalization_func:
text = normalization_func(text) text = normalization_func(text)
json_line = {"text": text} json_line = {column_name : text}
f.write(json.dumps(json_line, ensure_ascii=False) + "\n") f.write(json.dumps(json_line, ensure_ascii=False) + "\n")
print(f"Saved text chunk {i} to {output_path}") print(f"Saved text chunk {i} to {output_path}")