feat(dataset): 更新数据集处理和保存逻辑

This commit is contained in:
2025-07-20 11:21:18 +08:00
parent 365086f102
commit fcb5c2bcc8
3 changed files with 17 additions and 2 deletions
+1 -1
View File
@@ -32,6 +32,6 @@ if __name__ == "__main__":
dataset = load_dataset("BelleGroup/train_3.5M_CN")
process_dataset(
dataset_dict=dataset,
output_subdir="belle_sft",
output_subdir="belle-sft",
process_func=process_func,
)
+15
View File
@@ -0,0 +1,15 @@
from utils import dump_pkl_files, fetch_files
from tokenizer import BpeTokenizer
import os
if __name__ == "__main__":
tokenizer = BpeTokenizer("tokenizer.json")
base_dir = [
os.path.join("dataset", "belle-sft"),
]
base_out_dir = "pkl_output"
files = []
for dir_path in base_dir:
files.extend(fetch_files(dir_path))
dump_pkl_files(tokenizer, files, base_out_dir, packing_size=2048)
+1 -1
View File
@@ -38,7 +38,7 @@ def dump_pkl_files(
):
def process_line(line: str) -> Tensor:
line = json.loads(line)[key]
processed_line = encoder(line)
processed_line = encoder(line) if encoder else line
ids = tokenizer.encode(processed_line)
arrow = torch.tensor(ids, dtype=torch.int32)
return arrow