# Copyright 2020-2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import re
import textwrap
from pathlib import Path

import pytest
import transformers
from packaging.version import Version
from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoProcessor, AutoTokenizer

from trl import clone_chat_template
from trl.chat_template_utils import (
    _CHAT_TEMPLATES_DIR,
    _SUPPORTS_PARSE_RESPONSE_TOOLS,
    _SUPPORTS_RESPONSE_TEMPLATE,
    add_response_schema,
    cohere2_chat_template,
    cohere_chat_template,
    cohere_training_chat_template,
    gemma4_v2_chat_template,
    gemma4_v3_chat_template,
    gemma4_v4_chat_template,
    gemma4_v5_chat_template,
    get_training_chat_template,
    has_generation_markers,
    is_chat_template_prefix_preserving,
    is_chat_template_stop_token_trained,
    lfm2_2_5_v2_chat_template,
    lfm2_2_5_v4_chat_template,
    lfm2_2_5_vl_v2_chat_template,
    lfm2_chat_template,
    lfm2_v2_chat_template,
    parse_response,
    qwen3_training_chat_template,
    supports_tool_calling,
)
from trl.data_utils import prepare_multimodal_messages

from .testing_utils import TrlTestCase, require_response_parsing, require_vision


def multiply(a: int, b: int) -> int:
    """
    Multiplies two integers.

    Args:
        a: The first integer.
        b: The second integer.
    """
    return a * b


def addition(a: int, b: int) -> int:
    """
    Adds two integers.

    Args:
        a: The first integer.
        b: The second integer.
    """
    return a + b


def get_weather(zip_code: str, days: int, threshold: float, metric: bool, fields: list[str]) -> str:
    """
    Gets the weather forecast.

    Args:
        zip_code: The zip code of the location.
        days: The number of days to forecast.
        threshold: The rain probability above which rain is reported.
        metric: Whether to use metric units.
        fields: The fields to report.
    """
    return "sunny"


class TestCloneChatTemplate(TrlTestCase):
    def test_clone(self):
        # This tokenizer doesn't have a chat_template by default
        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-GPT2LMHeadModel")
        model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-GPT2LMHeadModel")
        # This one has a chat_template by default
        source = "trl-internal-testing/tiny-Qwen3ForCausalLM"
        _, modified_tokenizer, _ = clone_chat_template(model, tokenizer, source)

        # Check if special tokens are correctly set
        assert modified_tokenizer.eos_token == "<|im_end|>"

    def test_clone_with_resize(self):
        # This tokenizer doesn't have a chat_template by default
        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-GPT2LMHeadModel")
        model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-GPT2LMHeadModel")
        # This one has a chat_template by default
        source = "trl-internal-testing/tiny-Qwen3ForCausalLM"
        modified_model, modified_tokenizer, _ = clone_chat_template(
            model, tokenizer, source, resize_to_multiple_of=123
        )

        # Check that the input embeddings have been resized to a multiple of 123
        assert (modified_model.vocab_size % 123) == 0
        # Check that the input embeddings size matches the tokenizer vocabulary size
        assert model.vocab_size == len(modified_tokenizer.vocab)

    def test_clone_with_resize_and_extra_tokens_already_in_vocab(self):
        # This tokenizer doesn't have a chat_template by default
        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-GPT2LMHeadModel")
        model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-GPT2LMHeadModel")
        # This one has a chat_template by default
        source = "trl-internal-testing/tiny-Qwen3ForCausalLM"
        # This will add <extra_id_0>, <extra_id_1>, ... to the tokenizer
        modified_model, modified_tokenizer, _ = clone_chat_template(
            model, tokenizer, source, resize_to_multiple_of=123
        )
        # Try if we can resize a tokenizer that already has extra these extra tokens
        modified_model, modified_tokenizer, _ = clone_chat_template(
            modified_model, modified_tokenizer, source, resize_to_multiple_of=124
        )

        # Check that the input embeddings have been resized to a multiple of 123
        assert (modified_model.vocab_size % 124) == 0
        # Check that the input embeddings size matches the tokenizer vocabulary size
        assert model.vocab_size == len(modified_tokenizer.vocab)

    def test_apply_new_chat_template(self):
        # This tokenizer doesn't have a chat_template by default
        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-GPT2LMHeadModel")
        model = AutoModelForCausalLM.from_pretrained("trl-internal-testing/tiny-GPT2LMHeadModel")
        # This one has a chat_template by default
        source = "trl-internal-testing/tiny-Qwen3ForCausalLM"
        _, modified_tokenizer, _ = clone_chat_template(model, tokenizer, source)
        messages = [
            {"role": "system", "content": "You are helpful"},
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "Hi, how can I help you?"},
        ]
        prompt = modified_tokenizer.apply_chat_template(messages, tokenize=False)

        assert (
            prompt
            == "<|im_start|>system\nYou are helpful<|im_end|>\n<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nHi, how can I help you?<|im_end|>\n"
        )

    def test_clone_with_sequence_classification_model(self):
        # This tokenizer doesn't have a chat_template by default
        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-GPTNeoXForSequenceClassification")
        model = AutoModelForSequenceClassification.from_pretrained(
            "trl-internal-testing/tiny-GPTNeoXForSequenceClassification"
        )
        # This one has a chat_template by default
        source = "trl-internal-testing/tiny-Qwen3ForCausalLM"
        _, modified_tokenizer, _ = clone_chat_template(model, tokenizer, source)

        # Check if special tokens are correctly set
        assert modified_tokenizer.eos_token == "<|im_end|>"


