Update tokenizer_utils.py

This commit is contained in:
Daniel Han 2024-11-11 00:04:02 -08:00
commit 10d4187522

View file

@ -588,15 +588,21 @@ pass
def _fix_chat_template(chat_template):
endfor = "{% endfor %}"
where = chat_template.find(endfor)
if where == -1: return chat_template
if where == -1:
endfor = "{%- endfor %}"
where = chat_template.find(endfor)
if where == -1:
return chat_template
after_endfor = chat_template[where + len(endfor):]
if "{% if" not in after_endfor and "{% set " not in after_endfor and \
dash = "-" if endfor.startswith("{%-") else ""
if "{%" + dash + " if" not in after_endfor and "{%" + dash + " set " not in after_endfor and \
after_endfor.startswith("{{") and after_endfor.endswith("}}") and \
after_endfor.count("{{") == 1 and after_endfor.count("}}") == 1:
after_endfor = "{% if add_generation_prompt %}" + after_endfor + "{% endif %}"
after_endfor = "{%" + dash + " if add_generation_prompt %}" + after_endfor + endfor
chat_template = chat_template[:where + len(endfor)] + after_endfor
pass
@ -643,10 +649,12 @@ def fix_chat_template(tokenizer):
if no == yes:
# SAME?! That's not good! We check for add_generation_prompt
if "{% if add_generation_prompt %}" not in chat_template:
if "{% if add_generation_prompt %}" not in chat_template and \
"{%- if add_generation_prompt %}" not in chat_template:
# Try fixing it by adding it
new_chat_template = _fix_chat_template(chat_template)
if "{% if add_generation_prompt %}" not in new_chat_template:
if "{% if add_generation_prompt %}" not in new_chat_template and \
"{%- if add_generation_prompt %}" not in new_chat_template:
raise RuntimeError(
f"Unsloth: The tokenizer `{tokenizer.name_or_path}`\n"\
"does not have a {% if add_generation_prompt %} for generation purposes.\n"\