[discussion] apply chat template per message?

[arendu]

Hi Team, I see here and here that the chat template is applied individually to each message. This is different from HF usage where the entire conversation list gets chat templated.

These two ways might not be compatible depending on the tokenizer.

[SahilJain314]

This is intentional right now for multi-turn RL reasons.

We're worried about token merging across message boundaries. I.e

User:
What's the day today?
Assistant:
I'll use the internet to find out <tool: internet>Day today</toll: internet>

Then the tool is called, so the string becomes
User:
What's the day today?
Assistant:
I'll use the internet to find out <tool: internet>Day today</toll: internet>

However, if tokenized as a whole, we're worried that the and the </tool:internet> and/or the strings could get merged, which would break token alignment and logprobs would be incorrect

Instead of
in: [x x x], out: [y y y] being [x x x y y y], you could end up with [x x z y y].

If <tool: internet> and are delimited by special tokens (guaranteed not merged by tokenizer), then we wouldn't have this problem, but if we want to support arbitrary end strings (#57), we can't ensure that. (plus, I figure that'll be a common use case, since adding a special token would require some fine-tuning to get off the ground).

We don't really face this problem with single-turn things (we just tokenize the input and concatenate the output), but in multiturn cases, it becomes a question as to if you always concatenate (doing what we do), or detokenize and re-tokenize on every step, potentially running into the problems above.

@YUki-666 who had a similar question.

I'd love to hear your thoughts on this though. Do you have an alternative solution/believe this is a nonissue? I agree that this way of doing this looks nasty and breaks the nice property of $tokens == tokenizer(tokenizer^{-1}(tokens))$

[arendu]

Hmm, what would happen if we train our model with the current setup and then export it to HF and users in HF do the following for multi-turn prompts: tokenizer.apply_chat_template([message_list], add_generation_prompt=False)

Haven't we created a train/test breakdown?

Right now, I think aligner/nemo adds some asserts (at least I think so, been a while since I looked at this) to catch token merging across turns. This might be annoying because it essentially means we have to pick a chat template that will be merged correctly, but in the long run it seems safer.

Also, I think the token-merging is a problem only when we are working with new models from scratch(?). If we are working with instruct models with pre-defined tokenizers and chat_templates, the HF way of doing things should already work OOTB.

[yuki-97]

Since it's common to use tokenizer.apply_chat_template at the entire conversation, I think we should also do this at training to keep it consistent with generation.
Otherwise, if we apply template per message in training, we can't control how user use it in inference, and commonly they will tokenizer.apply_chat_template at the entire conversation. It's mismatched with training, so they will get bad result.
BTW, tokenizer.apply_chat_template per message will have multiple system prompts finally, which is definitely incorrect.

As for the boundary error, I'm a little bit confused now. (I can't understand well b/c maybe your <.bos> and <.eos> is hidden?)
As long as the chat_template is correct, there are <.bos> and <.eos> to split different message chunks, why will it be merged across message boundaries?