@pytest.mark.xfail(
    condition=Version(transformers.__version__) < Version("5.0.0"),
    reason="Response parsing is not supported in transformers versions below 5.0.0",
    strict=True,
)
@require_response_parsing
class TestAddResponseSchema:
    @pytest.mark.parametrize(
        "tokenizer_name",
        [
            pytest.param("trl-internal-testing/tiny-Glm4MoeForCausalLM", id="glm4moe"),
            pytest.param(
                "trl-internal-testing/tiny-GptOssForCausalLM",
                id="gptoss",
                marks=pytest.mark.xfail(
                    Version(transformers.__version__) < Version("5.5.0"),
                    reason="Upstream bug in response parsing (see #5753; fixed in transformers#45166)",
                    strict=True,
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5",
                id="lfm2-2.5",
                marks=pytest.mark.skipif(
                    not _SUPPORTS_RESPONSE_TEMPLATE,
                    reason="LFM2.5 only ships a new-style response template, which requires transformers>=5.13",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-Lfm2MoeForCausalLM-2.5",
                id="lfm2-2.5-moe",
                marks=pytest.mark.skipif(
                    not _SUPPORTS_RESPONSE_TEMPLATE,
                    reason="LFM2.5 only ships a new-style response template, which requires transformers>=5.13",
                ),
            ),
            pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3.1", id="llama3.1"),
            pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3.2", id="llama3.2"),
            pytest.param(
                "trl-internal-testing/tiny-NemotronHForCausalLM-nano",
                id="nemotron_3_nano",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.3.0"),
                    reason="Nemotron 3 tokenizer requires transformers>=5.3.0",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-NemotronHForCausalLM-super",
                id="nemotron_3_super",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.3.0"),
                    reason="Nemotron 3 tokenizer requires transformers>=5.3.0",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-NemotronHForCausalLM-ultra",
                id="nemotron_3_ultra",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.3.0"),
                    reason="Nemotron 3 tokenizer requires transformers>=5.3.0",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-NemotronHForCausalLM-3.5-lightning",
                id="nemotron_3_5_lightning",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.3.0"),
                    reason="Nemotron 3.5 tokenizer requires transformers>=5.3.0",
                ),
            ),
            pytest.param("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", id="qwen2.5"),
            pytest.param("trl-internal-testing/tiny-Qwen3MoeForCausalLM", id="qwen3"),
            pytest.param("trl-internal-testing/tiny-Qwen3ForCausalLM-Instruct-2507", id="qwen3_instruct_2507"),
        ],
    )
    def test_add_response_schema(self, tokenizer_name):
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        tokenizer = add_response_schema(tokenizer)
        messages = [
            {"role": "user", "content": "What is 3*4?"},
            {
                "role": "assistant",
                "tool_calls": [{"type": "function", "function": {"name": "multiply", "arguments": {"a": 3, "b": 4}}}],
            },
        ]
        prefix = tokenizer.apply_chat_template(messages[:1], tokenize=False, add_generation_prompt=True)
        text = tokenizer.apply_chat_template(messages, tokenize=False)
        response = text[len(prefix) :]
        # Here, we just test that the parsing doesn't raise an error.
        # The correctness of the parsing is tested in TestParseResponse
        # The new-style `response_template` parser requires the prompt prefix; the legacy `response_schema` parser
        # does not accept a `prefix=` argument.
        if _SUPPORTS_RESPONSE_TEMPLATE:
            tokenizer.parse_response(response, prefix=prefix)
        else:
            tokenizer.parse_response(response)

    @pytest.mark.parametrize(
        "tokenizer_name, chat_template",
        [
            pytest.param(
                "trl-internal-testing/tiny-Lfm2ForCausalLM",
                lfm2_v2_chat_template,
                id="lfm2-v2",
                marks=pytest.mark.skipif(
                    not _SUPPORTS_RESPONSE_TEMPLATE,
                    reason="LFM2 reuses the new-style response template of LFM2.5, which requires transformers>=5.13",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5",
                lfm2_2_5_v2_chat_template,
                id="lfm2-2.5-v2",
                marks=pytest.mark.skipif(
                    not _SUPPORTS_RESPONSE_TEMPLATE,
                    reason="LFM2.5 only ships a new-style response template, which requires transformers>=5.13",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5",
                lfm2_2_5_v4_chat_template,
                id="lfm2-2.5-v4",
                marks=pytest.mark.skipif(
                    not _SUPPORTS_RESPONSE_TEMPLATE,
                    reason="LFM2.5 only ships a new-style response template, which requires transformers>=5.13",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-Lfm2VlForConditionalGeneration-2.5",
                lfm2_2_5_vl_v2_chat_template,
                id="lfm2-2.5-vl-v2",
                marks=pytest.mark.skipif(
                    not _SUPPORTS_RESPONSE_TEMPLATE,
                    reason="LFM2.5 only ships a new-style response template, which requires transformers>=5.13",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-Qwen3MoeForCausalLM",
                qwen3_training_chat_template,
                id="qwen3-training",
            ),
        ],
    )
    def test_add_response_schema_template_revision(self, tokenizer_name, chat_template):
        # Same tiny model, with the chat template of a later Hub revision
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        tokenizer.chat_template = chat_template
        tokenizer = add_response_schema(tokenizer)
        tool_calls = [
            {"type": "function", "function": {"name": "multiply", "arguments": {"a": 3, "b": 4}}},
            {"type": "function", "function": {"name": "get_weather", "arguments": {"city": "Paris (FR), France"}}},
        ]
        messages = [
            {"role": "user", "content": "What is 3*4?"},
            {"role": "assistant", "content": "", "tool_calls": tool_calls},
        ]
        prefix = tokenizer.apply_chat_template(
            messages[:1], add_generation_prompt=True, tokenize=True, return_dict=True
        )
        text = tokenizer.apply_chat_template(messages, tokenize=True, return_dict=True)
        response = text.input_ids[len(prefix.input_ids) :]
        parsed = parse_response(tokenizer, response, prefix=prefix.input_ids)
        assert parsed == messages[-1]

    @pytest.mark.parametrize(
        "completion, expected",
        [
            pytest.param("</think>12<|im_end|>", {"role": "assistant", "content": "12"}, id="content"),
            pytest.param(
                "Hmmm.</think>12<|im_end|>",
                {"role": "assistant", "thinking": "Hmmm.", "content": "12"},
                id="thinking",
            ),
            pytest.param(
                "Hmmm.</think><|tool_call_start|>[multiply(a=3, b=4), get_weather(city='Paris', fields=[\"rain\"])]"
                "<|tool_call_end|><|im_end|>",
                {
                    "role": "assistant",
                    "thinking": "Hmmm.",
                    "content": "",
                    "tool_calls": [
                        {"type": "function", "function": {"name": "multiply", "arguments": {"a": 3, "b": 4}}},
                        {
                            "type": "function",
                            "function": {"name": "get_weather", "arguments": {"city": "Paris", "fields": ["rain"]}},
                        },
                    ],
                },
                id="tool_calls",
            ),
        ],
    )
    @pytest.mark.skipif(
        not _SUPPORTS_RESPONSE_TEMPLATE,
        reason="LFM2.5 only ships a new-style response template, which requires transformers>=5.13",
    )
    def test_add_response_schema_think_prefill(self, completion, expected):
        # The LFM2.5 v3 generation prompt ends with `<think>`, so the model's output starts inside the reasoning block,
        # and closes it right away when it skips reasoning. Rendered history turns have no such prefill, so the
        # completion can't be sliced off a rendered conversation.
        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Lfm2MoeForCausalLM-2.5")
        tokenizer = add_response_schema(tokenizer)
        messages = [{"role": "user", "content": "What is 3*4?"}]
        prefix = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True)
        response = tokenizer(completion, add_special_tokens=False).input_ids
        parsed = parse_response(tokenizer, response, prefix=prefix.input_ids)
        assert parsed == expected

    @pytest.mark.parametrize(
        "processor_name",
        [
            pytest.param(
                "trl-internal-testing/tiny-Lfm2VlForConditionalGeneration-2.5",
                id="lfm2-2.5-vl",
                marks=pytest.mark.skipif(
                    not _SUPPORTS_RESPONSE_TEMPLATE,
                    reason="LFM2.5 only ships a new-style response template, which requires transformers>=5.13",
                ),
            ),
            pytest.param("trl-internal-testing/tiny-Qwen3VLForConditionalGeneration", id="qwen3_vl"),
            pytest.param("trl-internal-testing/tiny-Qwen3_5ForConditionalGeneration-NoThink", id="qwen35-nothink"),
            pytest.param("trl-internal-testing/tiny-Qwen3_5ForConditionalGeneration-Think", id="qwen35-think"),
            pytest.param("trl-internal-testing/tiny-Qwen3_5MoeForConditionalGeneration-3.6", id="qwen36"),
            pytest.param("trl-internal-testing/tiny-Qwen3_5ForConditionalGeneration-3.8", id="qwen38"),
        ],
    )
    @require_vision
    def test_add_response_schema_vlm(self, processor_name):
        # For VLM processors, `add_response_schema` must set the template/schema on the inner tokenizer, since
        # `parse_response` is a tokenizer method that reads `self.response_template` / `self.response_schema` from the
        # tokenizer instance. Which attribute is set depends on the installed transformers version.
        processor = AutoProcessor.from_pretrained(processor_name)
        processor = add_response_schema(processor)
        if _SUPPORTS_RESPONSE_TEMPLATE:
            assert processor.tokenizer.response_template is not None
        else:
            assert processor.tokenizer.response_schema is not None
        messages = [
            {"role": "user", "content": [{"type": "text", "text": "What is 3*4?"}]},
            {
                "role": "assistant",
                # "content" is required here because VLM processors crash on tokenize=True without it
                # (KeyError in processing_utils.py). See huggingface/transformers#45290.
                "content": [{"type": "text", "text": ""}],
                "tool_calls": [{"type": "function", "function": {"name": "multiply", "arguments": {"a": 3, "b": 4}}}],
            },
        ]
        prefix = processor.apply_chat_template(messages[:1], tokenize=False, add_generation_prompt=True)
        text = processor.apply_chat_template(messages, tokenize=False)
        response = text[len(prefix) :]
        # Here, we just test that the parsing doesn't raise an error.
        # The correctness of the parsing is tested in TestParseResponse
        # The new-style `response_template` parser requires the prompt prefix; the legacy `response_schema` parser
        # does not accept a `prefix=` argument.
        if _SUPPORTS_RESPONSE_TEMPLATE:
            processor.tokenizer.parse_response(response, prefix=prefix)
        else:
            processor.tokenizer.parse_response(response)

    @pytest.mark.parametrize(
        "processor_name, chat_template",
        [
            pytest.param(
                "trl-internal-testing/tiny-Gemma4ForConditionalGeneration", gemma4_v2_chat_template, id="gemma4-v2"
            ),
            pytest.param(
                "trl-internal-testing/tiny-Gemma4ForConditionalGeneration", gemma4_v3_chat_template, id="gemma4-v3"
            ),
            pytest.param(
                "trl-internal-testing/tiny-Gemma4ForConditionalGeneration", gemma4_v4_chat_template, id="gemma4-v4"
            ),
            pytest.param(
                "trl-internal-testing/tiny-Gemma4ForConditionalGeneration", gemma4_v5_chat_template, id="gemma4-v5"
            ),
        ],
    )
    @pytest.mark.skipif(
        not _SUPPORTS_RESPONSE_TEMPLATE,
        reason="Gemma 4 response parsing is only provided as a new-style response template, which requires transformers>=5.13",
    )
    @require_vision
    def test_add_response_schema_vlm_template_revision(self, processor_name, chat_template):
        # Same tiny model, with the chat template of another Hub revision
        processor = AutoProcessor.from_pretrained(processor_name)
        processor.chat_template = chat_template
        processor = add_response_schema(processor)
        tool_calls = [{"type": "function", "function": {"name": "multiply", "arguments": {"a": 3, "b": 4}}}]
        messages = [
            {"role": "user", "content": "What is 3*4?"},
            {"role": "assistant", "content": "", "tool_calls": tool_calls},
        ]
        expected = messages[-1]
        messages = prepare_multimodal_messages(messages)
        prefix = processor.apply_chat_template(
            messages[:1], add_generation_prompt=True, tokenize=True, return_dict=True
        ).input_ids[0]
        text = processor.apply_chat_template(messages, tokenize=True, return_dict=True).input_ids[0]
        parsed = parse_response(processor.tokenizer, text[len(prefix) :], prefix=prefix)
        assert parsed == expected


class TestSupportsToolCalling:
    @pytest.mark.parametrize(
        "model_id",
        [
            pytest.param(
                "trl-internal-testing/tiny-BeamForCausalLM",
                id="beam",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.20.1"),
                    reason="Beam was introduced in transformers-5.20.1",
                ),
            ),
            pytest.param("trl-internal-testing/tiny-DeepseekV3ForCausalLM", id="deepseekv3"),
            pytest.param("trl-internal-testing/tiny-DeepseekV3ForCausalLM-0528", id="deepseekv3-0528"),
            pytest.param(
                "trl-internal-testing/tiny-Gemma4ForConditionalGeneration",
                id="gemma4",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.5.0"),
                    reason="Gemma4 models were introduced in transformers-5.5.0",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-Glm4MoeForCausalLM",
                id="glm4moe",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.0.0"),
                    reason="GLM4 tokenizer requires transformers>=5.0.0",
                ),
            ),
            pytest.param("trl-internal-testing/tiny-GptOssForCausalLM", id="gptoss"),
            pytest.param(
                "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5",
                id="lfm2-2.5",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.0.0"),
                    reason="LFM2.5 tokenizer requires transformers>=5.0.0",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-Lfm2MoeForCausalLM-2.5",
                id="lfm2-2.5-moe",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.0.0"),
                    reason="LFM2-MoE requires transformers>=5.0.0",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-Lfm2VlForConditionalGeneration-2.5",
                id="lfm2-2.5-vl",
                marks=[
                    require_vision,
                    pytest.mark.skipif(
                        Version(transformers.__version__) < Version("5.0.0"),
                        reason="LFM2.5-VL requires transformers>=5.0.0",
                    ),
                ],
            ),
            pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3.1", id="llama3.1"),
            pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3.2", id="llama3.2"),
            pytest.param(
                "trl-internal-testing/tiny-MistralLarge4ForConditionalGeneration",
                id="mistral_large4",
                marks=[
                    require_vision,
                    pytest.mark.skipif(
                        Version(transformers.__version__) < Version("5.22.0"),
                        reason="Mistral Large 4 was introduced in transformers-5.22.0",
                    ),
                ],
            ),
            pytest.param(
                "trl-internal-testing/tiny-Mistral3ForConditionalGeneration",
                id="mistral3",
                marks=[
                    require_vision,
                    pytest.mark.skipif(
                        Version(transformers.__version__) < Version("5.0.0"),
                        reason="Mistral Medium 3.5 tokenizer requires transformers>=5.0.0",
                    ),
                ],
            ),
            pytest.param(
                "trl-internal-testing/tiny-MuseGlimmerForConditionalGeneration",
                id="muse_glimmer",
                marks=[
                    require_vision,
                    pytest.mark.skipif(
                        Version(transformers.__version__) < Version("5.15.0"),
                        reason="Muse Glimmer was introduced in transformers-5.15.0",
                    ),
                ],
            ),
            pytest.param(
                "trl-internal-testing/tiny-NemotronHForCausalLM-nano",
                id="nemotron_3_nano",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.3.0"),
                    reason="Nemotron 3 tokenizer requires transformers>=5.3.0",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-NemotronHForCausalLM-super",
                id="nemotron_3_super",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.3.0"),
                    reason="Nemotron 3 tokenizer requires transformers>=5.3.0",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-NemotronHForCausalLM-ultra",
                id="nemotron_3_ultra",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.3.0"),
                    reason="Nemotron 3 tokenizer requires transformers>=5.3.0",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-NemotronHForCausalLM-3.5-lightning",
                id="nemotron_3_5_lightning",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.3.0"),
                    reason="Nemotron 3.5 tokenizer requires transformers>=5.3.0",
                ),
            ),
            pytest.param("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", id="qwen2.5"),
            pytest.param("trl-internal-testing/tiny-Qwen3ForCausalLM", id="qwen3"),
            pytest.param("trl-internal-testing/tiny-Qwen3ForCausalLM-Instruct-2507", id="qwen3_instruct_2507"),
            pytest.param("trl-internal-testing/tiny-Qwen3MoeForCausalLM", id="qwen3moe"),
            pytest.param(
                "trl-internal-testing/tiny-Qwen3VLForConditionalGeneration",
                id="qwen3_vl",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("4.57.0"),
                    reason="Qwen3-VL was introduced in transformers-4.57.0",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-Qwen3_5ForConditionalGeneration-NoThink",
                id="qwen35-nothink",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.0.0"),
                    reason="Qwen3.5 tokenizer requires transformers>=5.0.0",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-Qwen3_5ForConditionalGeneration-Think",
                id="qwen35-think",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.0.0"),
                    reason="Qwen3.5 tokenizer requires transformers>=5.0.0",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-Qwen3_5MoeForConditionalGeneration-3.6",
                id="qwen36",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.0.0"),
                    reason="Qwen3.5 tokenizer requires transformers>=5.0.0",
                ),
            ),
            pytest.param(
                "trl-internal-testing/tiny-Qwen3_5ForConditionalGeneration-3.8",
                id="qwen38",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.0.0"),
                    reason="Qwen3.5 tokenizer requires transformers>=5.0.0",
                ),
            ),
            pytest.param("trl-internal-testing/tiny-Lfm2ForCausalLM", id="lfm2"),
            pytest.param(
                "trl-internal-testing/tiny-MuseSparkForConditionalGeneration",
                id="muse_spark",
                marks=[
                    require_vision,
                    pytest.mark.skipif(
                        Version(transformers.__version__) < Version("5.20.0"),
                        reason="Muse Spark Open was introduced in transformers-5.20.0",
                    ),
                ],
            ),
        ],
    )
    def test_supports_tool_calling(self, model_id):
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        assert supports_tool_calling(tokenizer) is True

    @pytest.mark.parametrize(
        "model_id",
        [
            # No chat template
            pytest.param("trl-internal-testing/tiny-BartModel", id="bart"),
            pytest.param("trl-internal-testing/tiny-GPT2LMHeadModel", id="gpt2"),
            pytest.param("trl-internal-testing/tiny-GPTNeoXForCausalLM", id="gptneox"),
            pytest.param("trl-internal-testing/tiny-GPTNeoXForSequenceClassification", id="gptneox-seq"),
            pytest.param("trl-internal-testing/tiny-OPTForCausalLM", id="opt"),
            pytest.param("trl-internal-testing/tiny-T5ForConditionalGeneration", id="t5"),
            # TemplateError: rejects tool role sequence
            pytest.param("trl-internal-testing/tiny-CohereForCausalLM", id="cohere"),
            pytest.param("trl-internal-testing/tiny-FalconMambaForCausalLM", id="falconmamba"),
            pytest.param("trl-internal-testing/tiny-GemmaForCausalLM", id="gemma"),
            pytest.param("trl-internal-testing/tiny-Gemma2ForCausalLM", id="gemma2"),
            pytest.param("trl-internal-testing/tiny-Gemma3ForConditionalGeneration", id="gemma3"),
            pytest.param("trl-internal-testing/tiny-Idefics2ForConditionalGeneration", id="idefics2"),
            pytest.param("trl-internal-testing/tiny-Idefics3ForConditionalGeneration", id="idefics3"),
            pytest.param("trl-internal-testing/tiny-LlavaNextForConditionalGeneration", id="llava_next"),
            pytest.param("trl-internal-testing/tiny-MistralForCausalLM-0.1", id="mistral0.1"),
            pytest.param("trl-internal-testing/tiny-MistralForCausalLM-0.2", id="mistral0.2"),
            pytest.param("trl-internal-testing/tiny-SmolVLMForConditionalGeneration", id="smolvlm"),
            # Silently drops both tool_calls and tool messages
            pytest.param("trl-internal-testing/tiny-Cohere2ForCausalLM", id="cohere2"),
            pytest.param("trl-internal-testing/tiny-LlavaForConditionalGeneration", id="llava"),
            # Olmo3 uses a bespoke function-calling schema (a `functions`/`function_calls` string on the
            # message plus an `environment` role) instead of the standard `tools`/`tool_calls`/`tool`
            # interface, so a standard tool-calling conversation is silently dropped.
            pytest.param(
                "trl-internal-testing/tiny-Olmo3ForCausalLM",
                id="olmo3",
                marks=pytest.mark.skipif(
                    Version(transformers.__version__) < Version("4.57.0"),
                    reason="Olmo 3 was introduced in transformers>=4.57.0",
                ),
            ),
            pytest.param("trl-internal-testing/tiny-Phi3ForCausalLM-3", id="phi3"),
            pytest.param("trl-internal-testing/tiny-Phi3ForCausalLM-3.5", id="phi3.5"),
            # Renders tool message content as plain text but drops assistant tool_calls
            pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3", id="llama3"),
            # The earlier LFM2 revision renders `tools` into the system prompt and wraps tool message content in
            # <|tool_response_start|> / <|tool_response_end|>, but never reads `tool_calls`: the model is trained to
            # emit <|tool_call_start|> / <|tool_call_end|> as plain text inside `content`.
            pytest.param(("trl-internal-testing/tiny-Lfm2ForCausalLM", lfm2_chat_template), id="lfm2-v1"),
            # DeepSeek-R1-Distill renders `tool_calls` only when `content` is `None`, and never closes a single-call
            # turn with `<｜tool▁calls▁end｜>`.
            pytest.param("trl-internal-testing/tiny-Qwen2ForCausalLM-R1-Distill", id="r1_distill"),
            pytest.param("trl-internal-testing/tiny-Qwen2VLForConditionalGeneration", id="qwen2_vl"),
            pytest.param("trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration", id="qwen2.5_vl"),
        ],
    )
    def test_does_not_support_tool_calling(self, model_id):
        # A (model, chat template) pair stands for the same tiny model with the chat template of another Hub revision
        model_id, chat_template = model_id if isinstance(model_id, tuple) else (model_id, None)
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        if chat_template is not None:
            tokenizer.chat_template = chat_template
        assert supports_tool_calling(tokenizer) is False


