Add support for images and fix end of assistant issue (#22)

- add support for images and fix an issue with the end of assistant tokens that was not added if the last message was an assistant message. (f89f5ef02e1b9ecd18dd8b982016312f09c08663)


Co-authored-by: nathan ranchin <nathanrchn@users.noreply.huggingface.co>

chat_template.jinja

@@ -150,8 +150,9 @@
 {%- set outer_token = '<|inner_suffix|>' -%}
 {%- set tool_calls_token = '<|tools_prefix|>' -%}
 {%- set end_tool_calls_token = '<|tools_suffix|>' -%}
+{%- set image_token = '<|image|>' -%}
 
-{%- set ns = namespace(in_assistant=false, in_tool=false, in_inner=false, assistant_format=none) -%}
+{%- set ns = namespace(in_assistant=false, in_tool=false, in_inner=false, waiting_for_tool_outputs=false, assistant_format=none) -%}
 
 {%- if messages and messages[0].role == 'system' -%}
     {%- if "content" in messages[0] -%}
@@ -204,6 +205,8 @@
                 {%- for part in parts -%}
                     {%- if part.type == "text" -%}
                         {{ part.text }}
+                    {%- elif part.type == "image" -%}
+                        {{ image_token }}
                     {%- else -%}
                         {{- raise_exception("Invalid user part: " + part.type) -}}
                     {%- endif -%}
@@ -257,6 +260,7 @@
                             {%- endif -%}
                         {%- endfor -%}
                         {{ ']' + end_tool_calls_token }}
+                        {%- set ns.waiting_for_tool_outputs = true -%}
                     {%- elif block.type == 'tool_outputs' -%}
                         {%- if ns.in_tool -%}
                             {{- raise_exception("Cannot have both tool outputs as separate messages and tool outputs as blocks") -}}
@@ -269,6 +273,7 @@
                             {%- endif -%}
                         {%- endfor -%}
                         {{- ']' }}
+                        {%- set ns.waiting_for_tool_outputs = false -%}
                     {%- elif block.type == 'response' -%}
                         {%- if ns.in_tool -%}
                             {{ ']' }}
@@ -303,6 +308,7 @@
                 {%- endif -%}
             {%- endfor -%}
             {{ ']' + end_tool_calls_token }}
+            {%- set ns.waiting_for_tool_outputs = true -%}
         {%- endif -%}
     {%- elif message.role == 'tool' -%}
         {%- if not ns.in_assistant -%}
@@ -315,6 +321,7 @@
             {{ ", "}}
         {%- endif -%}
         {{ message.content }}
+        {%- set ns.waiting_for_tool_outputs = false -%}
     {%- else -%}
         {{- raise_exception("Invalid message role") -}}
     {%- endif -%}
@@ -322,6 +329,9 @@
 {%- if ns.in_tool -%}
     {{ ']' }}
 {%- endif -%}
+{%- if ns.in_assistant and not (continue_assistant_message is defined and continue_assistant_message) and not ns.waiting_for_tool_outputs -%}
+    {{ end_assistant_token }}
+{%- endif -%}
 {%- if add_generation_prompt -%}
     {{ assistant_token }}
 {%- endif -%}

tokenizer.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bb201fb226cde11f66c3cf51c5344fb37b1611f00c21e75c324546d854eff2e1
-size 17078480
+oid sha256:010095abf1dc6f52d4867584e7b3f0d4eece854593ae310220ec7782dd1b0a66
+size 17078474

tokenizer_config.json

@@ -588,7 +588,7 @@
       "special": true
     },
     "73": {
-      "content": "<SPECIAL_73>",
+      "content": "<|image|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,