fix add_generation_prompt in latest version

This commit is contained in:
Dejian Yang
2024-01-09 21:55:21 +08:00
parent d3bb741e25
commit c160d96860
2 changed files with 2 additions and 2 deletions

View File

@@ -66,7 +66,7 @@ def generate_one(example, tokenizer, model):
prompt = example['prompt']
inputs = tokenizer.apply_chat_template(
[{'role': 'user', 'content': prompt }],
return_tensors="pt"
return_tensors="pt", add_generation_prompt=True
).to(model.device)
stop_id = tokenizer.convert_tokens_to_ids("<|EOT|>")