modeling_colqwen35_bidirection.py

"""
ColQwen35Bidirection - Inheritance-based Qwen3.5 retrieval model.

Subclasses Qwen3_5ForConditionalGeneration directly and adds a projection
head for ColBERT-style multi-vector retrieval. Supports both causal (default)
and bidirectional attention modes via the `is_causal` parameter.

Key design decisions:
  - `is_causal` is popped in `from_pretrained` and NEVER forwarded to
    `super().from_pretrained(...)`. This prevents HuggingFace's
    PretrainedConfig catch-all from injecting `config.is_causal` onto
    the model config, which would alter Qwen 3.5's mask resolver path.
    With `is_causal=True`, the load path is bit-equivalent to
    ColQwen35Default (no config attribute mutation).
  - Bidirectional patches (`_configure_bidirectional_attention` on the
    config, `_patch_attention_is_causal` on attention modules) are
    applied around the parent load in `from_pretrained`, not forwarded
    through kwargs. The `__init__` retains idempotent guards for direct
    construction (e.g. tests).
  - No `self.is_causal` attribute is set before `super().__init__()`,
    avoiding shadowing of inherited attributes.
"""

from torch import nn
from transformers.models.qwen3_5 import Qwen3_5ForConditionalGeneration, Qwen3_5Config
from transformers.utils import is_torch_available, logging

if is_torch_available():
    import torch

logger = logging.get_logger(__name__)


