# Copyright 2020-2025 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 gc

import pytest
import torch
import transformers
from packaging import version
from parameterized import parameterized
from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, GenerationConfig

from trl import AutoModelForCausalLMWithValueHead, AutoModelForSeq2SeqLMWithValueHead, create_reference_model

from .testing_utils import TrlTestCase


ALL_CAUSAL_LM_MODELS = [
    "trl-internal-testing/tiny-BloomForCausalLM",
    "trl-internal-testing/tiny-CohereForCausalLM",
    "trl-internal-testing/tiny-DbrxForCausalLM",
    # "trl-internal-testing/tiny-FalconMambaForCausalLM",  # FalconMambaForCausalLM modeling seems to be broken for now
    "trl-internal-testing/tiny-Gemma2ForCausalLM",
    "trl-internal-testing/tiny-GemmaForCausalLM",
    "trl-internal-testing/tiny-GPT2LMHeadModel",
    "trl-internal-testing/tiny-GPTNeoXForCausalLM",
    "trl-internal-testing/tiny-LlamaForCausalLM-3.1",
    "trl-internal-testing/tiny-LlamaForCausalLM-3.2",
    "trl-internal-testing/tiny-LlamaForCausalLM-3",
    "trl-internal-testing/tiny-MistralForCausalLM-0.1",
    "trl-internal-testing/tiny-MistralForCausalLM-0.2",
    "trl-internal-testing/tiny-OPTForCausalLM",
    "trl-internal-testing/tiny-Phi3ForCausalLM",
    "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
]

ALL_SEQ2SEQ_MODELS = [
    "trl-internal-testing/tiny-T5ForConditionalGeneration",
    "trl-internal-testing/tiny-BartModel",
]


