mlx-community

/

pixtral-12b-4bit

like 7

Follow

MLX Community 2.94k

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 %}