class TestIsChatTemplatePrefixPreserving:
    def test_prefix_preserving_template(self):
        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Qwen3MoeForCausalLM")
        # docstyle-ignore
        tokenizer.chat_template = textwrap.dedent(r"""
        {%- for message in messages %}

        {%- if message.role == 'user' %}
            {{- '<|im_start|>user\n' + message.content + '<|im_end|>\n' }}
        {%- elif message.role == 'assistant' %}
            {{- '<|im_start|>assistant\n' + message.content }}
            {%- if message.tool_calls %}
                {%- for tool_call in message.tool_calls %}
                    {%- if tool_call.function %}
                        {%- set tool_call = tool_call.function %}
                    {%- endif %}
                    {{- '<tool_call>' + tool_call.name + '</tool_call>' }}
                {%- endfor %}
            {%- endif %}
            {{- '<|im_end|>\n' }}
        {%- elif message.role == 'tool' %}
            {{- '<|im_start|>tool\n' + message.content + '<|im_end|>\n' }}
        {%- endif %}

        {%- endfor %}

        {%- if add_generation_prompt %}
            {{- '<|im_start|>assistant\n' }}
        {%- endif %}""")
        assert is_chat_template_prefix_preserving(tokenizer) is True

    def test_non_prefix_preserving_template(self):
        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Qwen3MoeForCausalLM")
        # The following template is quite typical of models like Qwen3 and GPT-OSS, where the thinking part (even
        # empty) is only present for last assistant message, which makes it non-prefix-preserving: appending a tool
        # message changes the earlier output.
        # docstyle-ignore
        tokenizer.chat_template = textwrap.dedent(r"""
        {%- if messages[0].role == 'system' %}
            {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
        {%- endif %}
        {%- set ns = namespace(last_query_index=messages|length - 1) %}
        {%- for message in messages[::-1] %}
            {%- set index = (messages|length - 1) - loop.index0 %}
            {%- if message.role == "user" and message.content is string %}
                {%- set ns.last_query_index = index %}
                {%- break %}
            {%- endif %}
        {%- endfor %}
        {%- for message in messages %}
            {%- set content = message.content if message.content is string else '' %}
            {%- if message.role == "user" or (message.role == "system" and not loop.first) %}
                {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>\n' }}
            {%- elif message.role == "assistant" %}
                {%- set reasoning_content = '' %}
                {%- if message.reasoning_content is string %}
                    {%- set reasoning_content = message.reasoning_content %}
                {%- else %}
                    {%- if '</think>' in content %}
                        {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
                        {%- set content = content.split('</think>')[-1].lstrip('\n') %}
                    {%- endif %}
                {%- endif %}
                {%- if loop.index0 > ns.last_query_index %}
                    {%- if loop.last or (not loop.last and reasoning_content) %}
                        {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
                    {%- else %}
                        {{- '<|im_start|>' + message.role + '\n' + content }}
                    {%- endif %}
                {%- else %}
                    {{- '<|im_start|>' + message.role + '\n' + content }}
                {%- endif %}
                {%- if message.tool_calls %}
                    {%- for tool_call in message.tool_calls %}
                        {%- if tool_call.function %}
                            {%- set tool_call = tool_call.function %}
                        {%- endif %}
                        {{- '<tool_call>' + tool_call.name + '</tool_call>' }}
                    {%- endfor %}
                {%- endif %}
                {{- '<|im_end|>\n' }}
            {%- elif message.role == "tool" %}
                {{- '<|im_start|>tool\n' + content + '<|im_end|>\n' }}
            {%- endif %}
        {%- endfor %}
        {%- if add_generation_prompt %}
            {{- '<|im_start|>assistant\n' }}
            {%- if enable_thinking is defined and enable_thinking is false %}
                {{- '<think>\n\n</think>\n\n' }}
            {%- endif %}
        {%- endif %}""")
        assert is_chat_template_prefix_preserving(tokenizer) is False

    @require_vision
    def test_prefix_preserving_template_processor(self):
        processor = AutoProcessor.from_pretrained("trl-internal-testing/tiny-Qwen3VLForConditionalGeneration")
        # Simple prefix-preserving template that mirrors how Qwen-VL templates emit image tokens: a list-of-blocks
        # content is iterated, and `{"type": "image"}` blocks are rendered as `<|vision_start|><|image_pad|><|vision_end|>`.
        # docstyle-ignore
        processor.chat_template = textwrap.dedent(r"""
        {%- for message in messages %}

        {%- if message.role == 'user' %}
            {{- '<|im_start|>user\n' }}
            {%- if message.content is string %}
                {{- message.content }}
            {%- else %}
                {%- for content in message.content %}
                    {%- if content.type == 'image' or 'image' in content %}
                        {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
                    {%- elif 'text' in content %}
                        {{- content.text }}
                    {%- endif %}
                {%- endfor %}
            {%- endif %}
            {{- '<|im_end|>\n' }}
        {%- elif message.role == 'assistant' %}
            {{- '<|im_start|>assistant\n' }}
            {%- if message.content is string %}
                {{- message.content }}
            {%- else %}
                {%- for content in message.content %}
                    {%- if 'text' in content %}
                        {{- content.text }}
                    {%- endif %}
                {%- endfor %}
            {%- endif %}
            {%- if message.tool_calls %}
                {%- for tool_call in message.tool_calls %}
                    {%- if tool_call.function %}
                        {%- set tool_call = tool_call.function %}
                    {%- endif %}
                    {{- '<tool_call>' + tool_call.name + '</tool_call>' }}
                {%- endfor %}
            {%- endif %}
            {{- '<|im_end|>\n' }}
        {%- elif message.role == 'tool' %}
            {{- '<|im_start|>tool\n' }}
            {%- if message.content is string %}
                {{- message.content }}
            {%- else %}
                {%- for content in message.content %}
                    {%- if 'text' in content %}
                        {{- content.text }}
                    {%- endif %}
                {%- endfor %}
            {%- endif %}
            {{- '<|im_end|>\n' }}
        {%- endif %}

        {%- endfor %}

        {%- if add_generation_prompt %}
            {{- '<|im_start|>assistant\n' }}
        {%- endif %}""")
        assert is_chat_template_prefix_preserving(processor) is True


