ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Change _fix_chat_template in case a template has both endif and endfor (#1388)

Browse source
This commit is contained in:
Giulia Baldini 2024-12-20 11:23:30 +01:00 • committed by GitHub
commit f3e6a28c3f
1 changed files with 27 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

37
unsloth/tokenizer_utils.py
View file

@ -585,26 +585,43 @@ def load_correct_tokenizer(
pass
def _find_end_position(template, endfor, endif):
where_endfor = template.find(endfor)
where_endif = template.find(endif)
if where_endfor == where_endif == -1:
return None
elif where_endfor > where_endif:
return endfor
else:
return endif
pass
pass
def _fix_chat_template(chat_template):
endfor = "{% endif %}"
where = chat_template.find(endfor)
if where == -1:
endfor = "{%- endif %}"
where = chat_template.find(endfor)
if where == -1:
endfor = "{% endfor %}"
endif = "{% endif %}"
chosen_end = _find_end_position(chat_template, endfor, endif)
if chosen_end is None:
endfor = "{%- endfor %}"
endif = "{%- endif %}"
chosen_end = _find_end_position(chat_template, endfor, endif)
if chosen_end is None:
return chat_template
where = chat_template.find(chosen_end)
after_endfor = chat_template[where + len(endfor):]
after_endfor = chat_template[where + len(chosen_end):]
dash = "-" if endfor.startswith("{%-") else ""
dash = "-" if chosen_end.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 = "{%" + dash + " if add_generation_prompt %}" + after_endfor + endfor
after_endfor = "{%" + dash + " if add_generation_prompt %}" + after_endfor + endif
chat_template = chat_template[:where + len(endfor)] + after_endfor
chat_template = chat_template[:where + len(chosen_end)] + after_endfor
pass
return chat_template
pass