Fix issues with system_message rendering
When trying to pass a system_message in chat history, the template rendering fails with 'TypeError: can only concatenate str (not "list") to str'.
The underlying issue being the format of passed messages with "type" and "content" inside the contents of the message.
This case is handled correctly in user/assistant messages, inside the main loop, so I've adapted the template code from that part, assuming the system_message can be just a single text message.
Hey,
Thanks for the suggestion!
But I think this will break the processor that expects type.
I'm using this template right now, and haven't spotted the issue yet.
To further elaborate on the issue. Previous code will raise the TypeError: can only concatenate str (not "int") inside jinja2 template rendering function, when running the following code:
model, processor = load(name, processor_config={"trust_remote_code": True})
config = model.config
prompt = apply_chat_template(self.processor, self.config, [{"role": "system", "message": "some system message"}, {"role": "user", "content": "some prompt"}])
This is due to processor mapping these prompts into {"role": "system", "message": {"type": "text", "content": "some system message"}, {"role": "user", "content": {"type": "text", "content": "some prompt"}} before passing it to jinja.
So I've fixed this case, while leaving the room for old style, in case something changes inside the processor.
As the code is linearized, here's before and after versions of the part I've changed.
Before:
{%- if messages[0]["role"] == "system" %}
{%- set system_message = messages[0]["content"] %}
{%- set loop_messages = messages[1:] %}
{%- else %}
{%- set loop_messages = messages %}
{%- endif %}
After:
{%- if messages[0]["role"] == "system" %}
{%- if messages[0]["content"] is not string %}
{%- if messages[0]["content"]|length != 1 %}
{{- raise_exception("Multiple messages in system prompt are not supported!") }}
{%- endif %}
{%- if messages[0]["content"][0]["type"] == "text" %}
{%- set system_message = messages[0]["content"][0]["content"] %}
{%- else %}
{{- raise_exception("Unrecognized content type for system message!") }}
{%- endif %}
{%- else %}
{%- set system_message = messages[0]["content"] %}
{%- endif %}
{%- set loop_messages = messages[1:] %}
{%- else %}
{%- set loop_messages = messages %}
{%- endif %}
I think there are issues with the rest of chat template too, with system_message defined it puts some of the responses before system message?
For the following messages:
[
{'role': 'system', 'content': 'SYSTEM_MESSAGE'},
{'role': 'user', 'content': 'USER_REQUEST_1'},
{'role': 'assistant', 'content': 'ASSISTANT_RESPONSE'},
{'role': 'user', 'content': 'USER_REQUEST_2'}
]
the prompt will be generated by apply_chat_template function as follows:
<s>[INST][IMG]USER_REQUEST_1[/INST]ASSISTANT_RESPONSE</s>[INST]SYSTEM_MESSAGE
USER_REQUEST_2[/INST]
I'm still wrapping my head around it, but it definitely looks wrong:)
Here's the final version of the chat template that I've ended up using:
{%- if messages[0]["role"] == "system" %}
{%- if messages[0]["content"] is not string %}
{%- if messages[0]["content"]|length != 1 %}
{{- raise_exception("Multiple messages in system prompt are not supported!") }}
{%- endif %}
{%- if messages[0]["content"][0]["type"] == "text" %}
{%- set system_message = messages[0]["content"][0]["content"] %}
{%- else %}
{{- raise_exception("Unrecognized content type for system message!") }}
{%- endif %}
{%- else %}
{%- set system_message = messages[0]["content"] %}
{%- endif %}
{%- set loop_messages = messages[1:] %}
{%- else %}
{%- set loop_messages = messages %}
{%- endif %}
{%- for message in loop_messages %}
{%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
{{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }}
{%- endif %}
{%- if message["role"] == "user" %}
{%- if loop.first and system_message is defined %}
{{- bos_token }}
{{- "[INST]" + system_message + "[/INST][INST]" }}
{%- else %}
{{- bos_token }}{{- "[INST]" }}
{%- endif %}
{%- if message["content"] is not string %}
{%- for chunk in message["content"] %}
{%- if chunk["type"] == "text" %}
{{- chunk["content"] }}
{%- elif chunk["type"] == "image" %}
{{- "[IMG]" }}
{%- else %}
{{- raise_exception("Unrecognized content type!") }}
{%- endif %}
{%- endfor %}
{%- else %}
{{- message["content"] }}
{%- endif %}
{{- "[/INST]" }}
{%- elif message["role"] == "assistant" %}
{{- message["content"] + eos_token}}
{%- else %}
{{- raise_exception("Only user and assistant roles are supported, with the exception of an initial optional system message!") }}
{%- endif %}
{%- endfor %}