class BaseTester:
    class VHeadModelTester(TrlTestCase):
        all_model_names = None
        trl_model_class = None
        transformers_model_class = None

        def setup_method(self):
            self.device = "cuda" if torch.cuda.is_available() else "cpu"

        def test_value_head(self):
            r"""
            Test if the v-head is added to the model successfully
            """
            for model_name in self.all_model_names:
                if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
                    transformers.__version__
                ) < version.parse("4.58.0.dev0"):
                    # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
                    continue

                model = self.trl_model_class.from_pretrained(model_name)
                assert hasattr(model, "v_head")

        def test_value_head_shape(self):
            r"""
            Test if the v-head has the correct shape
            """
            for model_name in self.all_model_names:
                if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
                    transformers.__version__
                ) < version.parse("4.58.0.dev0"):
                    # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
                    continue

                model = self.trl_model_class.from_pretrained(model_name)
                assert model.v_head.summary.weight.shape[0] == 1

        def test_value_head_init_random(self):
            r"""
            Test if the v-head has been randomly initialized. We can check that by making sure the bias is different
            than zeros by default.
            """
            for model_name in self.all_model_names:
                if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
                    transformers.__version__
                ) < version.parse("4.58.0.dev0"):
                    # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
                    continue

                model = self.trl_model_class.from_pretrained(model_name)
                assert not torch.allclose(model.v_head.summary.bias, torch.zeros_like(model.v_head.summary.bias))

        def test_value_head_not_str(self):
            r"""
            Test if the v-head is added to the model successfully, by passing a non `PretrainedModel` as an argument to
            `from_pretrained`.
            """
            for model_name in self.all_model_names:
                if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
                    transformers.__version__
                ) < version.parse("4.58.0.dev0"):
                    # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
                    continue

                pretrained_model = self.transformers_model_class.from_pretrained(model_name)
                model = self.trl_model_class.from_pretrained(pretrained_model)
                assert hasattr(model, "v_head")

        def test_from_save_trl(self):
            """
            Test if the model can be saved and loaded from a directory and get the same weights, including the
            additional modules (e.g. v_head)
            """
            for model_name in self.all_model_names:
                if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
                    transformers.__version__
                ) < version.parse("4.58.0.dev0"):
                    # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
                    continue

                model = self.trl_model_class.from_pretrained(model_name)

                model.save_pretrained(self.tmp_dir)

                model_from_save = self.trl_model_class.from_pretrained(self.tmp_dir)

                # Check if the weights are the same
                for key in model_from_save.state_dict():
                    assert torch.allclose(model_from_save.state_dict()[key], model.state_dict()[key])

        def test_from_save_trl_sharded(self):
            """
            Test if the model can be saved and loaded from a directory and get the same weights - sharded case
            """
            for model_name in self.all_model_names:
                if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
                    transformers.__version__
                ) < version.parse("4.58.0.dev0"):
                    # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
                    continue

                model = self.trl_model_class.from_pretrained(model_name)

                model.save_pretrained(self.tmp_dir)

                model_from_save = self.trl_model_class.from_pretrained(self.tmp_dir)

                # Check if the weights are the same
                for key in model_from_save.state_dict():
                    assert torch.allclose(model_from_save.state_dict()[key], model.state_dict()[key])

        def test_from_save_transformers_sharded(self):
            """
            Test if the model can be saved and loaded using transformers and get the same weights - sharded case
            """
            for model_name in self.all_model_names:
                if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
                    transformers.__version__
                ) < version.parse("4.58.0.dev0"):
                    # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
                    continue

                transformers_model = self.trl_model_class.transformers_parent_class.from_pretrained(model_name)

                trl_model = self.trl_model_class.from_pretrained(model_name)

                trl_model.save_pretrained(self.tmp_dir, max_shard_size="1MB")
                transformers_model_from_save = self.trl_model_class.transformers_parent_class.from_pretrained(
                    self.tmp_dir
                )

                # Check if the weights are the same
                for key in transformers_model.state_dict():
                    assert torch.allclose(
                        transformers_model_from_save.state_dict()[key], transformers_model.state_dict()[key]
                    )

        def test_from_save_transformers(self):
            """
            Test if the model can be saved and loaded using transformers and get the same weights. We override the test
            of the super class to check if the weights are the same.
            """
            for model_name in self.all_model_names:
                if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
                    transformers.__version__
                ) < version.parse("4.58.0.dev0"):
                    # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
                    continue

                transformers_model = self.trl_model_class.transformers_parent_class.from_pretrained(model_name)

                trl_model = self.trl_model_class.from_pretrained(model_name)

                trl_model.save_pretrained(self.tmp_dir)
                transformers_model_from_save = self.trl_model_class.transformers_parent_class.from_pretrained(
                    self.tmp_dir
                )

                # Check if the weights are the same
                for key in transformers_model.state_dict():
                    assert torch.allclose(
                        transformers_model_from_save.state_dict()[key], transformers_model.state_dict()[key]
                    )

                # Check if the trl model has the same keys as the transformers model
                # except the v_head
                for key in trl_model.state_dict():
                    if "v_head" not in key:
                        assert key in transformers_model.state_dict()
                        # check if the weights are the same
                        assert torch.allclose(trl_model.state_dict()[key], transformers_model.state_dict()[key])

                # check if they have the same modules
                assert set(transformers_model_from_save.state_dict().keys()) == set(
                    transformers_model.state_dict().keys()
                )


