| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ColQwen3.5 late interaction model for multi-modal retrieval and reranking. |
| 4 | |
| 5 | ColQwen3.5 extends Qwen3.5 with a ColBERT-style late interaction head, |
| 6 | producing per-token embeddings for both text and image inputs. It uses |
| 7 | MaxSim scoring for retrieval/reranking tasks. |
| 8 | |
| 9 | This model supports the "token_embed" pooling task and is designed for |
| 10 | multi-vector retrieval of documents containing both text and images. |
| 11 | |
| 12 | Reference: https://arxiv.org/abs/2407.01449 (ColPali) |
| 13 | Based on: Qwen3.5 backbone with custom text projection |
| 14 | |
| 15 | Target models: |
| 16 | - athrael-soju/colqwen3.5-4.5B-v3 |
| 17 | - vultr/VultronRetrieverPrime-Qwen3.5-8B |
| 18 | """ |
| 19 | |
| 20 | from collections.abc import Iterable, Mapping |
| 21 | |
| 22 | import torch |
| 23 | import torch.nn as nn |
| 24 | from transformers.models.qwen3_vl import Qwen3VLProcessor |
| 25 | |
| 26 | from vllm.config import VllmConfig |
| 27 | from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed |
| 28 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 29 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 30 | |
| 31 | from .interfaces import SupportsLateInteraction |
| 32 | from .interfaces_base import default_pooling_type |
| 33 | from .qwen2_vl import Qwen2VLMultiModalDataParser |
| 34 | from .qwen3_5 import ( |
| 35 | Qwen3_5ForConditionalGeneration, |
| 36 | Qwen3_5ProcessingInfo, |
| 37 | ) |
| 38 | from .qwen3_vl import ( |
| 39 | Qwen3VLDummyInputsBuilder, |
| 40 | Qwen3VLMultiModalProcessor, |
| 41 | ) |
| 42 | from .utils import AutoWeightsLoader, WeightsMapper |
| 43 | |
| 44 | |
| 45 | class ColQwen3_5ProcessingInfo(Qwen3_5ProcessingInfo): |
| 46 | """Processing info for ColQwen3.5 models. |
| 47 | |
| 48 | ColQwen3.5 models use custom HuggingFace processors (e.g. |
| 49 | ColQwen3_5Processor) that are incompatible with vLLM's |
| 50 | Qwen3VLMultiModalProcessor. We override get_hf_config() and |
| 51 | get_hf_processor() to skip the strict type check and force the |
| 52 | standard Qwen3VLProcessor. |
| 53 | """ |
| 54 | |
| 55 | def get_hf_config(self): |
| 56 | return self.ctx.get_hf_config() |
| 57 | |
| 58 | def get_hf_processor(self, **kwargs: object) -> Qwen3VLProcessor: |
| 59 | return self.ctx.get_hf_processor( |
| 60 | Qwen3VLProcessor, |
| 61 | **kwargs, |
| 62 | ) |
| 63 | |
| 64 | @property |
| 65 | def _supports_video(self) -> bool: |
| 66 | """Check if the HF processor supports video inputs.""" |
| 67 | return hasattr(self.get_hf_processor(), "video_processor") |
| 68 | |
| 69 | def get_video_processor(self, **kwargs: object): |
| 70 | if not self._supports_video: |
| 71 | raise AttributeError( |
| 72 | f"The processor for {self.ctx.model_config.model} does not " |
| 73 | "support video inputs (no video_processor attribute)." |
| 74 | ) |
| 75 | return self.get_hf_processor(**kwargs).video_processor # type: ignore[attr-defined] |
| 76 | |
| 77 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 78 | limits: dict[str, int | None] = {"image": None} |
| 79 | if self._supports_video: |
| 80 | limits["video"] = None |
| 81 | return limits |
| 82 | |
| 83 | def get_mm_max_tokens_per_item( |
| 84 | self, |
| 85 | seq_len: int, |
| 86 | mm_counts: Mapping[str, int], |
| 87 | ) -> Mapping[str, int]: |
| 88 | max_image_tokens = self.get_max_image_tokens() |
| 89 | result: dict[str, int] = {"image": max_image_tokens} |
| 90 | if self._supports_video: |
| 91 | max_video_tokens = self.get_max_video_tokens(seq_len, mm_counts) |
| 92 | result["video"] = max_video_tokens |
| 93 | return result |
| 94 | |
| 95 | def get_data_parser(self): |
| 96 | hf_config = self.get_hf_config() |
| 97 | spatial_merge_size = hf_config.vision_config.spatial_merge_size |
| 98 | return Qwen2VLMultiModalDataParser( |
| 99 | spatial_merge_size, |
| 100 | video_needs_metadata=self._supports_video, |
| 101 | expected_hidden_size=self._get_expected_hidden_size(), |
| 102 | allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, |
| 103 | ) |
| 104 | |
| 105 | |
| 106 | @default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL") |
| 107 | @MULTIMODAL_REGISTRY.register_processor( |
| 108 | Qwen3VLMultiModalProcessor, |
| 109 | info=ColQwen3_5ProcessingInfo, |
| 110 | dummy_inputs=Qwen3VLDummyInputsBuilder, |
| 111 | ) |
| 112 | class ColQwen3_5Model( |
| 113 | Qwen3_5ForConditionalGeneration, |
| 114 | SupportsLateInteraction, |
| 115 | ): |
| 116 | """ColQwen3.5 late interaction model for multi-modal retrieval/reranking. |
| 117 | |
| 118 | This model extends Qwen3_5ForConditionalGeneration with a ColBERT-style |
| 119 | linear projection layer for per-token embeddings. It supports: |
| 120 | - "token_embed" task: Per-token embeddings for late interaction scoring |
| 121 | |
| 122 | The model produces L2-normalized per-token embeddings by: |
| 123 | 1. Running the Qwen3.5 backbone (vision + language) to get hidden states |
| 124 | 2. Projection and L2-normalization via the pooler (TokenEmbeddingPoolerHead) |
| 125 | |
| 126 | Attributes: |
| 127 | custom_text_proj: Linear projection from hidden_size to embed_dim. |
| 128 | This is passed to the pooler as pooler.head.projector. |
| 129 | |
| 130 | """ |
| 131 | |
| 132 | # Mark this as a pooling model so vLLM routes to pooler path |
| 133 | is_pooling_model = True |
| 134 | |
| 135 | # Override hf_to_vllm_mapper to handle ColQwen3.5 weight naming. |
| 136 | # ColPali saves weights as "language_model.*" but vLLM's |
| 137 | # Qwen3_5ForCausalLM has them under "language_model.model.*". |
| 138 | # Visual weights ("visual.*") already match the vLLM module path. |
| 139 | hf_to_vllm_mapper = WeightsMapper( |
| 140 | orig_to_new_prefix={ |
| 141 | "language_model.": "language_model.model.", |
| 142 | "mtp.": None, |
| 143 | } |
| 144 | ) |
| 145 | |
| 146 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 147 | super().__init__(vllm_config=vllm_config, prefix=prefix) |
| 148 | |
| 149 | config = vllm_config.model_config.hf_config |
| 150 | head_dtype = vllm_config.model_config.head_dtype |
| 151 | |
| 152 | hidden_size = getattr(config, "hidden_size", None) |
| 153 | if hidden_size is None and hasattr(config, "text_config"): |
| 154 | hidden_size = config.text_config.hidden_size |
| 155 | if hidden_size is None: |
| 156 | raise ValueError( |
| 157 | "Unable to determine text hidden size from config. " |
| 158 | "Expected 'hidden_size' or 'text_config.hidden_size'." |
| 159 | ) |
| 160 | |
| 161 | # (ColPali: dim, projection_dim, colbert_dim) |
| 162 | self.embed_dim: int = ( |
| 163 | getattr(config, "embed_dim", None) |
| 164 | or getattr(config, "dims", None) |
| 165 | or getattr(config, "dim", None) |
| 166 | or getattr(config, "projection_dim", None) |
| 167 | or getattr(config, "colbert_dim", None) |
| 168 | or 128 # default from reference implementation |
| 169 | ) |
| 170 | |
| 171 | # ColPali defines `custom_text_proj = nn.Linear(hidden, dim)`, i.e. |
| 172 | # bias=True by default, and the trained ColQwen3.5 checkpoints ship a |
| 173 | # `custom_text_proj.bias`. Construct with a bias and zero-initialize it: |
| 174 | # a (legacy) bias-less checkpoint then behaves identically to bias=False, |
| 175 | # while load_weights() below picks up a trained bias instead of silently |
| 176 | # dropping it (which shifts every per-token vector and the MaxSim ranking). |
| 177 | self.custom_text_proj = nn.Linear( |
| 178 | hidden_size, |
| 179 | self.embed_dim, |
| 180 | bias=True, |
| 181 | dtype=head_dtype, |
| 182 | ) |
| 183 | nn.init.zeros_(self.custom_text_proj.bias) |
| 184 | |
| 185 | pooler_config = vllm_config.model_config.pooler_config |
| 186 | assert pooler_config is not None |
| 187 | self.pooler = pooler_for_token_embed( |
| 188 | pooler_config, |
| 189 | projector=self.custom_text_proj, |
| 190 | ) |
| 191 | |
| 192 | def forward( |
| 193 | self, |
| 194 | input_ids: torch.Tensor | None, |
| 195 | positions: torch.Tensor, |
| 196 | intermediate_tensors=None, |
| 197 | inputs_embeds: torch.Tensor | None = None, |
| 198 | **kwargs: object, |
| 199 | ) -> torch.Tensor: |
| 200 | """Run forward pass returning hidden states for pooler.""" |
| 201 | return super().forward( |
| 202 | input_ids=input_ids, |
| 203 | positions=positions, |
| 204 | intermediate_tensors=intermediate_tensors, |
| 205 | inputs_embeds=inputs_embeds, |
| 206 | **kwargs, |
| 207 | ) |
| 208 | |
| 209 | # Names used for the projection layer across different ColQwen3.5 variants |
| 210 | _PROJ_LAYER_NAMES = { |
| 211 | "custom_text_proj", # ColPali naming |
| 212 | "embedding_proj_layer", # Alternative naming |
| 213 | } |
| 214 | |
| 215 | def _is_proj_weight(self, name: str) -> bool: |
| 216 | """Check if a weight name belongs to the projection layer.""" |
| 217 | return any(proj_name in name for proj_name in self._PROJ_LAYER_NAMES) |
| 218 | |
| 219 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 220 | """Load weights with special handling for projection layer.""" |
| 221 | weights_list = list(weights) |
| 222 | proj_weights: list[tuple[str, torch.Tensor]] = [] |
| 223 | model_weights: list[tuple[str, torch.Tensor]] = [] |
| 224 | |
| 225 | for name, weight in weights_list: |
| 226 | if self._is_proj_weight(name): |
| 227 | proj_weights.append((name, weight)) |
| 228 | else: |
| 229 | model_weights.append((name, weight)) |
| 230 | |
| 231 | loader = AutoWeightsLoader(self) |
| 232 | loaded = loader.load_weights(model_weights, mapper=self.hf_to_vllm_mapper) |
| 233 | |
| 234 | for name, weight in proj_weights: |
| 235 | param_name = name.split(".")[-1] |
| 236 | param = getattr(self.custom_text_proj, param_name, None) |
| 237 | if param is not None: |
| 238 | weight = weight.to(device=param.device, dtype=param.dtype) |
| 239 | default_weight_loader(param, weight) |
| 240 | loaded.add(f"custom_text_proj.{param_name}") |
| 241 | # Also mark as loaded under pooler path since custom_text_proj |
| 242 | # is assigned to pooler.head.projector |
| 243 | loaded.add(f"pooler.head.projector.{param_name}") |
| 244 | |
| 245 | return loaded |