class TestIsChatTemplateStopTokenTrained:
    def test_stop_token_trained(self):
        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Qwen3MoeForCausalLM")
        # The assistant turn is closed by <|im_end|> inside the generation span, so the end-of-turn token is masked
        # in and the model is trained to stop.
        # docstyle-ignore
        tokenizer.chat_template = textwrap.dedent(r"""
        {%- for message in messages %}
        {%- if message.role == 'user' %}
            {{- '<|im_start|>user\n' + message.content + '<|im_end|>\n' }}
        {%- elif message.role == 'assistant' %}
            {{- '<|im_start|>assistant\n' }}
            {%- generation %}{{- message.content + '<|im_end|>' }}{%- endgeneration %}
            {{- '\n' }}
        {%- endif %}
        {%- endfor %}
        {%- if add_generation_prompt %}
            {{- '<|im_start|>assistant\n' }}
        {%- endif %}""")
        assert is_chat_template_stop_token_trained(tokenizer) is True

    def test_stop_token_not_trained(self):
        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Qwen3MoeForCausalLM")
        # GLM-style: the assistant's end-of-turn token is emitted as the prefix of the following message, so the
        # generation span covers content only and the model is never trained to stop.
        # docstyle-ignore
        tokenizer.chat_template = textwrap.dedent(r"""
        {%- for message in messages %}
        {%- if message.role == 'user' %}
            {{- '<|im_start|>user\n' + message.content + '<|im_end|>\n' }}
        {%- elif message.role == 'assistant' %}
            {{- '<|im_start|>assistant\n' }}
            {%- generation %}{{- message.content }}{%- endgeneration %}
            {{- '<|im_end|>\n' }}
        {%- endif %}
        {%- endfor %}
        {%- if add_generation_prompt %}
            {{- '<|im_start|>assistant\n' }}
        {%- endif %}""")
        assert is_chat_template_stop_token_trained(tokenizer) is False

    def test_template_error_returns_false(self):
        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-Qwen3MoeForCausalLM")
        tokenizer.chat_template = "{{ raise_exception('probe rejected') }}"
        assert is_chat_template_stop_token_trained(tokenizer) is False


