feat(dump_sft_file): 更新数据集处理逻辑

This commit is contained in:
2025-08-01 16:59:47 +08:00
parent 835b12c36a
commit 8725a9ebd1
2 changed files with 7 additions and 16 deletions
+2 -2
View File
@@ -14,7 +14,7 @@ def get_processor(tokenizer: BpeTokenizer):
tokens = prefix_ids + suffix_ids tokens = prefix_ids + suffix_ids
tokens = torch.tensor(tokens, dtype=torch.int32) tokens = torch.tensor(tokens, dtype=torch.int32)
masks = torch.zeros_like(tokens, dtype=torch.bool) masks = torch.zeros_like(tokens, dtype=torch.bool)
masks[:len(prefix_ids)] = True masks[len(prefix_ids):] = True
return {"sequence": tokens, "mask": masks} return {"sequence": tokens, "mask": masks}
@@ -24,7 +24,7 @@ def get_processor(tokenizer: BpeTokenizer):
if __name__ == "__main__": if __name__ == "__main__":
tokenizer = BpeTokenizer("tokenizer.json") tokenizer = BpeTokenizer("tokenizer.json")
base_dir = [ base_dir = [
# os.path.join("dataset", "belle-sft"), os.path.join("dataset", "Ling-Coder-SFT"),
os.path.join("dataset", "chinese-instruct") os.path.join("dataset", "chinese-instruct")
] ]
base_out_dir = "pkl_output" base_out_dir = "pkl_output"
+5 -14
View File
@@ -2,15 +2,6 @@ from datasets import load_dataset
from utils import process_dataset from utils import process_dataset
def build_prompt(query, history) -> str:
ret_prompt = ""
if len(history) > 0:
for his_query, his_response in history:
ret_prompt += f"<|user|> {his_query} <|system|> <bos>{his_response}<eos>\n"
if query is not None:
ret_prompt += f"<|user|> {query} <|system|> <bos>"
return ret_prompt
def process_func(input_dict: dict): def process_func(input_dict: dict):
conversations = input_dict["conversations"] conversations = input_dict["conversations"]
@@ -20,12 +11,12 @@ def process_func(input_dict: dict):
for i in range(n): for i in range(n):
user_msg = conversations[2*i]["value"] user_msg = conversations[2*i]["value"]
assistant_msg = conversations[2*i+1]["value"] assistant_msg = conversations[2*i+1]["value"]
examples.append((user_msg, assistant_msg)) examples.append({
"query": user_msg,
"response": assistant_msg
})
content = { return examples
"text": build_prompt(None, examples)
}
return content
if __name__ == "__main__": if __name__ == "__main__":