class ColQwen35Bidirection(Qwen3_5ForConditionalGeneration):
    """
    Qwen 3.5 VLM adapted to return ColBERT-style token embeddings.

    Supports bidirectional attention (is_causal=False) for the 8 full
    attention layers while leaving the 24 GatedDeltaNet layers causal.
    When is_causal=True (default), the model uses Qwen 3.5's native
    attention without any modifications.
    """

    _checkpoint_conversion_mapping: dict[str, str] = {
        r"^base_model\.model\.embedding_proj_layer": "embedding_proj_layer",
    }

    def __init__(
        self,
        config: Qwen3_5Config,
        embedding_dim: int | None = None,
        is_causal: bool | None = None,
    ):
        if embedding_dim is None:
            embedding_dim = getattr(config, "embedding_dim", 128)
        if is_causal is None:
            is_causal = getattr(config, "is_causal", True)

        if not is_causal:
            self._configure_bidirectional_attention(config)

        super().__init__(config)

        if not is_causal:
            self._patch_attention_is_causal()

        self._is_causal_mode = is_causal

        hidden_size = getattr(self.config, "hidden_size", None)
        if hidden_size is None and hasattr(self.config, "text_config"):
            hidden_size = getattr(self.config.text_config, "hidden_size", None)
        if hidden_size is None:
            raise ValueError(
                f"Unable to determine text hidden size for {type(self.config).__name__}."
            )

        self.embedding_dim = embedding_dim
        self.config.embedding_dim = embedding_dim
        self.embedding_proj_layer = nn.Linear(hidden_size, self.embedding_dim)
        self.padding_side = "left"

        self.post_init()
        if getattr(self.config, "lm_head_removed", False):
            self.remove_lm_head()

    @staticmethod
    def _configure_bidirectional_attention(config: Qwen3_5Config) -> None:
        """
        Patch config for bidirectional mask generation.

        Sets is_causal=False on the config so that create_causal_mask()
        produces a padding-only mask instead of a lower-triangular causal mask.
        Must be called before super().__init__().
        """
        text_config = config.get_text_config()
        text_config.is_causal = False
        config.is_causal = False

    def _patch_attention_is_causal(self) -> None:
        """
        Patch is_causal on every Qwen3_5Attention module.

        Upstream Qwen3_5Attention.__init__ hardcodes self.is_causal = True.
        This flag is passed to Flash Attention and other attention backends,
        so it must be False for bidirectional behavior.
        """
        patched = 0
        for module in self.modules():
            if type(module).__name__ == "Qwen3_5Attention":
                module.is_causal = False
                patched += 1
        logger.info("Patched is_causal=False on %d Qwen3_5Attention modules", patched)

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        is_causal = kwargs.pop("is_causal", None)
        key_mapping = kwargs.pop("key_mapping", None)
        if key_mapping is None:
            key_mapping = dict(getattr(super(), "_checkpoint_conversion_mapping", {}))
            key_mapping.update(cls._checkpoint_conversion_mapping)

        if not is_causal:
            from transformers import AutoConfig
            config = kwargs.get("config")
            if config is None:
                path = args[0] if args else kwargs.get("pretrained_model_name_or_path")
                config = AutoConfig.from_pretrained(path)
                kwargs["config"] = config
            cls._configure_bidirectional_attention(config)

        instance = super().from_pretrained(
            *args, **kwargs, key_mapping=key_mapping
        )

        resolved_is_causal = (
            getattr(instance.config, "is_causal", True)
            if is_causal is None
            else is_causal
        )
        if not resolved_is_causal:
            instance._patch_attention_is_causal()

        instance._is_causal_mode = resolved_is_causal
        return instance

    def remove_lm_head(self) -> dict[str, int]:
        """Permanently remove the unused language-model output projection.

        The retrieval forward path calls ``self.model`` directly and consumes
        hidden states, so ``lm_head`` is never used to produce embeddings.
        Persisting ``lm_head_removed`` in the config ensures exported
        checkpoints reload with ``nn.Identity`` instead of recreating a large,
        randomly initialized output projection.
        """
        lm_head = getattr(self, "lm_head", None)
        if lm_head is None:
            raise AttributeError("Model has no lm_head attribute")
        if isinstance(lm_head, nn.Identity):
            return {"parameters": 0, "bytes": 0}

        parameters = {id(parameter): parameter for parameter in lm_head.parameters()}
        shared_elsewhere = {
            id(parameter)
            for name, parameter in self.named_parameters(remove_duplicate=False)
            if not name.startswith("lm_head.")
        }
        removed = {
            parameter_id: parameter
            for parameter_id, parameter in parameters.items()
            if parameter_id not in shared_elsewhere
        }
        removed_parameters = sum(parameter.numel() for parameter in removed.values())
        removed_bytes = sum(
            parameter.numel() * parameter.element_size()
            for parameter in removed.values()
        )

        self.lm_head = nn.Identity()
        self.config.lm_head_removed = True
        self.config.tie_word_embeddings = False

        tied_keys = getattr(self, "_tied_weights_keys", None)
        if isinstance(tied_keys, dict):
            self._tied_weights_keys = {
                key: value
                for key, value in tied_keys.items()
                if "lm_head" not in key and "lm_head" not in value
            }
        elif isinstance(tied_keys, (list, tuple, set)):
            self._tied_weights_keys = [
                key for key in tied_keys if "lm_head" not in key
            ]

        return {"parameters": removed_parameters, "bytes": removed_bytes}

    def forward(self, *args, **kwargs) -> torch.Tensor:
        inner = getattr(self.model, "model", self.model)
        if hasattr(inner, "rope_deltas"):
            inner.rope_deltas = None

        kwargs.pop("return_dict", None)
        kwargs.pop("output_hidden_states", None)
        kwargs.pop("use_cache", None)

        use_cache = not self.training

        last_hidden_state = (
            self.model.forward(
                *args, **kwargs,
                output_hidden_states=False,
                return_dict=True,
                use_cache=use_cache,
            )
            .last_hidden_state
        )

        proj_dtype = self.embedding_proj_layer.weight.dtype
        embeddings = self.embedding_proj_layer(last_hidden_state.to(proj_dtype))

        embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True).clamp(min=1e-6)
        embeddings = embeddings * kwargs["attention_mask"].unsqueeze(-1)

        return embeddings


__all__ = ["ColQwen35Bidirection"]