Add granite prompt format to example utils

This commit is contained in:
turboderp
2024-05-30 20:09:04 +02:00
parent 1fbf3643f1
commit 475c5b5e89

View File

@@ -11,10 +11,19 @@ def format_prompt(prompt_format, sp, p):
f"{p}<|eot_id|>"
f"<|start_header_id|>assistant<|end_header_id|>\n\n"
)
elif prompt_format == "granite":
return (
f"System:\n"
f"{sp}\n\n"
f"Question:\n"
f"{p}\n\n"
f"Answer:\n"
)
def get_stop_conditions(prompt_format, tokenizer):
if prompt_format == "llama":
return [tokenizer.eos_token_id]
elif prompt_format == "llama3":
return [tokenizer.single_id("<|eot_id|>")]
elif prompt_format == "granite":
return [tokenizer.eos_token_id, "\n\nQuestion:"]