fixes to templates to get unit tests passing

Signed-off-by: Travis Johnson <tsjohnso@us.ibm.com>
vllm-project · Nov 22, 2024 · 7c6c0ba · 7c6c0ba
1 parent ecbb268
commit 7c6c0ba
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/examples/tool_chat_template_llama3.1_json.jinja b/examples/tool_chat_template_llama3.1_json.jinja
@@ -59,7 +59,7 @@
         {%- if messages[0]['content'] is string %}
             {%- set first_user_message = messages[0]['content']|trim %}
         {%- else %}
-            {%- set first_user_message = messages[0]['content']['text']|trim %}
+            {%- set first_user_message = messages[0]['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | map('trim') | join('\n') %}
         {%- endif %}
         {%- set messages = messages[1:] %}
     {%- else %}
@@ -106,7 +106,11 @@
         {%- if message.content is string %}
             {{- { "output": message.content } | tojson }}
         {%- else %}
-            {{- message.content | tojson }}
+            {%- for content in message['content']  %}
+                {%- if content['type']  == 'text' %}
+                    {{- { "output": content['text']  } | tojson }}
+                {%- endif %}
+            {%- endfor %}
         {%- endif %}
         {{- "<|eot_id|>" }}
     {%- endif %}

diff --git a/examples/tool_chat_template_llama3.2_json.jinja b/examples/tool_chat_template_llama3.2_json.jinja
@@ -78,7 +78,7 @@
         {%- if messages[0]['content'] is string %}
             {%- set first_user_message = messages[0]['content']|trim %}
         {%- else %}
-            {%- set first_user_message = messages[0]['content']['text']|trim %}
+            {%- set first_user_message = messages[0]['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | map('trim') | join('\n') %}
         {%- endif %}
         {%- set messages = messages[1:] %}
     {%- else %}
@@ -127,7 +127,11 @@
         {%- if message.content is string %}
             {{- { "output": message.content } | tojson }}
         {%- else %}
-            {{- message.content | tojson }}
+            {%- for content in message['content']  %}
+                {%- if content['type']  == 'text' %}
+                    {{- { "output": content['text']  } | tojson }}
+                {%- endif %}
+            {%- endfor %}
         {%- endif %}
         {{- "<|eot_id|>" }}
     {%- endif %}