webAI-Official/webAI-ColVec1.1-4b

1"""2ColQwen35Bidirection - Inheritance-based Qwen3.5 retrieval model.3 4Subclasses Qwen3_5ForConditionalGeneration directly and adds a projection5head for ColBERT-style multi-vector retrieval. Supports both causal (default)6and bidirectional attention modes via the `is_causal` parameter.7 8Key design decisions:9  - `is_causal` is popped in `from_pretrained` and NEVER forwarded to10    `super().from_pretrained(...)`. This prevents HuggingFace's11    PretrainedConfig catch-all from injecting `config.is_causal` onto12    the model config, which would alter Qwen 3.5's mask resolver path.13    With `is_causal=True`, the load path is bit-equivalent to14    ColQwen35Default (no config attribute mutation).15  - Bidirectional patches (`_configure_bidirectional_attention` on the16    config, `_patch_attention_is_causal` on attention modules) are17    applied around the parent load in `from_pretrained`, not forwarded18    through kwargs. The `__init__` retains idempotent guards for direct19    construction (e.g. tests).20  - No `self.is_causal` attribute is set before `super().__init__()`,21    avoiding shadowing of inherited attributes.22"""23 24from torch import nn25from transformers.models.qwen3_5 import Qwen3_5ForConditionalGeneration, Qwen3_5Config26from transformers.utils import is_torch_available, logging27 28if is_torch_available():29    import torch30 31logger = logging.get_logger(__name__)32 33 34class ColQwen35Bidirection(Qwen3_5ForConditionalGeneration):35    """36    Qwen 3.5 VLM adapted to return ColBERT-style token embeddings.37 38    Supports bidirectional attention (is_causal=False) for the 8 full39    attention layers while leaving the 24 GatedDeltaNet layers causal.40    When is_causal=True (default), the model uses Qwen 3.5's native41    attention without any modifications.42    """43 44    _checkpoint_conversion_mapping: dict[str, str] = {45        r"^base_model\.model\.embedding_proj_layer": "embedding_proj_layer",46    }47 48    def __init__(49        self,50        config: Qwen3_5Config,51        embedding_dim: int | None = None,52        is_causal: bool | None = None,53    ):54        if embedding_dim is None:55            embedding_dim = getattr(config, "embedding_dim", 128)56        if is_causal is None:57            is_causal = getattr(config, "is_causal", True)58 59        if not is_causal:60            self._configure_bidirectional_attention(config)61 62        super().__init__(config)63 64        if not is_causal:65            self._patch_attention_is_causal()66 67        self._is_causal_mode = is_causal68 69        hidden_size = getattr(self.config, "hidden_size", None)70        if hidden_size is None and hasattr(self.config, "text_config"):71            hidden_size = getattr(self.config.text_config, "hidden_size", None)72        if hidden_size is None:73            raise ValueError(74                f"Unable to determine text hidden size for {type(self.config).__name__}."75            )76 77        self.embedding_dim = embedding_dim78        self.config.embedding_dim = embedding_dim79        self.embedding_proj_layer = nn.Linear(hidden_size, self.embedding_dim)80        self.padding_side = "left"81 82        self.post_init()83        if getattr(self.config, "lm_head_removed", False):84            self.remove_lm_head()85 86    @staticmethod87    def _configure_bidirectional_attention(config: Qwen3_5Config) -> None:88        """89        Patch config for bidirectional mask generation.90 91        Sets is_causal=False on the config so that create_causal_mask()92        produces a padding-only mask instead of a lower-triangular causal mask.93        Must be called before super().__init__().94        """95        text_config = config.get_text_config()96        text_config.is_causal = False97        config.is_causal = False98 99    def _patch_attention_is_causal(self) -> None:100        """101        Patch is_causal on every Qwen3_5Attention module.102 103        Upstream Qwen3_5Attention.__init__ hardcodes self.is_causal = True.104        This flag is passed to Flash Attention and other attention backends,105        so it must be False for bidirectional behavior.106        """107        patched = 0108        for module in self.modules():109            if type(module).__name__ == "Qwen3_5Attention":110                module.is_causal = False111                patched += 1112        logger.info("Patched is_causal=False on %d Qwen3_5Attention modules", patched)113 114    @classmethod115    def from_pretrained(cls, *args, **kwargs):116        is_causal = kwargs.pop("is_causal", None)117        key_mapping = kwargs.pop("key_mapping", None)118        if key_mapping is None:119            key_mapping = dict(getattr(super(), "_checkpoint_conversion_mapping", {}))120            key_mapping.update(cls._checkpoint_conversion_mapping)121 122        if not is_causal:123            from transformers import AutoConfig124            config = kwargs.get("config")125            if config is None:126                path = args[0] if args else kwargs.get("pretrained_model_name_or_path")127                config = AutoConfig.from_pretrained(path)128                kwargs["config"] = config129            cls._configure_bidirectional_attention(config)130 131        instance = super().from_pretrained(132            *args, **kwargs, key_mapping=key_mapping133        )134 135        resolved_is_causal = (136            getattr(instance.config, "is_causal", True)137            if is_causal is None138            else is_causal139        )140        if not resolved_is_causal:141            instance._patch_attention_is_causal()142 143        instance._is_causal_mode = resolved_is_causal144        return instance145 146    def remove_lm_head(self) -> dict[str, int]:147        """Permanently remove the unused language-model output projection.148 149        The retrieval forward path calls ``self.model`` directly and consumes150        hidden states, so ``lm_head`` is never used to produce embeddings.151        Persisting ``lm_head_removed`` in the config ensures exported152        checkpoints reload with ``nn.Identity`` instead of recreating a large,153        randomly initialized output projection.154        """155        lm_head = getattr(self, "lm_head", None)156        if lm_head is None:157            raise AttributeError("Model has no lm_head attribute")158        if isinstance(lm_head, nn.Identity):159            return {"parameters": 0, "bytes": 0}160 161        parameters = {id(parameter): parameter for parameter in lm_head.parameters()}162        shared_elsewhere = {163            id(parameter)164            for name, parameter in self.named_parameters(remove_duplicate=False)165            if not name.startswith("lm_head.")166        }167        removed = {168            parameter_id: parameter169            for parameter_id, parameter in parameters.items()170            if parameter_id not in shared_elsewhere171        }172        removed_parameters = sum(parameter.numel() for parameter in removed.values())173        removed_bytes = sum(174            parameter.numel() * parameter.element_size()175            for parameter in removed.values()176        )177 178        self.lm_head = nn.Identity()179        self.config.lm_head_removed = True180        self.config.tie_word_embeddings = False181 182        tied_keys = getattr(self, "_tied_weights_keys", None)183        if isinstance(tied_keys, dict):184            self._tied_weights_keys = {185                key: value186                for key, value in tied_keys.items()187                if "lm_head" not in key and "lm_head" not in value188            }189        elif isinstance(tied_keys, (list, tuple, set)):190            self._tied_weights_keys = [191                key for key in tied_keys if "lm_head" not in key192            ]193 194        return {"parameters": removed_parameters, "bytes": removed_bytes}195 196    def forward(self, *args, **kwargs) -> torch.Tensor:197        inner = getattr(self.model, "model", self.model)198        if hasattr(inner, "rope_deltas"):199            inner.rope_deltas = None200 201        kwargs.pop("return_dict", None)202        kwargs.pop("output_hidden_states", None)203        kwargs.pop("use_cache", None)204 205        use_cache = not self.training206 207        last_hidden_state = (208            self.model.forward(209                *args, **kwargs,210                output_hidden_states=False,211                return_dict=True,212                use_cache=use_cache,213            )214            .last_hidden_state215        )216 217        proj_dtype = self.embedding_proj_layer.weight.dtype218        embeddings = self.embedding_proj_layer(last_hidden_state.to(proj_dtype))219 220        embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True).clamp(min=1e-6)221        embeddings = embeddings * kwargs["attention_mask"].unsqueeze(-1)222 223        return embeddings224 225 226__all__ = ["ColQwen35Bidirection"]227