class TestCausalLMValueHeadModel(BaseTester.VHeadModelTester, TrlTestCase):
    """
    Testing suite for v-head models.
    """

    all_model_names = ALL_CAUSAL_LM_MODELS
    trl_model_class = AutoModelForCausalLMWithValueHead
    transformers_model_class = AutoModelForCausalLM

    def teardown_method(self):
        # free memory
        gc.collect()

    def test_inference(self):
        r"""
        Test if the model can be used for inference and outputs 3 values
        - logits, loss, and value states
        """
        EXPECTED_OUTPUT_SIZE = 3

        for model_name in self.all_model_names:
            if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
                transformers.__version__
            ) < version.parse("4.58.0.dev0"):
                # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
                continue

            model = self.trl_model_class.from_pretrained(model_name).to(self.device)
            input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
            outputs = model(input_ids)

            # Check if the outputs are of the right size - here
            # we always output 3 values - logits, loss, and value states
            assert len(outputs) == EXPECTED_OUTPUT_SIZE

    def test_dropout_config(self):
        r"""
        Test if we instantiate a model by adding `summary_drop_prob` to the config it will be added to the v_head
        """
        for model_name in self.all_model_names:
            if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
                transformers.__version__
            ) < version.parse("4.58.0.dev0"):
                # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
                continue

            pretrained_model = self.transformers_model_class.from_pretrained(model_name)
            pretrained_model.config.summary_dropout_prob = 0.5
            model = self.trl_model_class.from_pretrained(pretrained_model)

            # Check if v head of the model has the same dropout as the config
            assert model.v_head.dropout.p == pretrained_model.config.summary_dropout_prob

    def test_dropout_kwargs(self):
        r"""
        Test if we instantiate a model by adding `summary_drop_prob` to the config it will be added to the v_head
        """
        for model_name in self.all_model_names:
            if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
                transformers.__version__
            ) < version.parse("4.58.0.dev0"):
                # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
                continue
            v_head_kwargs = {"summary_dropout_prob": 0.5}

            model = self.trl_model_class.from_pretrained(model_name, **v_head_kwargs)

            # Check if v head of the model has the same dropout as the config
            assert model.v_head.dropout.p == 0.5

            model = self.trl_model_class.from_pretrained(model_name, summary_dropout_prob=0.5)

            # Check if v head of the model has the same dropout as the config
            assert model.v_head.dropout.p == 0.5

    @parameterized.expand(ALL_CAUSAL_LM_MODELS)
    def test_generate(self, model_name):
        r"""
        Test if `generate` works for every model
        """
        if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
            transformers.__version__
        ) < version.parse("4.58.0.dev0"):
            # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
            pytest.xfail("DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version")

        generation_config = GenerationConfig(max_new_tokens=9)
        model = self.trl_model_class.from_pretrained(model_name).to(self.device)
        input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)

        # Just check if the generation works
        _ = model.generate(input_ids, generation_config=generation_config)

    def test_transformers_bf16_kwargs(self):
        r"""
        Test if the transformers kwargs are correctly passed. Here we check that loading a model in half precision
        works as expected, i.e. the weights of the `pretrained_model` attribute is loaded in half precision and you can
        run a dummy forward pass without any issue.
        """
        for model_name in self.all_model_names:
            if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
                transformers.__version__
            ) < version.parse("4.58.0.dev0"):
                # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
                continue

            trl_model = self.trl_model_class.from_pretrained(model_name, dtype=torch.bfloat16).to(self.device)

            lm_head_namings = ["lm_head", "embed_out", "output_layer"]

            assert any(hasattr(trl_model.pretrained_model, lm_head_naming) for lm_head_naming in lm_head_namings), (
                "Can't test the model because it doesn't have any of the expected lm_head namings"
            )

            for lm_head_naming in lm_head_namings:
                if hasattr(trl_model.pretrained_model, lm_head_naming):
                    assert getattr(trl_model.pretrained_model, lm_head_naming).weight.dtype == torch.bfloat16

            dummy_input = torch.LongTensor([[0, 1, 0, 1]]).to(self.device)

            # check dummy forward pass works in half precision
            _ = trl_model(dummy_input)

    @pytest.mark.skip(reason="This test needs to be run manually due to HF token issue.")
    def test_push_to_hub(self):
        for model_name in self.all_model_names:
            if model_name == "trl-internal-testing/tiny-DbrxForCausalLM" and version.parse(
                transformers.__version__
            ) < version.parse("4.58.0.dev0"):
                # DbrxConfig generated after 4.58.0 isn't compatible with modeling code before this version
                continue

            model = AutoModelForCausalLMWithValueHead.from_pretrained(model_name)
            if "sharded" in model_name:
                model.push_to_hub(model_name + "-ppo", use_auth_token=True, max_shard_size="1MB")
            else:
                model.push_to_hub(model_name + "-ppo", use_auth_token=True)

            model_from_pretrained = AutoModelForCausalLMWithValueHead.from_pretrained(model_name + "-ppo")
            # check all keys
            assert model.state_dict().keys() == model_from_pretrained.state_dict().keys()

            for name, param in model.state_dict().items():
                assert torch.allclose(param, model_from_pretrained.state_dict()[name]), (
                    f"Parameter {name} is not the same after push_to_hub and from_pretrained"
                )