Assume this is first turn:
[{role=system, content=system_prompt}, {role=user, content=What's the day today?}]
After apply_chat_template:
<.bos>System
system_prompt<.eos>
<.bos>User
What's the day today?<.eos>
<.bos>Assistant

Second turn:
[{role=system, content=system_prompt}, {role=user, content=What's the day today?}, {role=assistant, content=I'll use the internet to find out <tool: internet>Day today</toll: internet>}, {role=user, content=xxx}]
After apply_chat_template:
<.bos>System
system_prompt<.eos>
<.bos>User
What's the day today?<.eos>
<.bos>Assistant
I'll use the internet to find out <tool: internet>Day today</toll: internet><.eos>
<.bos>User
xxx<.eos>
<.bos>Assistant

[arendu]

I can see why message by message is useful (it makes masking easy). Can we however, throw an assertion that we should have apply_chat_template(list_of_message) == [apply_chat_template(msg) for msg in list_of_message]
Would some kind of check like this make sense?

On the other hand, it is best to keep implementation simple/consistent and restrict the types of tokens that signify the boundaries. For example -- why can we enforce that only single-token end-of-sequence is supported? or only single-token message-turn-end is supported?

[yuki-97]

apply_chat_template(list_of_message) == [apply_chat_template(msg) for msg in list_of_message] is impossible at least on Qwen's own chat_template.
You will always get a <|im_start|>system when calling apply_chat_template no matter whether you have the role "system".

Maybe we can use custom chat_template to avoid it, but how we ensure that user will use our chat_template and apply it per message when inference?
And if we force user to use a custom chat_template, it's very inconvenient; else if we support ourselves, it's also inconvenient since there are so many models.

Still, I'd like to know why here is a boundary error, can someone give me an example?
If we have a message list like [{role=xx, content=xx}, {role=xx, content=xx}, ...], then apply_chat_template to the entire conversation, we will have <.bos> and <.eos> to separate each sequence well. (detail can be seen at my previous comment)
Even if use arbitrary end strings, when we get new sequence end by the end string, we can add {role=assistant, content=new_sequence} to previous message list, then use apply_chat_template to the entire conversation, we will still get <.bos> and <.eos> for the new sequence part.

For example Qwen/Qwen2.5-Math-1.5B-Instruct's template:

"{%- if tools %}\n    {{- '<\|im_start\|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'Please reason step by step, and put your final answer within \\\\boxed{}.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool \| tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><\|im_end\|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<\|im_start\|>system\\n' + messages[0]['content'] + '<\|im_end\|>\\n' }}\n    {%- else %}\n        {{- '<\|im_start\|>system\\nPlease reason step by step, and put your final answer within \\\\boxed{}.<\|im_end\|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<\|im_start\|>' + message.role + '\\n' + message.content + '<\|im_end\|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<\|im_start\|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments \| tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<\|im_end\|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<\|im_start\|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<\|im_end\|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<\|im_start\|>assistant\\n' }}\n{%- endif %}\n"

Convert it to a python script. (by AI)

def format_chat_messages(messages, tools=None, add_generation_prompt=False):
    """Format chat messages with tools and special formatting.
    
    Args:
        messages: List of message dictionaries with 'role' and 'content' keys
        tools: Optional list of tool definitions
        add_generation_prompt: Whether to add an assistant prompt at the end
        
    Returns:
        Formatted string with all messages and tools
    """
    result = []
    
    # Handle system message and tools
    if tools:
        # Add system message
        if messages and messages[0]['role'] == 'system':
            result.append('<|im_start|>system\n' + messages[0]['content'])
        else:
            result.append('<|im_start|>system\nPlease reason step by step, and put your final answer within \\boxed{}.')
            
        # Add tools section
        result.append('\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\n'
                     'You are provided with function signatures within <tools></tools> XML tags:\n<tools>')
        
        # Add each tool
        for tool in tools:
            result.append('\n' + str(tool))
            
        result.append('\n</tools>\n\n'
                     'For each function call, return a json object with function name and arguments within '
                     '<tool_call></tool_call> XML tags:\n<tool_call>\n'
                     '{"name": <function-name>, "arguments": <args-json-object>}\n'
                     '</tool_call><|im_end|>\n')
    else:
        # Handle system message without tools
        if messages and messages[0]['role'] == 'system':
            result.append('<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n')
        else:
            result.append('<|im_start|>system\nPlease reason step by step, and put your final answer within \\boxed{}.<|im_end|>\n')
    
    # Process all messages
    for i, message in enumerate(messages):
        role = message['role']
        
        if role in ['user', 'system'] and (role != 'system' or i > 0):
            # Handle user and non-first system messages
            result.append(f'<|im_start|>{role}\n{message["content"]}<|im_end|>\n')
            
        elif role == 'assistant':
            # Handle assistant messages
            result.append(f'<|im_start|>{role}')
            
            if message.get('content'):
                result.append('\n' + message['content'])
                
            # Handle tool calls
            if 'tool_calls' in message:
                for tool_call in message['tool_calls']:
                    if 'function' in tool_call:
                        tool_call = tool_call['function']
                    result.append(f'\n<tool_call>\n{{"name": "{tool_call["name"]}", "arguments": {tool_call["arguments"]}}}\n</tool_call>')
                    
            result.append('<|im_end|>\n')
            
        elif role == 'tool':
            # Handle tool responses
            if i == 0 or messages[i-1]['role'] != 'tool':
                result.append('<|im_start|>user')
            result.append(f'\n<tool_response>\n{message["content"]}\n</tool_response>')
            if i == len(messages)-1 or messages[i+1]['role'] != 'tool':
                result.append('<|im_end|>\n')
    
    # Add generation prompt if requested
    if add_generation_prompt:
        result.append('<|im_start|>assistant\n')
        
    return ''.join(result)

[arendu]

apply_chat_template(list_of_message) == [apply_chat_template(msg) for msg in list_of_message] is impossible at least on Qwen's own chat_template.

If this is the case, then I think we should not do [apply_chat_template(msg) for msg in list_of_message] and find another way to resolve the "boundary issue".