import torch.nn as nn

from transformers.models.qwen3_vl.modeling_qwen3_vl import (
    Qwen3VLForConditionalGeneration,
    Qwen3VLModel,
)


# The model was trained with transformers==4.57.1, where
# `Qwen3VLForConditionalGeneration(...).hidden_states[-1]` was the pre-final-norm
# state of the text decoder. In transformers 5.x that field is now the post-norm
# `last_hidden_state`. Replacing the text model's final RMSNorm with a no-op
# restores the representation the model was trained on.
_NORM_KEY_PATTERN = r"^model\.language_model\.norm\.weight$"


class EagerEmbedModel(Qwen3VLModel):
    _keys_to_ignore_on_load_unexpected = [_NORM_KEY_PATTERN]

    def __init__(self, config):
        super().__init__(config)
        self.language_model.norm = nn.Identity()


class EagerEmbedForConditionalGeneration(Qwen3VLForConditionalGeneration):
    _keys_to_ignore_on_load_unexpected = [_NORM_KEY_PATTERN]

    def __init__(self, config):
        super().__init__(config)
        self.model.language_model.norm = nn.Identity()