File size: 1,205 Bytes
2ada2c7 | 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 | """Bidirectional attention for ColQwen3.5 retrieval.
Released `colpali-engine` (through 0.3.17) builds ColQwen3.5 with the causal
Qwen3.5 masks it inherits from the generative backbone. EVIE was trained and
evaluated with the full-attention layers encoder-ized, so the checkpoint must be
switched before it reproduces the reported scores.
Qwen3.5 interleaves GatedDeltaNet (`linear_attention`) and `full_attention`
layers. Only the full-attention layers are flipped here; the recurrent layers
are order-dependent by construction and are left untouched.
"""
from typing import Any
_ATTENTION_CLASSES = ("Qwen3_5Attention", "Qwen3Attention")
def enable_bidirectional_attention(model: Any) -> None:
"""Encoder-ize the full-attention layers of a ColQwen3.5 model, in place."""
config = getattr(model, "config", None)
for cfg in (config, getattr(config, "text_config", None)):
# `create_causal_mask` falls back to `create_bidirectional_mask` on this flag.
if cfg is not None:
cfg.is_causal = False
for module in model.modules():
if module.__class__.__name__ in _ATTENTION_CLASSES and hasattr(module, "is_causal"):
module.is_causal = False
|