Return assistant generated tokens mask in apply_chat_template 
Author: yonigottesman
Pull request: yes

# What does this PR do?
This PR addresses issue huggingface/transformers#28950 and enhances the functionality of the `tokenizer.apply_chat_template` method when finetuning on chat datasets.

The method tokenizer.apply_chat_template is recommended for maintaining consistency with the model's original template during both training and inference phases. This practice ensures that conversations are processed in a uniform manner.

Moreover, during the finetuning process on chat datasets, it is crucial to exclude tokens from the "user" or "system" segments of the conversation. This exclusion is necessary because including these tokens would train the model to predict not only the "assistant" responses but also potential user queries, which is undesirable (and strange).

Currently, the `tokenizer.apply_chat_template` method does not provide a way to identify which tokens belong to the "assistant" response. To address this, the PR introduces a new parameter called `return_assistant_mask`. This parameter returns a mask that identifies tokens generated by the assistant, allowing for the appropriate creation of a labels arrays with ignore (-100) values during training.

Additionally, this PR proposes the introduction of a new keyword `generation` (name open for discussion) in the jinja2 chat template. This keyword is used to encapsulate the assistant’s response within your chat template.

Here is an example of the new api:
```
template = (
    "{% for message in messages %}"
    "{% if (message['role'] != 'assistant') %}"
    "{{'
' + message['role'] + '\n' + message['content'] + '
' + '\n'}}"
    "{% elif (message['role'] == 'assistant')%}"
    "{{'
' + message['role'] + '\n'}}"
    "{% generation %}"
    "{{message['content'] + '
'}}"
    "{% endgeneration %}"
    "{{'\n'}}"
    "{% endif %}"
    "{% endfor %}"
)
dummy_conversation = [
      {"role": "system", "content": "system message"},
      {"role": "user", "content": "user message"},
      {"role": "assistant", "content": "assistant\nmessage"},
      {"role": "user", "content": "user message 2"},
      {"role": "assistant", "content": "assistant message 2"},
]

output = tokenizer_r.apply_chat_template(
    dummy_conversations,
    chat_template=dummy_template,
    tokenize=True,
    return_assistant_mask=True,
    return_dict=True,
  )

labels = [output["input_ids"][index] if mask == 1 else -100 for index, mask in enumerate(output["assistant_mask"])]
```

There are some issues I would want to discuss during this pr:

