Automatically add generation tags to chat template for assistant_only_loss=True training (TRL Issue #4879)

[Neelectric]

What does this PR do?

This is a very first step at implementing the feature requested in #4879. For now the approach would be trying to recognise the model type and then replacing the chat template with a 'fixed' version. This would be helpful for the most common models (Qwen, Llama, OLMo etc) but I wonder if there's a more fancy way to go about this that automatically recognises when the assistant starts messaging, which we would want to train on.

Implements #4879 (feature request)

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? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes?
• 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.
@qgallouedec

[qgallouedec]

I don't think we can do it 100% automatically, but I think that there are not thousands of different chat templates out there, if we support the top 5 chat templates, we may covers a good majority of use cases

[qgallouedec]

This looks solid so far. A couple things still missing:

• A unit test.
• In SFTTrainer, we should call this function when assistant_only_loss is enabled (and when {% generation %} isn't in the chat template.

Also, in SFTConfig, we should explicitly mention that toggling assistant_only_loss may modify the tokenizer’s chat template.

---

Ideally, this feature wouldn’t mutate the tokenizer: we can avoid that by calling apply_chat_template with the chat_template argument set to the enhanced template (this is what we do in the GRPO trainer). That said, I’m okay with modifying the chat template for now, it keeps the implementation simpler.