Update chat_templates.py

This commit is contained in:
Daniel Han-Chen 2024-02-29 00:09:00 +11:00
commit b1892b5511

View file

@ -221,6 +221,7 @@ CHAT_TEMPLATES["alpaca"] = (alpaca_template, alpaca_eos_token,)
# Notice we must use |trim for lstrip and rstrip. <start_of_turn> maps to 106.
# <end_of_turn> maps to 107. user and model are normal 1 word tokens.
gemma_template = \
"{{ bos_token }}"\
"{% for message in messages %}"\
"{% if message['role'] == 'user' %}"\
"{{'<start_of_turn>user\n' + message['content'] | trim + '<end_of_turn>\n'}}"\
@ -238,7 +239,7 @@ CHAT_TEMPLATES["gemma"] = (gemma_template, gemma_eos_token,)
# Gemma with ChatML instead
gemma_chatml_template = chatml_template
gemma_chatml_template = "{{ bos_token }}" + chatml_template
gemma_chatml_eos_token = (
{"<start_of_turn>" : "<|im_start|>", "<end_of_turn>" : "<|im_end|>"},
"<|im_end|>",