* Is this API fine? maybe we should return the a `labels` key in the dict already and not bother with the intermediate mask.
* Name of the new tag? currently `generation` but maybe should be `assistant_response`? or anything you like.
* I think maybe I should add a warning if a user runs with `return_assistant_mask` but the tokenizer chat template hasn't changed yet to support this new tag. That way users will know the are probably training on wrong tokens.
* In 99% of finetuning examples I see people using the [trl trainer](https://github.com/huggingface/trl/blob/main/trl/trainer/sft_trainer.py#L131) with `packing=True`. My new changes wont be usable easily if people use that parameter and maybe we should think of my API while taking into consideration a refactor of the `packing` affect.





## 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](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 and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

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:
cc @lewtun and @xenova to this as well!

My thoughts on your questions:

> Is this API fine? maybe we should return the a labels key in the dict already and not bother with the intermediate mask.

I prefer just returning the labels with masking applied, rather than returning the mask for the user to apply.

>Name of the new tag? currently `generation` but maybe should be `assistant_response`? or anything you like.

I think `generation` is fine - `assistant_response` is very long!

> I think maybe I should add a warning if a user runs with return_assistant_mask but the tokenizer chat template hasn't changed yet to support this new tag. That way users will know the are probably training on wrong tokens. 

Agreed! I guess the easiest way to check this is to just do a string search for `{% generation %}` tags? Be careful, because you'll also need to check for variants like `{-` 

> n 99% of finetuning examples I see people using the [trl trainer](https://github.com/huggingface/trl/blob/main/trl/trainer/sft_trainer.py#L131) with `packing=True`. My new changes wont be usable easily if people use that parameter and maybe we should think of my API while taking into consideration a refactor of the `packing` affect.

Yes, there's already a [DataCollatorForCompletionOnlyLM](https://huggingface.co/docs/trl/en/sft_trainer#train-on-completions-only) which also requires `packing=False`. I feel like we can slot in with that easily enough!

I want to hear from @xenova and ideally someone using minijinja as well, though - how easily can we support this extension? Since it's only useful in training, maybe it's less critical to have it in `huggingface/jinja` or `TGI`, but at the very least we should be able to gracefully ignore the `generation` tags.

> I prefer just returning the labels with masking applied, rather than returning the mask for the user to apply.

I agree, but then what should be the ignore label? -100 (pytorch)?. Im not sure its a good idea to add another parameter `ignore_label`

I think `-100` is correct, yes! This is the standard value for Torch and Transformers, so we don't need an extra arg to change it.

yea i just thought of non pytorch users where -100 is not the default. 
Anyways I updated the code to return `labels`

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

@yonigottesman Thanks for working on this, this is a feature I am very much looking forward to. Hope this can be merged soon.

Yes, sorry for not checking in @yonigottesman! Do you have other features you want to add, or should we treat this as review-ready?

@Rocketknight1 this is ready to be reviewed yes :) 

On it!

@yonigottesman while I'm reviewing can you rebase/resolve the merge conflict? It's nothing major, but it'll block us merging the PR until it's ready. (Edit: Probably better to rebase because your branch is a little out of date by now, a rebase will catch any other issues before merging)

I agree it should be `assistant_mask` and not labels. I feel like the collator should be added here and not `trl` what do you think?

Yes, agree! It's also fine to leave that for a separate PR, and just add the mask functionality in this PR.

ok. fixed to now return mask

Got it! Ping me whenever you're ready for re-review.

ready 😀

This is amazing!! Looking forward to this new change!

fixed your suggestions. Do you think the docstring should contain a small example ? like this is the phi template with the new token:
```
    "{{ bos_token }}"
    "{% for message in messages %}"
    "{% if (message['role'] == 'user') %}"
    "{{'
' + '\n' + message['content'] + '
' + '\n' + '
' + '\n'}}"
    "{% elif (message['role'] == 'assistant') %}"
    "{%generation%}"
    "{{message['content'] + '
' + '\n'}}"
    "{%endgeneration%}"
    "{% endif %}"
    "{% endfor %}"
```

Yes! We could also add a section to the end of the [chat template documentation](https://huggingface.co/docs/transformers/main/en/chat_templating), under "Template Writing Tips" - it's a good way to make this feature visible to people. The source for it is `en/chat_templating.md`

I've updated that doc recently, though - make sure you rebase before editing it!

@yonigottesman Could you rebase on `main` to include upstream changes? This should resolve a lot of the current CI failures 

@amyeroberts done

thanks @amyeroberts. You are right the whole sharing of the two lists between the AssistantTracker and the code outside is fishy and dangerous. I came up with a more pythonic solution with context managers tell me what you think.
This way the lists are only created and managed outside the object and you need to "activate" the object with the lists to get it to work. no encapsulation broken by code touching inner object members, and kind of safe as long as we use it with `with` statement


@amyeroberts anything else I should update? 😄 

@yonigottesman There's been a update on main which should fix the hub tests. Could you try rebasing, this should hopefully resolve 

Amazing contribution! 🎉 🎉 🎉 It helps me a lot!

> Here is an example of the new api:
> 
> ```
> template = (
>     "{% for message in messages %}"
>     "{% if (message['role'] != 'assistant') %}"
>     "{{'
' + message['role'] + '\n' + message['content'] + '
' + '\n'}}"
>     "{% elif (message['role'] == 'assistant')%}"
>     "{{'
' + message['role'] + '\n'}}"
>     "{% generation %}"
>     "{{message['content'] + '
'}}"
>     "{% endgeneration %}"
>     "{{'\n'}}"
>     "{% endif %}"
>     "{% endfor %}"
> )
> dummy_conversation = [
>       {"role": "system", "content": "system message"},
>       {"role": "user", "content": "user message"},
>       {"role": "assistant", "content": "assistant\nmessage"},
>       {"role": "user", "content": "user message 2"},
>       {"role": "assistant", "content": "assistant message 2"},
> ]
> 
> output = tokenizer_r.apply_chat_template(
>     dummy_conversations,
>     chat_template=dummy_template,
>     tokenize=True,
>     return_assistant_mask=True,
>     return_dict=True,
>   )
> 
> labels = [output["input_ids"][index] if mask == 1 else -100 for index, mask in enumerate(output["assistant_mask"])]
> ```


Some spelling mistakes in this example:

> output = tokenizer_r.apply_chat_template(
>     dummy_conversations,
>     chat_template=dummy_template,
>     tokenize=True,
>     return_assistant_mask=True,
>     return_dict=True,
>   )

`dummy_conversation` instead of `dummy_conversations`, `template` instead of `dummy_template`

> labels = [output["input_ids"][index] if mask == 1 else -100 for index, mask in enumerate(output["assistant_mask"])]

`assistant_masks` instead of `assistant_mask`


Thank you for your work on this!

I'm having some issues though. When I run the example script from the tests, I don't seem to get any assistant tokens:

```python
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct")

dummy_template = (
    "{% for message in messages %}"
    "{% if (message['role'] != 'assistant') %}"
    "{{'
' + message['role'] + '\n' + message['content'] + '
' + '\n'}}"
    "{% elif (message['role'] == 'assistant')%}"
    "{{'
' + message['role'] + '\n'}}"
    "{% generation %}"
    "{{message['content'] + '
'}}"
    "{% endgeneration %}"
    "{{'\n'}}"
    "{% endif %}"
    "{% endfor %}"
)
conversations = [
    [
        {"role": "system", "content": "system message"},
        {"role": "user", "content": "user message"},
        {"role": "assistant", "content": "start turn 1 assistant message. end turn 1"},
        {"role": "user", "content": "user message 2"},
        {"role": "assistant", "content": "start turn 2 assistant message. end turn 2"},
    ],
    [
        {"role": "system", "content": "system message 3"},
        {"role": "user", "content": "user message 3"},
        {"role": "assistant", "content": "start turn 3 assistant message. end turn 3"},
        {"role": "user", "content": "user message 4"},
        {"role": "assistant", "content": "start turn 4 assistant message. end turn 4"},
    ],
]

output = tokenizer.apply_chat_template(
    conversations[0],
    chat_template=dummy_template,
    tokenize=True,
    return_assistant_tokens_mask=True,
    return_dict=True,
)
print("".join(map(str, output["assistant_masks"])))
```

For me, this prints out `0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000`

I think this bug is being caused by other tokens being printed out before the `{% generation %}`, within the same turn. For example, if I change the chat template to:
```python
dummy_template = (
    "{% for message in messages %}"
    "{% if (message['role'] != 'assistant') %}"
    "{{'
' + message['role'] + '\n' + message['content'] + '
' + '\n'}}"
    "{% elif (message['role'] == 'assistant')%}"
    "{% generation %}"
    "{{'
' + message['role'] + '\n'}}"
    "{{message['content'] + '
'}}"
    "{% endgeneration %}"
    "{{'\n'}}"
    "{% endif %}"
    "{% endfor %}"
)
```

it works correctly, printing out `0000000000000000000000000000000011111111111111111111111110000000000000000001111111111111111111111111`

I am running the latest version of transformers, 4.44.0

@avicooper1 there is a bug but its not about tokens before "generation" in the same turn. If you try a different tokenizer it will work.
There is something strange about the llama3 tokenizer (`PreTrainedTokenizerFast`) for some reason the `char_to_token` function isn't working as expected and my implementation is based on its result.
I opened an issue huggingface/tokenizers#1620. 

Given the issue, is there a workaround to get the assistant mask?

sadly for llama3 i dont think so :( 
other models that use the same tokenizer class `PreTrainedTokenizerFast` (but different config) do work for example `tiiuae/falcon-mamba-7b-instruct`. so i guess its something specific to the llama3 configuration 

Unfortunately the fact that the template needs to contain the `{% generation %}` part makes it very inflexible to use. Would it be somehow possible to just generate the mask base on the provided `user` `assistant` inputs?

@psinger there were several issues where some solutions were discussed but nothing was flexible enough as every model can have its own chat template and special tokens. see huggingface/transformers#28950 and huggingface/transformers#27609.
having a `{% generation %}` keyword was the only thing i could come up with, if you have any better idea that could be great ill be happy to try and implement it

I can also confirm that there's something strange about llama3's tokenizer such that it just can't work regardless of the template.

So in that case I should manually add {% generation %} for each tokenizer's template?

Hi! I could confirm this works on Mistral7B tokenizer, but doesn't work on any of the Llama tokenizers (tried 3 and 3.1 LlamaTokenizer). 

Same here, tried 3.2 model as well.

@thepowerfuldeez @kwanUm can you update the `tokenizers` package to latest and check? there was a fix https://github.com/huggingface/tokenizers/pull/1640  than should have fixed this issue.

Still doesn't work @yonigottesman, appereantly I need to add the {% generation %} tag manually? is there an existing fucntionality for it in hf or somewhere else?

Here's a reproduction
```
from transformers import AutoTokenizer

# Load the tokenizer for LLaMA 3.1
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")

# Define the conversation
conversation = [
    {"role": "system", "content": "blablabla"},
    {"role": "user", "content": "blablabla"},
    {"role": "assistant", "content": "blablabla"},
]



# Apply the chat template with tokenization and assistant mask
output = tokenizer.apply_chat_template(
    conversation,
    tokenize=True,
    return_assistant_tokens_mask=True,
    return_dict=True,
)

# Print the assistant_mask
print("With default chat template:")
print(output['assistant_masks'])

# Define the chat template with the {% generation %} tag
chat_template = (
    "{% for message in messages %}"
    "{% if message['role'] != 'assistant' %}"
    "{{ '
' + message['role'] + '\n' + message['content'] + '
\n' }}"
    "{% elif message['role'] == 'assistant' %}"
    "{{ '
' + message['role'] + '\n' }}"
    "{% generation %}"
    "{{ message['content'] + '
\n' }}"
    "{% endgeneration %}"
    "{% endif %}"
    "{% endfor %}"
)

# Apply the chat template with tokenization and assistant mask
output = tokenizer.apply_chat_template(
    conversation,
    chat_template=chat_template,
    tokenize=True,
    return_assistant_tokens_mask=True,
    return_dict=True,
)


# Print the assistant_mask
print("With custom chat template containing the {% generation %} tag:")
print(output['assistant_masks'])

Output:
```
With default chat template:
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
With custom chat template containing the {% generation %} tag:
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
```


```
pip list | grep -E "token|transfor|accele|sentence"

accelerate                1.0.1
asttokens                 2.4.1
tiktoken                  0.8.0
tokenizers                0.20.1
transformers              4.46.0

```

until tokenizers start adding the `generation` token to the tokenizer_config.json, you will have to add this manually

The following chat template is what I am using for Llama 3.2 1B and 3B models. While there are a few differences between Llama 3.1 and 3.2 chat templates, one can use this prompt as a reference point.

```python
LLAMA32_CHAT_TEMPLATE = """{{- bos_token }}
{%- if custom_tools is defined %}
    {%- set tools = custom_tools %}
{%- endif %}
{%- if not tools_in_user_message is defined %}
    {%- set tools_in_user_message = true %}
{%- endif %}
{%- if not date_string is defined %}
    {%- if strftime_now is defined %}
        {%- set date_string = strftime_now("%d %b %Y") %}
    {%- else %}
        {%- set date_string = "26 Jul 2024" %}
    {%- endif %}
{%- endif %}
{%- if not tools is defined %}
    {%- set tools = none %}
{%- endif %}

{#- This block extracts the system message, so we can slot it into the right place. #}
{%- if messages[0]['role'] == 'system' %}
    {%- set system_message = messages[0]['content']|trim %}
    {%- set messages = messages[1:] %}
{%- else %}
    {%- set system_message = "" %}
{%- endif %}

{#- System message #}
{{- "
system
\n\n" }}
{%- if tools is not none %}
    {{- "Environment: ipython\n" }}
{%- endif %}
{{- "Cutting Knowledge Date: December 2023\n" }}
{{- "Today Date: " + date_string + "\n\n" }}
{%- if tools is not none and not tools_in_user_message %}
    {{- "You have access to the following functions. To call a function, please respond with JSON for a function call." }}
    {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}.' }}
    {{- "Do not use variables.\n\n" }}
    {%- for t in tools %}
        {{- t | tojson(indent=4) }}
        {{- "\n\n" }}
    {%- endfor %}
{%- endif %}
{{- system_message }}
{{- "
" }}

{#- Custom tools are passed in a user message with some extra guidance #}
{%- if tools_in_user_message and not tools is none %}
    {#- Extract the first user message so we can plug it in here #}
    {%- if messages | length != 0 %}
        {%- set first_user_message = messages[0]['content']|trim %}
        {%- set messages = messages[1:] %}
    {%- else %}
        {{- raise_exception("Cannot put tools in the first user message when there's no first user message!") }}
{%- endif %}
    {{- '
user
\n\n' -}}
    {{- "Given the following functions, please respond with a JSON for a function call " }}
    {{- "with its proper arguments that best answers the given prompt.\n\n" }}
    {{- 'Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}.' }}
    {{- "Do not use variables.\n\n" }}
    {%- for t in tools %}
        {{- t | tojson(indent=4) }}
        {{- "\n\n" }}
    {%- endfor %}
    {{- first_user_message + "
"}}
{%- endif %}

{%- for message in messages %}
    {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
            {%- if message.role  != 'assistant' %}
                  {{- '
' + message['role'] + '
\n\n'+ message['content'] | trim + '
' }}
            {%- elif message.role  == 'assistant' %}
                  {{- '
' + message['role'] + '
\n\n'}}
                 {% generation %}
                 {{- message['content'] | trim + '
' }}
                 {% endgeneration %}
           {%- endif %}
    {%- elif 'tool_calls' in message %}
        {%- if not message.tool_calls|length == 1 %}
            {{- raise_exception("This model only supports single tool-calls at once!") }}
        {%- endif %}
        {%- set tool_call = message.tool_calls[0].function %}
        {{- '
assistant
\n\n' -}}
        {{- '{"name": "' + tool_call.name + '", ' }}
        {{- '"parameters": ' }}
        {{- tool_call.arguments | tojson }}
        {{- "}" }}
        {{- "
" }}
    {%- elif message.role == "tool" or message.role == "ipython" %}
        {{- "
ipython
\n\n" }}
        {%- if message.content is mapping or message.content is iterable %}
            {{- message.content | tojson }}
        {%- else %}
            {{- message.content }}
        {%- endif %}
        {{- "
" }}
    {%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
    {{- '
assistant
\n\n' }}
{%- endif %}"""
```

I found that if the assistant is the first turn, the assistant mask simply omits that. This actually has some use scenarios, so if it's a 'safe' logic implemented in the codes, it might be a bit over optimized

FYI, models still don’t add generation tags in their chat templates, e.g. due to potential incompatibilities with other frameworks like GGUF. 

See detailed discussion for Qwen3: https://huggingface.co/Qwen/Qwen3-8B/discussions/14#68532a1344245e3283b96c8a