@pytest.mark.parametrize(
    "tokenizer_name",
    [
        pytest.param(
            "trl-internal-testing/tiny-BeamForCausalLM",
            id="beam",
            marks=pytest.mark.skipif(
                Version(transformers.__version__) < Version("5.20.1"),
                reason="Beam was introduced in transformers-5.20.1",
            ),
        ),
        pytest.param("trl-internal-testing/tiny-CohereForCausalLM", id="cohere"),
        # Real Cohere checkpoints ship their chat template as named variants, which load as a dict
        pytest.param(
            ("trl-internal-testing/tiny-CohereForCausalLM", {"default": cohere_chat_template}),
            id="cohere-named-variants",
        ),
        pytest.param("trl-internal-testing/tiny-Cohere2ForCausalLM", id="cohere2"),
        pytest.param(
            ("trl-internal-testing/tiny-Cohere2ForCausalLM", {"default": cohere2_chat_template}),
            id="cohere2-named-variants",
        ),
        pytest.param("trl-internal-testing/tiny-Qwen2ForCausalLM-R1-Distill", id="r1_distill"),
        pytest.param("trl-internal-testing/tiny-DeepseekV3ForCausalLM", id="deepseekv3"),
        pytest.param(
            "trl-internal-testing/tiny-DiffusionGemmaForBlockDiffusion",
            id="diffusion_gemma",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.11.0"),
                    reason="DiffusionGemma was introduced in transformers 5.11.0",
                ),
            ],
        ),
        pytest.param("trl-internal-testing/tiny-GemmaForCausalLM", id="gemma"),
        pytest.param("trl-internal-testing/tiny-Gemma2ForCausalLM", id="gemma2"),
        pytest.param("trl-internal-testing/tiny-Gemma3ForConditionalGeneration", id="gemma3", marks=require_vision),
        pytest.param(
            "trl-internal-testing/tiny-Glm4MoeForCausalLM",
            id="glm4moe",
            marks=pytest.mark.skipif(
                Version(transformers.__version__) < Version("5.0.0"),
                reason="GLM4 tokenizer requires transformers>=5.0.0",
            ),
        ),
        pytest.param("trl-internal-testing/tiny-GptOssForCausalLM", id="gptoss"),
        pytest.param(
            "trl-internal-testing/tiny-Idefics3ForConditionalGeneration", id="idefics3", marks=require_vision
        ),
        pytest.param("trl-internal-testing/tiny-Lfm2ForCausalLM", id="lfm2"),
        pytest.param(("trl-internal-testing/tiny-Lfm2ForCausalLM", lfm2_chat_template), id="lfm2-v1"),
        pytest.param(("trl-internal-testing/tiny-Lfm2ForCausalLM", lfm2_v2_chat_template), id="lfm2-v2"),
        pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3", id="llama3"),
        pytest.param("trl-internal-testing/tiny-LlavaForConditionalGeneration", id="llava", marks=require_vision),
        pytest.param(
            "trl-internal-testing/tiny-LlavaNextForConditionalGeneration", id="llava_next", marks=require_vision
        ),
        pytest.param(
            "trl-internal-testing/tiny-MistralLarge4ForConditionalGeneration",
            id="mistral_large4",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.22.0"),
                    reason="Mistral Large 4 was introduced in transformers-5.22.0",
                ),
            ],
        ),
        pytest.param(
            "trl-internal-testing/tiny-Mistral3ForConditionalGeneration",
            id="mistral3",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.0.0"),
                    reason="Mistral Medium 3.5 tokenizer requires transformers>=5.0.0",
                ),
            ],
        ),
        pytest.param(
            "trl-internal-testing/tiny-MuseGlimmerForConditionalGeneration",
            id="muse_glimmer",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.15.0"),
                    reason="Muse Glimmer was introduced in transformers-5.15.0",
                ),
            ],
        ),
        pytest.param(
            "trl-internal-testing/tiny-NemotronHForCausalLM-nano",
            id="nemotron_3_nano",
            marks=pytest.mark.skipif(
                Version(transformers.__version__) < Version("5.3.0"),
                reason="Nemotron 3 tokenizer requires transformers>=5.3.0",
            ),
        ),
        pytest.param(
            "trl-internal-testing/tiny-NemotronHForCausalLM-super",
            id="nemotron_3_super",
            marks=pytest.mark.skipif(
                Version(transformers.__version__) < Version("5.3.0"),
                reason="Nemotron 3 tokenizer requires transformers>=5.3.0",
            ),
        ),
        pytest.param(
            "trl-internal-testing/tiny-NemotronHForCausalLM-ultra",
            id="nemotron_3_ultra",
            marks=pytest.mark.skipif(
                Version(transformers.__version__) < Version("5.3.0"),
                reason="Nemotron 3 tokenizer requires transformers>=5.3.0",
            ),
        ),
        pytest.param(
            "trl-internal-testing/tiny-NemotronHForCausalLM-3.5-lightning",
            id="nemotron_3_5_lightning",
            marks=pytest.mark.skipif(
                Version(transformers.__version__) < Version("5.3.0"),
                reason="Nemotron 3.5 tokenizer requires transformers>=5.3.0",
            ),
        ),
        pytest.param("trl-internal-testing/tiny-Phi3ForCausalLM-3", id="phi3"),
        pytest.param("trl-internal-testing/tiny-Phi3ForCausalLM-3.5", id="phi3.5"),
        pytest.param("trl-internal-testing/tiny-Qwen2VLForConditionalGeneration", id="qwen2_vl", marks=require_vision),
        pytest.param("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", id="qwen2.5"),
        pytest.param(
            "trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration", id="qwen2.5_vl", marks=require_vision
        ),
        pytest.param("trl-internal-testing/tiny-Qwen3MoeForCausalLM", id="qwen3"),
        pytest.param("trl-internal-testing/tiny-Qwen3ForCausalLM-Instruct-2507", id="qwen3_instruct_2507"),
        pytest.param(
            "trl-internal-testing/tiny-Qwen3VLForConditionalGeneration",
            id="qwen3_vl",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    Version(transformers.__version__) < Version("4.57.0"),
                    reason="Qwen3-VL was introduced in transformers-4.57.0",
                ),
            ],
        ),
        pytest.param(
            "trl-internal-testing/tiny-Qwen3_5ForConditionalGeneration-NoThink",
            id="qwen35-nothink",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.0.0"),
                    reason="Qwen3.5 tokenizer requires transformers>=5.0.0",
                ),
            ],
        ),
        pytest.param(
            "trl-internal-testing/tiny-Qwen3_5ForConditionalGeneration-Think",
            id="qwen35-think",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.0.0"),
                    reason="Qwen3.5 tokenizer requires transformers>=5.0.0",
                ),
            ],
        ),
        pytest.param(
            "trl-internal-testing/tiny-Qwen3_5MoeForConditionalGeneration-3.6",
            id="qwen36",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.0.0"),
                    reason="Qwen3.5 tokenizer requires transformers>=5.0.0",
                ),
            ],
        ),
        pytest.param(
            "trl-internal-testing/tiny-Qwen3_5ForConditionalGeneration-3.8",
            id="qwen38",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.0.0"),
                    reason="Qwen3.5 tokenizer requires transformers>=5.0.0",
                ),
            ],
        ),
        pytest.param("trl-internal-testing/tiny-SmolVLMForConditionalGeneration", id="smolvlm", marks=require_vision),
        pytest.param(
            "trl-internal-testing/tiny-MuseSparkForConditionalGeneration",
            id="muse_spark",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.20.0"),
                    reason="Muse Spark Open was introduced in transformers-5.20.0",
                ),
            ],
        ),
    ],
)
class TestGetTrainingChatTemplate:
    def _load(self, model_name):
        # A (model, chat template) pair stands for the same tiny model with another form of its chat template
        model_name, chat_template = model_name if isinstance(model_name, tuple) else (model_name, None)
        if "ForCausalLM" in model_name:
            self.is_vlm = False
            processing_class = AutoTokenizer.from_pretrained(model_name)
        elif "ForConditionalGeneration" in model_name or "ForBlockDiffusion" in model_name:
            self.is_vlm = True
            processing_class = AutoProcessor.from_pretrained(model_name)
        if chat_template is not None:
            processing_class.chat_template = chat_template

        return processing_class

    def test_new_chat_template_is_prefix_preserving(self, tokenizer_name):
        tokenizer = self._load(tokenizer_name)
        new_chat_template = get_training_chat_template(tokenizer)
        if new_chat_template is not None:
            tokenizer.chat_template = new_chat_template
        # Prefix-preservation is only meaningful for templates that actually support tool messages — the check
        # itself renders one. Skip the assertion for tool-less templates (e.g. Gemma).
        if not supports_tool_calling(tokenizer):
            pytest.skip("Template does not support tool calling; prefix-preservation check is not applicable.")
        assert is_chat_template_prefix_preserving(tokenizer) is True

    def test_new_chat_template_trains_stop_token(self, tokenizer_name, request):
        if tokenizer_name == "trl-internal-testing/tiny-LlavaForConditionalGeneration":
            reason = "Llava's official chat template emits no end-of-turn token after the assistant turn."
            request.node.add_marker(pytest.mark.xfail(strict=True, reason=reason))
        tokenizer = self._load(tokenizer_name)
        new_chat_template = get_training_chat_template(tokenizer)
        assert is_chat_template_stop_token_trained(tokenizer, chat_template=new_chat_template) is True

    def test_behavior_unchanged_single_user_no_generation_prompt(self, tokenizer_name):
        tokenizer = self._load(tokenizer_name)
        messages = [{"role": "user", "content": "What color is the sky?"}]
        if self.is_vlm:
            messages = prepare_multimodal_messages(messages)

        before = tokenizer.apply_chat_template(messages, tokenize=False)
        new_chat_template = get_training_chat_template(tokenizer)
        after = tokenizer.apply_chat_template(messages, tokenize=False, chat_template=new_chat_template)
        assert before == after

    def test_behavior_unchanged_single_user_with_generation_prompt(self, tokenizer_name):
        tokenizer = self._load(tokenizer_name)
        messages = [{"role": "user", "content": "What color is the sky?"}]
        if self.is_vlm:
            messages = prepare_multimodal_messages(messages)

        before = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
        new_chat_template = get_training_chat_template(tokenizer)
        after = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            chat_template=new_chat_template,
        )
        assert before == after

    def test_behavior_unchanged_single_user_and_final_assistant_plain_content(self, tokenizer_name):
        tokenizer = self._load(tokenizer_name)
        messages = [
            {"role": "user", "content": "What color is the sky?"},
            {"role": "assistant", "content": "It is blue."},
        ]
        if self.is_vlm:
            messages = prepare_multimodal_messages(messages)

        before = tokenizer.apply_chat_template(messages, tokenize=False)
        new_chat_template = get_training_chat_template(tokenizer)
        after = tokenizer.apply_chat_template(messages, tokenize=False, chat_template=new_chat_template)
        if tokenizer_name == "trl-internal-testing/tiny-Glm4MoeForCausalLM":
            # GLM's native template doesn't terminate an assistant turn with an end-of-turn token; the turn is ended
            # by the following message's role marker. The training template appends that terminator to the final
            # assistant turn so the stop token is trained — here the `<|user|>` that would open the next turn.
            assert after == before + "<|user|>"
        else:
            assert before == after

    def test_behavior_unchanged_final_assistant_with_reasoning_content(self, tokenizer_name):
        tokenizer = self._load(tokenizer_name)
        messages = [
            {"role": "user", "content": "What color is the sky?"},
            {
                "role": "assistant",
                "content": "It is blue.",
                "reasoning_content": "The sky appears blue due to Rayleigh scattering.",
            },
        ]
        if self.is_vlm:
            messages = prepare_multimodal_messages(messages)

        before = tokenizer.apply_chat_template(messages, tokenize=False)
        new_chat_template = get_training_chat_template(tokenizer)
        after = tokenizer.apply_chat_template(messages, tokenize=False, chat_template=new_chat_template)
        if tokenizer_name == "trl-internal-testing/tiny-Glm4MoeForCausalLM":
            # GLM's native template doesn't terminate an assistant turn with an end-of-turn token; the turn is ended
            # by the following message's role marker. The training template appends that terminator to the final
            # assistant turn so the stop token is trained — here the `<|user|>` that would open the next turn.
            assert after == before + "<|user|>"
        else:
            assert before == after

    def test_behavior_unchanged_final_assistant_with_existing_think_tags(self, tokenizer_name):
        tokenizer = self._load(tokenizer_name)
        messages = [
            {"role": "user", "content": "What color is the sky?"},
            {
                "role": "assistant",
                "content": "<think>\nThe sky scatters shorter wavelengths.\n</think>\n\nIt is blue.",
            },
        ]
        if self.is_vlm:
            messages = prepare_multimodal_messages(messages)

        before = tokenizer.apply_chat_template(messages, tokenize=False)
        new_chat_template = get_training_chat_template(tokenizer)
        after = tokenizer.apply_chat_template(messages, tokenize=False, chat_template=new_chat_template)
        if tokenizer_name == "trl-internal-testing/tiny-Glm4MoeForCausalLM":
            # GLM's native template doesn't terminate an assistant turn with an end-of-turn token; the turn is ended
            # by the following message's role marker. The training template appends that terminator to the final
            # assistant turn so the stop token is trained — here the `<|user|>` that would open the next turn.
            assert after == before + "<|user|>"
        elif tokenizer_name == "trl-internal-testing/tiny-Qwen2ForCausalLM-R1-Distill":
            # DeepSeek-R1-Distill's native template drops everything up to and including `</think>` when re-rendering
            # an assistant turn. The training template keeps the reasoning, so it stays in the SFT target.
            reasoning = "<think>\nThe sky scatters shorter wavelengths.\n</think>"
            assert after == before.replace("<｜Assistant｜>", "<｜Assistant｜>" + reasoning)
        else:
            assert before == after

    def test_behavior_unchanged_assistant_with_tool_calls(self, tokenizer_name):
        tokenizer = self._load(tokenizer_name)
        tool_calls = [{"type": "function", "function": {"name": "multiply", "arguments": {"a": 3, "b": 4}}}]
        messages = [
            {"role": "user", "content": "Multiply 3 by 4."},
            {"role": "assistant", "content": "I will call a tool.", "tool_calls": tool_calls},
        ]
        if self.is_vlm:
            messages = prepare_multimodal_messages(messages)

        messages_before = copy.deepcopy(messages)
        if tokenizer_name == "trl-internal-testing/tiny-DeepseekV3ForCausalLM":
            # Best-effort fallback for templates that reject dict args (e.g. DeepSeek-V3). This is a chat template
            # bug (see transformers#45419), and the training chat template fixes it to avoid blocking users.
            messages_before[1]["tool_calls"][0]["function"]["arguments"] = '{"a": 3, "b": 4}'

        before = tokenizer.apply_chat_template(messages_before, tokenize=False)
        new_chat_template = get_training_chat_template(tokenizer)
        after = tokenizer.apply_chat_template(messages, tokenize=False, chat_template=new_chat_template)
        if tokenizer_name == "trl-internal-testing/tiny-Glm4MoeForCausalLM":
            # GLM's native template doesn't terminate an assistant turn with an end-of-turn token; the turn is ended
            # by the following message's role marker. The training template appends that terminator to the final
            # assistant turn so the stop token is trained — here `<|observation|>`, which closes a tool call.
            assert after == before + "<|observation|>"
        else:
            assert before == after

    def test_behavior_unchanged_with_tools_with_and_without_system_message(self, tokenizer_name):
        tokenizer = self._load(tokenizer_name)
        tools = [
            {
                "type": "function",
                "function": {
                    "name": "multiply",
                    "description": "Multiply two numbers.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "a": {"type": "number"},
                            "b": {"type": "number"},
                        },
                        "required": ["a", "b"],
                    },
                },
            }
        ]
        messages = [{"role": "user", "content": "Multiply 3 by 4."}]
        if self.is_vlm:
            messages = prepare_multimodal_messages(messages)

        before = tokenizer.apply_chat_template(messages, tokenize=False, tools=tools)
        new_chat_template = get_training_chat_template(tokenizer)
        after = tokenizer.apply_chat_template(messages, tokenize=False, tools=tools, chat_template=new_chat_template)
        assert before == after

    def test_behavior_unchanged_with_tools_with_system_message(self, tokenizer_name):
        tokenizer = self._load(tokenizer_name)
        if not supports_tool_calling(tokenizer):
            pytest.skip("Template does not support tool calling; skipping tool_calls test.")
        tools = [
            {
                "type": "function",
                "function": {
                    "name": "multiply",
                    "description": "Multiply two numbers.",
                    "parameters": {
                        "type": "object",
                        "properties": {"a": {"type": "number"}, "b": {"type": "number"}},
                        "required": ["a", "b"],
                    },
                },
            }
        ]
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "Multiply 3 by 4."},
        ]
        if self.is_vlm:
            messages = prepare_multimodal_messages(messages)

        before = tokenizer.apply_chat_template(messages, tokenize=False, tools=tools)
        new_chat_template = get_training_chat_template(tokenizer)
        after = tokenizer.apply_chat_template(messages, tokenize=False, tools=tools, chat_template=new_chat_template)
        assert before == after

    def test_behavior_unchanged_generation_prompt_with_enable_thinking_false(self, tokenizer_name):
        tokenizer = self._load(tokenizer_name)
        if "enable_thinking" not in tokenizer.chat_template:
            pytest.skip("This chat template doesn't use enable_thinking.")
        messages = [{"role": "user", "content": "What color is the sky?"}]
        if self.is_vlm:
            messages = prepare_multimodal_messages(messages)

        before = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
        )
        new_chat_template = get_training_chat_template(tokenizer)
        after = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=False,
            chat_template=new_chat_template,
        )
        assert before == after

    def test_assistant_masks(self, tokenizer_name, request):
        if tokenizer_name == "trl-internal-testing/tiny-LlavaForConditionalGeneration" and Version(
            transformers.__version__
        ) < Version("5.18.0.dev0"):
            request.node.add_marker(
                pytest.mark.xfail(
                    reason="Llava's official chat template `{% generation %}` markers don't yield assistant masks "
                    "through the processor path. It is not a supported training template.",
                    strict=True,
                )
            )
        tokenizer = self._load(tokenizer_name)
        messages = [
            {"role": "user", "content": "What color is the sky?"},
            {"role": "assistant", "content": "It is blue."},
        ]
        if self.is_vlm:
            messages = prepare_multimodal_messages(messages)

        chat_template = get_training_chat_template(tokenizer)
        result = tokenizer.apply_chat_template(
            messages, chat_template=chat_template, return_assistant_tokens_mask=True, return_dict=True, tokenize=True
        )
        masks = result["assistant_masks"]
        if self.is_vlm:  # VLM processors return batched output
            masks = masks[0]
        # The first tokens (user turn) should not be masked
        assert masks[0] == 0
        # Should have one masked region (the assistant turn)
        region_starts = sum(1 for i in range(1, len(masks)) if masks[i] == 1 and masks[i - 1] == 0)
        assert region_starts == 1

    def test_assistant_masks_multi_turn(self, tokenizer_name, request):
        if tokenizer_name == "trl-internal-testing/tiny-LlavaForConditionalGeneration" and Version(
            transformers.__version__
        ) < Version("5.18.0.dev0"):
            request.node.add_marker(
                pytest.mark.xfail(
                    reason="Llava's official chat template `{% generation %}` markers don't yield assistant masks "
                    "through the processor path. It is not a supported training template.",
                    strict=True,
                )
            )
        tokenizer = self._load(tokenizer_name)
        messages = [
            {"role": "user", "content": "Hi"},
            {"role": "assistant", "content": "Hello!"},
            {"role": "user", "content": "Bye"},
            {"role": "assistant", "content": "Goodbye!"},
        ]
        if self.is_vlm:
            messages = prepare_multimodal_messages(messages)

        chat_template = get_training_chat_template(tokenizer)
        result = tokenizer.apply_chat_template(
            messages, chat_template=chat_template, return_assistant_tokens_mask=True, return_dict=True, tokenize=True
        )
        masks = result["assistant_masks"]
        if self.is_vlm:  # VLM processors return batched output
            masks = masks[0]
        # Should have two masked regions (two assistant turns)
        region_starts = sum(1 for i in range(1, len(masks)) if masks[i] == 1 and masks[i - 1] == 0)
        assert region_starts == 2