class TestSeq2SeqValueHeadModel(BaseTester.VHeadModelTester, TrlTestCase):
    """
    Testing suite for v-head models.
    """

    all_model_names = ALL_SEQ2SEQ_MODELS
    trl_model_class = AutoModelForSeq2SeqLMWithValueHead
    transformers_model_class = AutoModelForSeq2SeqLM

    def teardown_method(self):
        # free memory
        gc.collect()

    def test_inference(self):
        r"""
        Test if the model can be used for inference and outputs 3 values
        - logits, loss, and value states
        """
        EXPECTED_OUTPUT_SIZE = 3

        for model_name in self.all_model_names:
            model = self.trl_model_class.from_pretrained(model_name).to(self.device)
            input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
            decoder_input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
            outputs = model(input_ids, decoder_input_ids=decoder_input_ids)

            # Check if the outputs are of the right size - here
            # we always output 3 values - logits, loss, and value states
            assert len(outputs) == EXPECTED_OUTPUT_SIZE

    def test_dropout_config(self):
        r"""
        Test if we instantiate a model by adding `summary_drop_prob` to the config it will be added to the v_head
        """
        for model_name in self.all_model_names:
            pretrained_model = self.transformers_model_class.from_pretrained(model_name)
            pretrained_model.config.summary_dropout_prob = 0.5
            model = self.trl_model_class.from_pretrained(pretrained_model)

            # Check if v head of the model has the same dropout as the config
            assert model.v_head.dropout.p == pretrained_model.config.summary_dropout_prob

    def test_dropout_kwargs(self):
        r"""
        Test if we instantiate a model by adding `summary_drop_prob` to the config it will be added to the v_head
        """
        for model_name in self.all_model_names:
            v_head_kwargs = {"summary_dropout_prob": 0.5}

            model = self.trl_model_class.from_pretrained(model_name, **v_head_kwargs)

            # Check if v head of the model has the same dropout as the config
            assert model.v_head.dropout.p == 0.5

            model = self.trl_model_class.from_pretrained(model_name, summary_dropout_prob=0.5)

            # Check if v head of the model has the same dropout as the config
            assert model.v_head.dropout.p == 0.5

    @parameterized.expand(ALL_SEQ2SEQ_MODELS)
    def test_generate(self, model_name):
        r"""
        Test if `generate` works for every model
        """
        generation_config = GenerationConfig(max_new_tokens=9)
        model = self.trl_model_class.from_pretrained(model_name).to(self.device)
        input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)
        decoder_input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], device=self.device)

        # Just check if the generation works
        _ = model.generate(input_ids, decoder_input_ids=decoder_input_ids, generation_config=generation_config)

    @pytest.mark.skip(reason="This test needs to be run manually due to HF token issue.")
    def test_push_to_hub(self):
        for model_name in self.all_model_names:
            model = self.trl_model_class.from_pretrained(model_name)
            if "sharded" in model_name:
                model.push_to_hub(model_name + "-ppo", use_auth_token=True, max_shard_size="1MB")
            else:
                model.push_to_hub(model_name + "-ppo", use_auth_token=True)

            model_from_pretrained = self.trl_model_class.from_pretrained(model_name + "-ppo")
            # check all keys
            assert model.state_dict().keys() == model_from_pretrained.state_dict().keys()

            for name, param in model.state_dict().items():
                assert torch.allclose(param, model_from_pretrained.state_dict()[name]), (
                    f"Parameter {name} is not the same after push_to_hub and from_pretrained"
                )

    def test_transformers_bf16_kwargs(self):
        r"""
        Test if the transformers kwargs are correctly passed. Here we check that loading a model in half precision
        works as expected, i.e. the weights of the `pretrained_model` attribute is loaded in half precision and you can
        run a dummy forward pass without any issue.
        """
        for model_name in self.all_model_names:
            trl_model = self.trl_model_class.from_pretrained(model_name, dtype=torch.bfloat16).to(self.device)

            lm_head_namings = self.trl_model_class.lm_head_namings

            assert any(hasattr(trl_model.pretrained_model, lm_head_naming) for lm_head_naming in lm_head_namings)

            for lm_head_naming in lm_head_namings:
                if hasattr(trl_model.pretrained_model, lm_head_naming):
                    assert getattr(trl_model.pretrained_model, lm_head_naming).weight.dtype == torch.bfloat16

            dummy_input = torch.LongTensor([[0, 1, 0, 1]]).to(self.device)

            # check dummy forward pass works in half precision
            _ = trl_model(input_ids=dummy_input, decoder_input_ids=dummy_input)


