Skip to content

Commit

Permalink
fix: more chatml-function-calling fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
abetlen committed Feb 14, 2024
1 parent b1637c2 commit 345215a
Showing 1 changed file with 7 additions and 5 deletions.
12 changes: 7 additions & 5 deletions llama_cpp/llama_chat_format.py
Original file line number Diff line number Diff line change
Expand Up @@ -2061,12 +2061,12 @@ def chatml_function_calling(
"\nfunctions.<function_name>:"
'\n{ "arg1": "value1", "arg2": "value2" }'
"{% endif %}"
"\n<|im_end|>\n"
"<|im_end|>\n"
"{% endif %}"
# User message
"{% if message.role == 'user' %}"
"{{ message.content }}"
"\n<|im_end|>\n"
"<|im_end|>\n"
"{% endif %}"
# Assistant message
"{% if message.role == 'assistant' %}"
Expand All @@ -2076,19 +2076,19 @@ def chatml_function_calling(
"message:\n"
"{% endif %}"
"{{ message.content }}"
"\n<|im_end|>\n"
"<|im_end|>\n"
"{% endif %}"
## Function calls
"{% if 'tool_calls' in message %}"
"{% for tool_call in message.tool_calls %}"
"functions.{{ tool_call.function.name }}:\n"
"{{ tool_call.function.arguments }}"
"{% endfor %}"
"\n<|im_end|>\n"
"<|im_end|>\n"
"{% endif %}"
"{% endif %}"
"{% endfor %}"
"{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
"{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
)
template_renderer = jinja2.Environment(
loader=jinja2.BaseLoader(),
Expand Down Expand Up @@ -2120,6 +2120,8 @@ def chatml_function_calling(
},
}

stop = [stop, "<|im_end|>"] if isinstance(stop, str) else stop + ["<|im_end|>"] if stop else ["<|im_end|>"]

# Case 1: No tool choice by user
if (
tool_choice is None
Expand Down

0 comments on commit 345215a

Please sign in to comment.