[lewtun]

What does this PR do?

This PR addresses an issue with chat templates producing incorrect inputs when continue_final_message=True and we wish to terminate generation on newlines like \n (as is common in process supervision with math reward models like this one).

In this case, the prefills had their newlines stripped at each step of generation and produced results that disagree with pure greedy decoding to the EOS token. Here's an example:

import copy
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "Qwen/Qwen2.5-Math-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id).to("cuda")

### Inputs
system_prompt = "Please reason step by step, and put your final answer within \\boxed{}."
query = "Which number is larger, 9.8 or 9.11?"
messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": query}]

### Reference - greedy decoding to EOS
inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", return_dict=True, add_generation_prompt=True).to("cuda")
output = model.generate(**inputs, max_new_tokens=256, do_sample=False)
ref_output = tokenizer.decode(output[0][inputs.input_ids.shape[-1]:])
print(f"\n=== Ref completion ===\n{ref_output}")

### Iterative decoding with prefill
current_messages = copy.deepcopy(messages)

for i in range(5):
    if current_messages[-1]["role"] == "user":
        inputs = tokenizer.apply_chat_template(current_messages, return_tensors="pt", return_dict=True, add_generation_prompt=True).to("cuda")
    else:
        inputs = tokenizer.apply_chat_template(current_messages, return_tensors="pt", return_dict=True, continue_final_message=True).to("cuda")
    output = model.generate(**inputs, max_new_tokens=256, do_sample=False, stop_strings=["\n"], tokenizer=tokenizer)
    output_decoded = tokenizer.decode(output[0][inputs.input_ids.shape[-1]:])
    print(f"\n=== Completion {i+1} ===\n{repr(output_decoded)}")
    if current_messages[-1]["role"] == "user":
        current_messages.append({"role": "assistant", "content": output_decoded})
    else:
        current_messages[-1]["content"] += output_decoded

# we strip the final output to remove any trailing newlines produced by terminating on `\n`
final_output = current_messages[-1]["content"].rstrip()
print(f"\n=== Final output ===\n{final_output}")
assert ref_output == final_output

Here's the reference completion from greedy decoding until EOS (note the new lines):

To determine which number is larger between 9.8 and 9.11, we can compare them digit by digit starting from the leftmost digit.

1. Both numbers have the same digit in the units place, which is 9.
2. Next, we compare the digits in the tenths place. The digit in the tenths place of 9.8 is 8, and the digit in the tenths place of 9.11 is 1. Since 8 is greater than 1, we can conclude that 9.8 is greater than 9.11.

Therefore, the larger number is \boxed{9.8}.<|im_end|>

Without the fix we get bad generations because the prefill is continually stripping the trailing new lines per iteration:

=== Completion 1 ===
'To determine which number is larger between 9.8 and 9.11, we can compare them digit by digit starting from the leftmost digit.\n\n'

=== Completion 2 ===
" Here's the step-by-step process:\n\n"

=== Completion 3 ===
' \n\n'

=== Completion 4 ===
' \n\n'

=== Completion 5 ===
' \n\n'

=== Final output ===
To determine which number is larger between 9.8 and 9.11, we can compare them digit by digit starting from the leftmost digit.

 Here's the step-by-step process:

With the proposed fix, order is restored and the resulting completion is sane:

=== Completion 1 ===
'To determine which number is larger between 9.8 and 9.11, we can compare them digit by digit starting from the leftmost digit.\n\n'

=== Completion 2 ===
'1. Both numbers have the same digit in the units place, which is 9.\n'

=== Completion 3 ===
'2. Next, we compare the digits in the tenths place. The digit in the tenths place of 9.8 is 8, and the digit in the tenths place of 9.11 is 1. Since 8 is greater than 1, we can conclude that 9.8 is greater than 9.11.\n\n'

=== Completion 4 ===
'Therefore, the larger number is \\boxed{9.8}.<|im_end|>'

=== Completion 5 ===
'\n'

=== Final output ===
To determine which number is larger between 9.8 and 9.11, we can compare them digit by digit starting from the leftmost digit.

1. Both numbers have the same digit in the units place, which is 9.
2. Next, we compare the digits in the tenths place. The digit in the tenths place of 9.8 is 8, and the digit in the tenths place of 9.11 is 1. Since 8 is greater than 1, we can conclude that 9.8 is greater than 9.11.

Therefore, the larger number is \boxed{9.8}.<|im_end|>

Questions

• I am not sure why stripping was included in the original implementation. If it's important to retain, could we instead expose a strip_prefill arg in apply_chat_template that is True by default but can be disabled in cases like process supervision?
• I wasn't sure how to write a regression test for this behaviour since the tiny-gpt2 models aren't prone to emitting newlines and running a 1.5B Qwen model in the CI is presumably discouraged.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[lewtun]

CI seems to be failing for reasons unrelated to this PR

[Rocketknight1]

Wow, this issue was only introduced extremely recently, in #34214 ! I completely didn't foresee this being an issue when I merged that one.

The fundamental problem is that some templates strip() input text, and others don't. This means that the string-matching we do in continue_final_message isn't guaranteed to find a match if we don't also strip the search string. However, as you've noted, this can cause other problems.

Maybe the right solution is to first try to match the input with trailing whitespace, and then fallback to searching for the stripped version if that fails?

[lewtun]

Maybe the right solution is to first try to match the input with trailing whitespace, and then fallback to searching for the stripped version if that fails?

I like this approach - let me amend my PR. Do we also need to retain the rstrip() that's applied to rendered_chat?

[Rocketknight1]

No, I don't think so! I think it would break your task if we did, and I really just included it because I didn't realize the trailing whitespace could be meaningful.

[Rocketknight1]

cc @schoennenbeck, who wrote #34214 - you don't need to do anything here, just making you aware that we might change the behaviour again!

[lewtun]

@Rocketknight1 I've added a try/except in 0fd98e1

I've also tested that the Llama code still works

# From: https://github.com/huggingface/transformers/pull/34214
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4")
messages=[
    {"role": "user", "content": "What is the capital of France?"},
    {"role": "assistant", "content": "Great question! The capital of France is "}
]
s = tokenizer.apply_chat_template(
    messages, add_generation_prompt=False, 
    continue_final_message=True,
    tokenize=False
)
print(s)

Would you like me to include this as a regression test? I'm also not sure if we need something for the Qwen model, but for that we'd want to run it on GPU

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[Rocketknight1]

@lewtun yeah, this feels like an area where other changes could easily break things again. A regression test would definitely be appreciated!

[lewtun]

Added a regression test here @Rocketknight1

[Rocketknight1]

This looks like a great PR to me, and the regression test is nice! I left two nits, but both are just code style rather than anything that actually affects the output.

[Rocketknight1]

Nice! This is a clean test.

[lewtun]

@Rocketknight1 I've included your suggestions and the CI seems to be failing for reasons unrelated to this PR 😢 . What do you recommend doing in this case?

[Rocketknight1]

@lewtun CI looks green now, so I'm happy! cc @ArthurZucker / @LysandreJik for core maintainer review and we should be good to go.

[LysandreJik]

Thanks @lewtun for the PR and @Rocketknight1 for the review!