from typing import Any

from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig


class Qwen2_5_VLCAConfig(Qwen2_5_VLConfig):
    """Qwen2.5-VL config augmented with cross-attention options"""

    model_type = "qwen2_5vl_ca"

    def __init__(
        self,
        *args: Any,
        # Our fusion mechanism
        xa_layers: None | tuple = None,
        cross_attention: bool = False,
        xa_share_qkvo: bool = False,
        **kwargs: Any,
    ):
        super().__init__(*args, **kwargs)
        self.head_dim = self.hidden_size // self.num_attention_heads
        # Our fusion mechanisms
        self.xa_layers = xa_layers
        self.cross_attention = cross_attention
        self.xa_share_qkvo = xa_share_qkvo
        # Derived from vision_start/end_token_id set by parent
        self.pre_image_tokens = (
            [self.vision_start_token_id]
            if getattr(self, "vision_start_token_id", None) is not None
            else []
        )
        self.post_image_tokens = (
            [self.vision_end_token_id]
            if getattr(self, "vision_end_token_id", None) is not None
            else []
        )