""" 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"]