fix(utils): 重命名变量并统一张量转换逻辑

This commit is contained in:
2025-09-29 14:13:46 +08:00
parent c3ebc4f681
commit e844680025
+5 -4
View File
@@ -121,10 +121,10 @@ def dump_pkl_files(
def get_pt_processor(tokenizer: BpeTokenizer): def get_pt_processor(tokenizer: BpeTokenizer):
def processor(intput_dict: dict) -> dict: def processor(intput_dict: dict) -> dict:
segment = intput_dict["text"] segment = intput_dict["text"]
ids = tokenizer.encode(f"{segment} <eos>") tokens = tokenizer.encode(f"{segment}<eos>")
t_ids = torch.tensor(ids, dtype=torch.int32) tokens = torch.tensor(tokens, dtype=torch.int32)
return {'sequence': t_ids} return {'sequence': tokens}
return processor return processor
@@ -133,6 +133,7 @@ def get_sft_processor(tokenizer: BpeTokenizer):
def processor(input_dict: dict): def processor(input_dict: dict):
query, response = input_dict["query"], input_dict["response"] query, response = input_dict["query"], input_dict["response"]
tokens = tokenizer.encode(f"<|user|> {query} <|system|> <bos>{response}<eos>\n") tokens = tokenizer.encode(f"<|user|> {query} <|system|> <bos>{response}<eos>\n")
tokens = torch.tensor(tokens, dtype=torch.int32)
return {"sequence": tokens} return {"sequence": tokens}
@@ -169,7 +170,7 @@ def cache_files(tokenizer, files, base_out_dir, cache_type):
keys = ["sequence"] keys = ["sequence"]
elif cache_type == "sft": elif cache_type == "sft":
processor = get_sft_processor(tokenizer) processor = get_sft_processor(tokenizer)
keys = ["sequence", "mask"] keys = ["sequence"]
elif cache_type == "dpo": elif cache_type == "dpo":
processor = get_dpo_processor(tokenizer) processor = get_dpo_processor(tokenizer)
keys = ["chosen", "chosen_mask", "rejected", "rejected_mask"] keys = ["chosen", "chosen_mask", "rejected", "rejected_mask"]