class TestReferenceModel(TrlTestCase):
    def setup_method(self):
        self.model = AutoModelForCausalLMWithValueHead.from_pretrained("trl-internal-testing/tiny-GPT2LMHeadModel")
        self.test_input = torch.tensor([[0, 1, 2, 3]])
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1)
        self.layer_format = "pretrained_model.transformer.h.{layer}.attn.c_attn.weight"

    def test_independent_reference(self):
        layer_0 = self.layer_format.format(layer=0)
        layer_1 = self.layer_format.format(layer=1)

        ref_model = create_reference_model(self.model)

        first_layer_before = self.model.get_parameter(layer_0).data.clone()
        last_layer_before = self.model.get_parameter(layer_1).data.clone()  # the model only has 2 layers

        first_ref_layer_before = ref_model.get_parameter(layer_0).data.clone()
        last_ref_layer_before = ref_model.get_parameter(layer_1).data.clone()

        output = self.model(input_ids=self.test_input, labels=self.test_input)
        output[1].backward()
        self.optimizer.step()

        first_layer_after = self.model.get_parameter(layer_0).data.clone()
        last_layer_after = self.model.get_parameter(layer_1).data.clone()

        first_ref_layer_after = ref_model.get_parameter(layer_0).data.clone()
        last_ref_layer_after = ref_model.get_parameter(layer_1).data.clone()

        # before optimization ref and model are identical
        assert (first_layer_before == first_ref_layer_before).all()
        assert (last_layer_before == last_ref_layer_before).all()

        # ref model stays identical after optimization
        assert (first_ref_layer_before == first_ref_layer_after).all()
        assert (last_ref_layer_before == last_ref_layer_after).all()

        # optimized model changes
        assert not (first_layer_before == first_layer_after).all()
        assert not (last_layer_before == last_layer_after).all()

    def test_shared_layers(self):
        layer_0 = self.layer_format.format(layer=0)
        layer_1 = self.layer_format.format(layer=1)

        ref_model = create_reference_model(self.model, num_shared_layers=1)

        first_layer_before = self.model.get_parameter(layer_0).data.clone()
        second_layer_before = self.model.get_parameter(layer_1).data.clone()

        first_ref_layer_before = ref_model.get_parameter(layer_0).data.clone()
        second_ref_layer_before = ref_model.get_parameter(layer_1).data.clone()

        output = self.model(input_ids=self.test_input, labels=self.test_input)
        output[1].backward()
        self.optimizer.step()

        first_layer_after = self.model.get_parameter(layer_0).data.clone()
        second_layer_after = self.model.get_parameter(layer_1).data.clone()

        first_ref_layer_after = ref_model.get_parameter(layer_0).data.clone()
        second_ref_layer_after = ref_model.get_parameter(layer_1).data.clone()

        # before optimization ref and model are identical
        assert (first_layer_before == first_ref_layer_before).all()
        assert (second_layer_before == second_ref_layer_before).all()

        # ref model stays identical after optimization
        assert (first_ref_layer_before == first_ref_layer_after).all()
        assert (second_ref_layer_before == second_ref_layer_after).all()

        # first layer of optimized model stays the same
        assert (first_layer_before == first_layer_after).all()

        # other layers in optimized model change
        assert not (second_layer_before == second_layer_after).all()
