feat(dataset): 更新数据集处理和保存逻辑
This commit is contained in:
+1
-1
@@ -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
@@ -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)
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user