class TestNamedChatTemplateVariants:
    # A chat template shipped as named variants loads as a dict, e.g. {"default": ..., "tool_use": ..., "rag": ...}
    def test_has_generation_markers(self):
        assert has_generation_markers({"default": cohere_training_chat_template}) is True
        assert has_generation_markers({"default": cohere_chat_template}) is False

    def test_unsupported_default_variant(self):
        tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/tiny-CohereForCausalLM")
        tokenizer.chat_template = {"default": "{{ messages[0]['content'] }}", "tool_use": cohere_chat_template}
        with pytest.raises(ValueError, match="not training-compatible"):
            get_training_chat_template(tokenizer)


@pytest.mark.parametrize(
    "model_name",
    [
        pytest.param(
            "trl-internal-testing/tiny-BeamForCausalLM",
            id="beam",
            marks=pytest.mark.skipif(
                Version(transformers.__version__) < Version("5.20.1"),
                reason="Beam was introduced in transformers-5.20.1",
            ),
        ),
        pytest.param("trl-internal-testing/tiny-Glm4MoeForCausalLM", id="glm4moe"),
        pytest.param("trl-internal-testing/tiny-GptOssForCausalLM", id="gptoss"),
        pytest.param(
            "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5",
            id="lfm2-2.5",
            marks=pytest.mark.skipif(
                not _SUPPORTS_RESPONSE_TEMPLATE,
                reason="LFM2.5 only ships a new-style response template, which requires transformers>=5.13",
            ),
        ),
        pytest.param(
            "trl-internal-testing/tiny-Lfm2VlForConditionalGeneration-2.5",
            id="lfm2-2.5-vl",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    not _SUPPORTS_RESPONSE_TEMPLATE,
                    reason="LFM2.5 only ships a new-style response template, which requires transformers>=5.13",
                ),
            ],
        ),
        pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3.1", id="llama3.1"),
        pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3.2", id="llama3.2"),
        pytest.param(
            "trl-internal-testing/tiny-MuseGlimmerForConditionalGeneration",
            id="muse_glimmer",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.15.0"),
                    reason="Muse Glimmer was introduced in transformers-5.15.0",
                ),
            ],
        ),
        pytest.param(
            "trl-internal-testing/tiny-NemotronHForCausalLM-nano",
            id="nemotron_3_nano",
            marks=pytest.mark.skipif(
                Version(transformers.__version__) < Version("5.3.0"),
                reason="Nemotron 3 tokenizer requires transformers>=5.3.0",
            ),
        ),
        pytest.param(
            "trl-internal-testing/tiny-NemotronHForCausalLM-super",
            id="nemotron_3_super",
            marks=pytest.mark.skipif(
                Version(transformers.__version__) < Version("5.3.0"),
                reason="Nemotron 3 tokenizer requires transformers>=5.3.0",
            ),
        ),
        pytest.param(
            "trl-internal-testing/tiny-NemotronHForCausalLM-ultra",
            id="nemotron_3_ultra",
            marks=pytest.mark.skipif(
                Version(transformers.__version__) < Version("5.3.0"),
                reason="Nemotron 3 tokenizer requires transformers>=5.3.0",
            ),
        ),
        pytest.param(
            "trl-internal-testing/tiny-NemotronHForCausalLM-3.5-lightning",
            id="nemotron_3_5_lightning",
            marks=pytest.mark.skipif(
                Version(transformers.__version__) < Version("5.3.0"),
                reason="Nemotron 3.5 tokenizer requires transformers>=5.3.0",
            ),
        ),
        pytest.param("trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", id="qwen2.5"),
        pytest.param("trl-internal-testing/tiny-Qwen3MoeForCausalLM", id="qwen3"),
        pytest.param(
            ("trl-internal-testing/tiny-Qwen3MoeForCausalLM", qwen3_training_chat_template),
            id="qwen3-training",
        ),
        pytest.param(
            # Same model as `qwen3`, but shipping the response template itself instead of relying on
            # `add_response_schema` to supply it — so `_load` must leave it alone.
            "trl-internal-testing/tiny-Qwen3MoeForCausalLM-ResponseTemplate",
            id="qwen3-response-template",
            marks=pytest.mark.skipif(
                not _SUPPORTS_RESPONSE_TEMPLATE,
                reason="Fixture ships a new-style response template, which requires transformers>=5.13",
            ),
        ),
        pytest.param("trl-internal-testing/tiny-Qwen3ForCausalLM-Instruct-2507", id="qwen3_instruct_2507"),
        pytest.param("trl-internal-testing/tiny-Qwen3VLForConditionalGeneration", id="qwen3_vl", marks=require_vision),
        pytest.param(
            "trl-internal-testing/tiny-Qwen3_5ForConditionalGeneration-NoThink",
            id="qwen35-nothink",
            marks=require_vision,
        ),
        pytest.param(
            "trl-internal-testing/tiny-Qwen3_5ForConditionalGeneration-Think", id="qwen35-think", marks=require_vision
        ),
        pytest.param(
            "trl-internal-testing/tiny-Qwen3_5MoeForConditionalGeneration-3.6", id="qwen36", marks=require_vision
        ),
        pytest.param(
            "trl-internal-testing/tiny-Qwen3_5ForConditionalGeneration-3.8", id="qwen38", marks=require_vision
        ),
        pytest.param(
            "trl-internal-testing/tiny-Gemma4ForConditionalGeneration",
            id="gemma4",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.5.0"),
                    reason="Gemma4 models were introduced in transformers-5.5.0",
                ),
            ],
        ),
        pytest.param(
            "trl-internal-testing/tiny-MuseSparkForConditionalGeneration",
            id="muse_spark",
            marks=[
                require_vision,
                pytest.mark.skipif(
                    Version(transformers.__version__) < Version("5.20.0"),
                    reason="Muse Spark Open was introduced in transformers-5.20.0",
                ),
            ],
        ),
    ],
)
@pytest.mark.xfail(
    condition=Version(transformers.__version__) < Version("5.0.0"),
    reason="Response parsing is not supported in transformers versions below 5.0.0",
    strict=True,
)
@require_response_parsing
class TestParseResponse:
    def _load(self, model_name):
        # A (model, chat template) pair stands for the same tiny model with another form of its chat template
        model_name, chat_template = model_name if isinstance(model_name, tuple) else (model_name, None)
        if "ForCausalLM" in model_name:
            self.is_vlm = False
            processing_class = AutoTokenizer.from_pretrained(model_name)
            tokenizer = processing_class
        elif "ForConditionalGeneration" in model_name:
            self.is_vlm = True
            processing_class = AutoProcessor.from_pretrained(model_name)
            tokenizer = processing_class.tokenizer

        if chat_template is not None:
            processing_class.chat_template = chat_template

        # Nothing to add for models that already ship a new-style `response_template` or a legacy
        # `response_schema`; `add_response_schema` only knows a fixed set of chat templates and raises otherwise.
        has_template = getattr(tokenizer, "response_template", None) is not None
        has_schema = getattr(tokenizer, "response_schema", None) is not None
        if not has_template and not has_schema:
            processing_class = add_response_schema(processing_class)

        return processing_class

    def test_parse_response(self, model_name):
        if model_name in ("trl-internal-testing/tiny-GptOssForCausalLM",) and Version(
            transformers.__version__
        ) < Version("5.5.0"):
            pytest.skip("Upstream bug in response parsing (see #5753; fixed in transformers#45166)")
        processing_class = self._load(model_name)
        messages = [
            {"role": "user", "content": "What is 3*4?"},
            {"role": "assistant", "content": "12"},
        ]
        expected = messages[-1]
        messages = prepare_multimodal_messages(messages) if self.is_vlm else messages
        prefix = processing_class.apply_chat_template(
            messages[:1], add_generation_prompt=True, tokenize=True, return_dict=True
        ).input_ids
        text = processing_class.apply_chat_template(messages, tokenize=True, return_dict=True).input_ids
        if self.is_vlm:  # VLM processors return batched output
            prefix = prefix[0]
            text = text[0]
        response = text[len(prefix) :]
        tokenizer = processing_class.tokenizer if self.is_vlm else processing_class
        parsed = parse_response(tokenizer, response, prefix=prefix)
        assert parsed == expected

    def test_parse_response_with_reasoning_content(self, model_name):
        if model_name in (
            "trl-internal-testing/tiny-Gemma4ForConditionalGeneration",
            "trl-internal-testing/tiny-LlamaForCausalLM-3.1",
            "trl-internal-testing/tiny-LlamaForCausalLM-3.2",
            "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
            "trl-internal-testing/tiny-Qwen3ForCausalLM-Instruct-2507",
            "trl-internal-testing/tiny-Qwen3VLForConditionalGeneration",
        ):
            pytest.skip("This tokenizer doesn't support inline reasoning_content.")
        if model_name == "trl-internal-testing/tiny-GptOssForCausalLM" and not _SUPPORTS_RESPONSE_TEMPLATE:
            # gpt-oss surfaces the analysis channel as `thinking` (separate from `content`) only with the new-style
            # `response_template` parser. The legacy `gptoss_schema` folds the analysis channel into `content`.
            pytest.skip("gpt-oss thinking/content separation requires the response_template parser (>= 5.13).")

        processing_class = self._load(model_name)
        # gpt-oss, Beam and LFM2.5 use the `thinking` field name (matching their own chat templates, which read it back
        # off the message) rather than the `reasoning_content` convention used by other models.
        reasoning_field = (
            "thinking"
            if model_name
            in (
                "trl-internal-testing/tiny-BeamForCausalLM",
                "trl-internal-testing/tiny-GptOssForCausalLM",
                "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5",
                "trl-internal-testing/tiny-Lfm2VlForConditionalGeneration-2.5",
            )
            else "reasoning_content"
        )
        messages = [
            {"role": "user", "content": "What is 3*4?"},
            {"role": "assistant", reasoning_field: "Hmmm.", "content": "12"},
        ]
        expected = messages[-1]
        messages = prepare_multimodal_messages(messages) if self.is_vlm else messages
        # enable_thinking=True is required here because the Qwen3.5 NoThink fixture disables thinking by default
        # for the generation prompt. Only pass it to templates that use it: processors warn otherwise (#7811).
        template_kwargs = {"enable_thinking": True} if "enable_thinking" in processing_class.chat_template else {}
        prefix = processing_class.apply_chat_template(
            messages[:1], add_generation_prompt=True, tokenize=True, return_dict=True, **template_kwargs
        ).input_ids
        text = processing_class.apply_chat_template(messages, tokenize=True, return_dict=True).input_ids
        if self.is_vlm:  # VLM processors return batched output
            prefix = prefix[0]
            text = text[0]
        response = text[len(prefix) :]
        tokenizer = processing_class.tokenizer if self.is_vlm else processing_class
        parsed = parse_response(tokenizer, response, prefix=prefix)
        assert parsed == expected

    def test_parse_response_tool_call(self, model_name):
        if model_name in (
            "trl-internal-testing/tiny-GptOssForCausalLM",
            "trl-internal-testing/tiny-LlamaForCausalLM-3.1",
            "trl-internal-testing/tiny-LlamaForCausalLM-3.2",
        ) and Version(transformers.__version__) < Version("5.5.0"):
            pytest.skip("Upstream bug in response parsing (see #5753; fixed in transformers#45166)")
        processing_class = self._load(model_name)
        tool_calls = [{"type": "function", "function": {"name": "multiply", "arguments": {"a": 3, "b": 4}}}]
        messages = [
            {"role": "user", "content": "What is 3*4?"},
            {
                "role": "assistant",
                # "content" is required here because VLM processors crash on tokenize=True without it
                # (KeyError in processing_utils.py). See huggingface/transformers#45290.
                "content": "",
                "tool_calls": tool_calls,
            },
        ]
        expected = messages[-1]
        messages = prepare_multimodal_messages(messages) if self.is_vlm else messages
        prefix = processing_class.apply_chat_template(
            messages[:1], add_generation_prompt=True, tokenize=True, return_dict=True
        ).input_ids
        text = processing_class.apply_chat_template(messages, tokenize=True, return_dict=True).input_ids
        if self.is_vlm:  # VLM processors return batched output
            prefix = prefix[0]
            text = text[0]
        response = text[len(prefix) :]
        tokenizer = processing_class.tokenizer if self.is_vlm else processing_class
        parsed = parse_response(tokenizer, response, prefix=prefix, tools=[multiply, addition])
        assert parsed == expected

    def test_parse_response_tool_call_with_content(self, model_name):
        if model_name in (
            "trl-internal-testing/tiny-LlamaForCausalLM-3.1",
            "trl-internal-testing/tiny-LlamaForCausalLM-3.2",
        ):
            pytest.skip("Llama 3.1 / 3.2 templates only allow a single tool call per assistant turn, with no content.")
        if model_name == "trl-internal-testing/tiny-GptOssForCausalLM" and not _SUPPORTS_RESPONSE_TEMPLATE:
            # gpt-oss surfaces the content-before-tool-call (analysis channel) as `thinking` only with the new-style
            # `response_template` parser. The legacy `gptoss_schema` folds the analysis channel into `content`.
            pytest.skip("gpt-oss thinking-before-tool-call requires the response_template parser (>= 5.13).")
        processing_class = self._load(model_name)
        tool_calls = [{"type": "function", "function": {"name": "multiply", "arguments": {"a": 3, "b": 4}}}]
        assistant = {"role": "assistant", "tool_calls": tool_calls}
        if model_name in (
            "trl-internal-testing/tiny-BeamForCausalLM",
            "trl-internal-testing/tiny-GptOssForCausalLM",
        ):
            assistant["thinking"] = "Let's call the tool."
            assistant["content"] = ""
        elif model_name == "trl-internal-testing/tiny-MuseGlimmerForConditionalGeneration":
            # Muse Glimmer renders one ATEM channel per assistant turn: a turn carrying tool calls has no `to=user`
            # channel, so plain content is dropped. Text alongside a tool call belongs in the `to=self` channel.
            assistant["reasoning_content"] = "Let's call the tool."
            assistant["content"] = ""
        else:
            assistant["content"] = "Let's call the tool."
        messages = [{"role": "user", "content": "What is 3*4?"}, assistant]
        expected = messages[-1]
        messages = prepare_multimodal_messages(messages) if self.is_vlm else messages
        prefix = processing_class.apply_chat_template(
            messages[:1], add_generation_prompt=True, tokenize=True, return_dict=True
        ).input_ids
        text = processing_class.apply_chat_template(messages, tokenize=True, return_dict=True).input_ids
        if self.is_vlm:  # VLM processors return batched output
            prefix = prefix[0]
            text = text[0]
        response = text[len(prefix) :]
        tokenizer = processing_class.tokenizer if self.is_vlm else processing_class
        parsed = parse_response(tokenizer, response, prefix=prefix, tools=[multiply, addition])
        assert parsed == expected

    def test_parse_response_tool_call_without_arguments(self, model_name):
        if model_name in ("trl-internal-testing/tiny-GptOssForCausalLM",) and Version(
            transformers.__version__
        ) < Version("5.5.0"):
            pytest.skip("Upstream bug in response parsing (see #5753; fixed in transformers#45166)")
        processing_class = self._load(model_name)
        tool_calls = [{"type": "function", "function": {"name": "ping", "arguments": {}}}]
        messages = [
            {"role": "user", "content": "Ping the service."},
            {
                "role": "assistant",
                # "content" is required here because VLM processors crash on tokenize=True without it
                # (KeyError in processing_utils.py). See huggingface/transformers#45290.
                "content": "",
                "tool_calls": tool_calls,
            },
        ]
        expected = messages[-1]
        messages = prepare_multimodal_messages(messages) if self.is_vlm else messages
        prefix = processing_class.apply_chat_template(
            messages[:1], add_generation_prompt=True, tokenize=True, return_dict=True
        ).input_ids
        text = processing_class.apply_chat_template(messages, tokenize=True, return_dict=True).input_ids
        if self.is_vlm:  # VLM processors return batched output
            prefix = prefix[0]
            text = text[0]
        response = text[len(prefix) :]
        tokenizer = processing_class.tokenizer if self.is_vlm else processing_class
        parsed = parse_response(tokenizer, response, prefix=prefix)
        assert parsed == expected

    def test_parse_response_multiple_tool_calls(self, model_name):
        if model_name in (
            "trl-internal-testing/tiny-GptOssForCausalLM",
            "trl-internal-testing/tiny-LlamaForCausalLM-3.1",
            "trl-internal-testing/tiny-LlamaForCausalLM-3.2",
        ):
            pytest.skip("This template only renders one tool call per assistant message.")
        processing_class = self._load(model_name)
        tool_calls = [
            {"type": "function", "function": {"name": "multiply", "arguments": {"a": 3, "b": 4}}},
            {"type": "function", "function": {"name": "addition", "arguments": {"a": 4, "b": 3}}},
        ]
        messages = [
            {"role": "user", "content": "What is 3*4?"},
            {
                "role": "assistant",
                # "content" is required here because VLM processors crash on tokenize=True without it
                # (KeyError in processing_utils.py). See huggingface/transformers#45290.
                "content": "",
                "tool_calls": tool_calls,
            },
        ]
        expected = messages[-1]
        messages = prepare_multimodal_messages(messages) if self.is_vlm else messages
        prefix = processing_class.apply_chat_template(
            messages[:1], add_generation_prompt=True, tokenize=True, return_dict=True
        ).input_ids
        text = processing_class.apply_chat_template(messages, tokenize=True, return_dict=True).input_ids
        if self.is_vlm:  # VLM processors return batched output
            prefix = prefix[0]
            text = text[0]
        response = text[len(prefix) :]
        tokenizer = processing_class.tokenizer if self.is_vlm else processing_class
        parsed = parse_response(tokenizer, response, prefix=prefix, tools=[multiply, addition])
        assert parsed == expected

    @pytest.mark.skipif(
        not _SUPPORTS_PARSE_RESPONSE_TOOLS,
        reason="Typing tool-call arguments with the tool's JSON schema requires transformers>=5.15",
    )
    def test_parse_response_tool_call_argument_types(self, model_name):
        if model_name == "trl-internal-testing/tiny-MuseGlimmerForConditionalGeneration":
            pytest.skip("Muse Glimmer ships its own response template, which JSON-parses `75001` as an int.")
        processing_class = self._load(model_name)
        if model_name == "trl-internal-testing/tiny-MuseSparkForConditionalGeneration":
            pytest.skip("The Muse Spark Open response template reads argument values as JSON, before the tool schema.")
        # A string that looks like a number must stay a string, and every other JSON type must come back typed
        arguments = {"zip_code": "75001", "days": 3, "threshold": 0.5, "metric": True, "fields": ["rain", "wind"]}
        tool_calls = [{"type": "function", "function": {"name": "get_weather", "arguments": arguments}}]
        messages = [
            {"role": "user", "content": "What's the weather in Paris?"},
            {
                "role": "assistant",
                # "content" is required here because VLM processors crash on tokenize=True without it
                # (KeyError in processing_utils.py). See huggingface/transformers#45290.
                "content": "",
                "tool_calls": tool_calls,
            },
        ]
        expected = messages[-1]
        messages = prepare_multimodal_messages(messages) if self.is_vlm else messages
        prefix = processing_class.apply_chat_template(
            messages[:1], tools=[get_weather], add_generation_prompt=True, tokenize=True, return_dict=True
        ).input_ids
        text = processing_class.apply_chat_template(
            messages, tools=[get_weather], tokenize=True, return_dict=True
        ).input_ids
        if self.is_vlm:  # VLM processors return batched output
            prefix = prefix[0]
            text = text[0]
        response = text[len(prefix) :]
        tokenizer = processing_class.tokenizer if self.is_vlm else processing_class
        parsed = parse_response(tokenizer, response, prefix=prefix, tools=[get_weather])
        assert parsed == expected

    def test_parse_response_malformed_tool_call(self, model_name):
        if model_name != "trl-internal-testing/tiny-Qwen3MoeForCausalLM":
            pytest.skip("For simplicity, we only test the malformed tool call case on one tokenizer.")
        processing_class = self._load(model_name)
        text = '<tool_call>\n{"name": "multiply", "arguments": {"a": 3, "b": 4}\n</tool_call><|im_end|>'
        assistant_text = processing_class(text)["input_ids"]
        parsed = parse_response(processing_class, assistant_text)
        expected = {
            "role": "assistant",
            "content": '<tool_call>\n{"name": "multiply", "arguments": {"a": 3, "b": 4}\n</tool_call>',
        }

        assert parsed == expected

    def test_parse_response_truncated(self, model_name):
        processing_class = self._load(model_name)
        # Here we use 2 tool calls as it seems to be a more common source of failure when truncated.
        # Llama 3.1 / 3.2 templates only allow a single tool call per assistant turn, so fall back to one.
        tool_calls = [{"type": "function", "function": {"name": "multiply", "arguments": {"a": 3, "b": 4}}}]
        if model_name not in (
            "trl-internal-testing/tiny-LlamaForCausalLM-3.1",
            "trl-internal-testing/tiny-LlamaForCausalLM-3.2",
        ):
            tool_calls.append({"type": "function", "function": {"name": "addition", "arguments": {"a": 4, "b": 3}}})
        messages = [
            {"role": "user", "content": "What is 3*4?"},
            {"role": "assistant", "content": "", "tool_calls": tool_calls},
        ]
        messages = prepare_multimodal_messages(messages) if self.is_vlm else messages
        prefix = processing_class.apply_chat_template(
            messages[:1], add_generation_prompt=True, tokenize=True, return_dict=True
        ).input_ids
        text = processing_class.apply_chat_template(messages, tokenize=True, return_dict=True).input_ids
        if self.is_vlm:  # VLM processors return batched output
            prefix = prefix[0]
            text = text[0]
        response = text[len(prefix) :]
        tokenizer = processing_class.tokenizer if self.is_vlm else processing_class
        # Truncate the response mid-tool-call and just check that parsing doesn't crash.
        for end in range(1, len(response)):
            parse_response(tokenizer, response[:end])


class TestChatTemplatesDocumentation:
    def _documented(self, path):
        # Templates documented in a `### ` heading, which may list several of them
        headings = [line for line in path.read_text(encoding="utf-8").splitlines() if line.startswith("### ")]
        return set(re.findall(r"`([\w.]+\.jinja)`", "\n".join(headings)))

    def test_every_template_in_readme(self):
        documented = self._documented(_CHAT_TEMPLATES_DIR / "README.md")
        missing = sorted(path.name for path in _CHAT_TEMPLATES_DIR.glob("*.jinja") if path.name not in documented)
        assert missing == []

    def test_every_training_template_in_docs(self):
        documented = self._documented(Path(__file__).parents[1] / "docs" / "source" / "chat_templates.md")
        missing = sorted(
            path.name for path in _CHAT_TEMPLATES_DIR.glob("*_training.jinja") if path.name not in documented
        )
        assert missing == []