Retain newlines in chat template when `continue_final_message=True`
Author: lewtun
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

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](https://arxiv.org/abs/2211.14275) with math reward models like [this one](https://huggingface.co/peiyi9979/math-shepherd-mistral-7b-prm)). 

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:

```py
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}.

```

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}.
'

=== 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}.

```

**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).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? 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](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] 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.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @ylacombe, @eustlb
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
CI seems to be failing for reasons unrelated to this PR

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?

> 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`?


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.

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

@Rocketknight1 I've added a try/except in https://github.com/huggingface/transformers/pull/34253/commits/0fd98e1a05bb987b9ee0cdcb2246de67d3665dd7

I've also tested that the Llama code still works

```py
# 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

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_34253). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

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

@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?

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