fix: more chatml-function-calling fixes

abetlen · abetlen · commit 345215a76cf5 · 2024-02-13T23:02:50.000-05:00
diff --git a/llama_cpp/llama_chat_format.py b/llama_cpp/llama_chat_format.py
@@ -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' %}"
@@ -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(),
@@ -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
