fix-apply-chat-template

chat_template.jinja

@@ -1,3 +1,19 @@
+{% macro render_content(msg) -%}
+    {%- set c = msg.get('content') -%}
+    {%- if c is string -%}
+      {{ c }}
+    {%- elif c is not none -%}
+      {% for content in c -%}
+        {% if content['type'] == 'image' or 'image' in content or 'image_url' in content -%}
+          <|media_start|>image<|media_content|><|media_pad|><|media_end|>
+        {% else -%}
+          {{ content['text'] }}
+        {%- endif -%}
+      {%- endfor -%}
+    {%- endif -%}
+{%- endmacro %}
+
+
 {%- if tools -%}
   <|im_system|>tool_declare<|im_middle|>{{ tools | tojson(separators=(',', ':')) }}<|im_end|>
 {%- endif -%}
@@ -16,26 +32,18 @@
   {%- endif -%}
 
   {%- if message['role'] == 'assistant' and message.get('tool_calls') -%}
-    {%- if message['content'] -%}{{ message['content'] }}{%- endif -%}
-    <|tool_calls_section_begin|>
+    {{render_content(message)}}<|tool_calls_section_begin|>
     {%- for tool_call in message['tool_calls'] -%}
       {%- set formatted_id = tool_call['id'] -%}
       <|tool_call_begin|>{{ formatted_id }}<|tool_call_argument_begin|>{% if tool_call['function']['arguments'] is string %}{{ tool_call['function']['arguments'] }}{% else %}{{ tool_call['function']['arguments'] | tojson }}{% endif %}<|tool_call_end|>
     {%- endfor -%}
     <|tool_calls_section_end|>
   {%- elif message['role'] == 'tool' -%}
-    ## Return of {{ message.tool_call_id }}
-{{ message['content'] }}
-  {%- elif message['content'] is string -%}
-    {{ message['content'] }}
+    {%- set tool_call_id = message.tool_call_id if message.tool_call_id.startswith('functions.') else 'functions.' + message.tool_call_id -%}
+    ## Return of {{ tool_call_id }}
+{{render_content(message)}}
   {%- elif message['content'] is not none -%}
-    {% for content in message['content'] -%}
-      {% if content['type'] == 'image' or 'image' in content or 'image_url' in content -%}
-        <|media_start|>image<|media_content|><|media_pad|><|media_end|>
-      {% else -%}
-        {{ content['text'] }}
-      {%- endif -%}
-    {%- endfor -%}
+    {{render_content(message)}}
   {%- endif -%}
   <|im_end|>
 {%- endfor -%}

tokenization_kimi.py

@@ -326,10 +326,17 @@ class TikTokenTokenizer(PreTrainedTokenizer):
 
     
     def apply_chat_template(
-        self, conversation, tools: Optional[list[dict]] = None, **kwargs
+        self, conversation, tools: Optional[list[dict]] = None, 
+        tokenize: bool = False, 
+        add_generation_prompt: bool = True, 
+        **kwargs
     ):
         tools = deep_sort_dict(tools)
-        return super().apply_chat_template(conversation, tools=tools, **kwargs)
+        return super().apply_chat_template(conversation, 
+                                           tools=tools,
+                                           tokenize=tokenize,
+                                           add_generation_prompt=add_generation_prompt, 
+                                           **kwargs)
     
 
 def deep_sort_dict(obj: Any) -> Any: