Add {% generation %} markers for assistant-only loss masking

#2
by qgallouedec HF Staff - opened

Wraps the assistant branch of the chat template in {% generation %} / {% endgeneration %} so apply_chat_template(..., return_assistant_tokens_mask=True) yields a correct assistant-only loss mask (used e.g. by TRL's SFTConfig(assistant_only_loss=True)).

The <|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|> prompt cue stays outside the block (it is not generated by the model); the content and the trailing <|END_OF_TURN_TOKEN|> are inside, so the model is trained to stop.

Verified: rendering is byte-identical to the current template for both text and multimodal messages; the mask covers exactly the assistant content plus its end-of-turn token.

d-rau changed pull request status to merged

Sign up or log in to comment