File size: 8,590 Bytes
e8491df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 | """
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"]
|