| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ColQwen3 late interaction model for multi-modal retrieval and reranking. |
| 4 | |
| 5 | ColQwen3 extends Qwen3-VL 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-VL backbone with custom text projection |
| 14 | |
| 15 | Target models: |
| 16 | - TomoroAI/tomoro-colqwen3-embed-8b |
| 17 | - OpenSearch-AI/Ops-Colqwen3-4B |
| 18 | - nvidia/nemotron-colembed-vl-4b-v2 |
| 19 | """ |
| 20 | |
| 21 | from collections.abc import Iterable, Mapping |
| 22 | |
| 23 | import torch |
| 24 | import torch.nn as nn |
| 25 | from transformers.models.qwen3_vl import Qwen3VLProcessor |
| 26 | |
| 27 | from vllm.config import VllmConfig |
| 28 | from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed |
| 29 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 30 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 31 | |
| 32 | from .interfaces import SupportsLateInteraction |
| 33 | from .interfaces_base import default_pooling_type |
| 34 | from .qwen2_vl import Qwen2VLMultiModalDataParser |
| 35 | from .qwen3_vl import ( |
| 36 | Qwen3VLDummyInputsBuilder, |
| 37 | Qwen3VLForConditionalGeneration, |
| 38 | Qwen3VLMultiModalProcessor, |
| 39 | Qwen3VLProcessingInfo, |
| 40 | ) |
| 41 | from .utils import AutoWeightsLoader, WeightsMapper |
| 42 | |
| 43 | |
| 44 | class ColQwen3ProcessingInfo(Qwen3VLProcessingInfo): |
| 45 | """Processing info for ColQwen3 models. |
| 46 | |
| 47 | ColQwen3 models (TomoroAI, OpenSearch-AI, etc.) use custom HuggingFace |
| 48 | configs (e.g. ColQwen3Config, OpsColQwen3Config) that are not instances |
| 49 | of Qwen3VLConfig. We override get_hf_config() and get_hf_processor() |
| 50 | to skip the strict type check, similar to OpenCUAProcessingInfo. |
| 51 | """ |
| 52 | |
| 53 | def get_hf_config(self): |
| 54 | return self.ctx.get_hf_config() |
| 55 | |
| 56 | def get_hf_processor(self, **kwargs: object) -> Qwen3VLProcessor: |
| 57 | # Force standard Qwen3VLProcessor even when trust_remote_code=True. |
| 58 | # ColQwen3 custom processors (e.g. ColQwen3Processor) have |
| 59 | # incompatible interfaces with vLLM's Qwen3VLMultiModalProcessor. |
| 60 | # The standard Qwen3VLProcessor handles both text and image inputs |
| 61 | # correctly for the Qwen3-VL backbone. |
| 62 | return self.ctx.get_hf_processor( |
| 63 | Qwen3VLProcessor, |
| 64 | **kwargs, |
| 65 | ) |
| 66 | |
| 67 | @property |
| 68 | def _supports_video(self) -> bool: |
| 69 | """Check if the HF processor supports video inputs.""" |
| 70 | return hasattr(self.get_hf_processor(), "video_processor") |
| 71 | |
| 72 | def get_video_processor(self, **kwargs: object): |
| 73 | if not self._supports_video: |
| 74 | raise AttributeError( |
| 75 | f"The processor for {self.ctx.model_config.model} does not " |
| 76 | "support video inputs (no video_processor attribute)." |
| 77 | ) |
| 78 | return self.get_hf_processor(**kwargs).video_processor # type: ignore[attr-defined] |
| 79 | |
| 80 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 81 | limits: dict[str, int | None] = {"image": None} |
| 82 | if self._supports_video: |
| 83 | limits["video"] = None |
| 84 | return limits |
| 85 | |
| 86 | def get_mm_max_tokens_per_item( |
| 87 | self, |
| 88 | seq_len: int, |
| 89 | mm_counts: Mapping[str, int], |
| 90 | ) -> Mapping[str, int]: |
| 91 | max_image_tokens = self.get_max_image_tokens() |
| 92 | result: dict[str, int] = {"image": max_image_tokens} |
| 93 | if self._supports_video: |
| 94 | max_video_tokens = self.get_max_video_tokens(seq_len, mm_counts) |
| 95 | result["video"] = max_video_tokens |
| 96 | return result |
| 97 | |
| 98 | def get_data_parser(self): |
| 99 | hf_config = self.get_hf_config() |
| 100 | spatial_merge_size = hf_config.vision_config.spatial_merge_size |
| 101 | return Qwen2VLMultiModalDataParser( |
| 102 | spatial_merge_size, |
| 103 | video_needs_metadata=self._supports_video, |
| 104 | expected_hidden_size=self._get_expected_hidden_size(), |
| 105 | allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, |
| 106 | ) |
| 107 | |
| 108 | |
| 109 | @default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL") |
| 110 | @MULTIMODAL_REGISTRY.register_processor( |
| 111 | Qwen3VLMultiModalProcessor, |
| 112 | info=ColQwen3ProcessingInfo, |
| 113 | dummy_inputs=Qwen3VLDummyInputsBuilder, |
| 114 | ) |
| 115 | class ColQwen3Model(Qwen3VLForConditionalGeneration, SupportsLateInteraction): |
| 116 | """ColQwen3 late interaction model for multi-modal retrieval/reranking. |
| 117 | |
| 118 | This model extends Qwen3VLForConditionalGeneration 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-VL backbone (vision + language) to get hidden states |
| 124 | 2. Projecting hidden states through a linear layer (hidden_size -> embed_dim) |
| 125 | 3. L2-normalizing the projected embeddings |
| 126 | |
| 127 | ColBERT-style MaxSim scoring is computed externally, either client-side |
| 128 | or via the late interaction scoring path in ServingScores. |
| 129 | |
| 130 | Attributes: |
| 131 | custom_text_proj: Linear projection from hidden_size to embed_dim |
| 132 | |
| 133 | """ |
| 134 | |
| 135 | # Mark this as a pooling model so vLLM routes to pooler path |
| 136 | is_pooling_model = True |
| 137 | # Mark as false to avoid collision with the checkpoint's custom code path |
| 138 | supports_mm_device_do_normalize = False |
| 139 | |
| 140 | # Override hf_to_vllm_mapper to handle ColQwen3 weight naming. |
| 141 | # NOTE: WeightsMapper applies ALL matching prefix rules sequentially |
| 142 | # (no early exit), so more-specific prefixes must come first. |
| 143 | # TomoroAI: "vlm.model.visual.", "vlm.model.language_model." |
| 144 | # ColPali: "model.visual.", "model.language_model." |
| 145 | # OpenSearch: "visual.", "language_model." (no outer prefix, |
| 146 | # re-prefixed to "model.*" in load_weights) |
| 147 | hf_to_vllm_mapper = WeightsMapper( |
| 148 | orig_to_new_prefix={ |
| 149 | # TomoroAI naming convention (most specific first) |
| 150 | "vlm.model.visual.": "visual.", |
| 151 | "vlm.lm_head.": "language_model.lm_head.", |
| 152 | "vlm.model.language_model.": "language_model.model.", |
| 153 | # ColPali / nvidia naming convention |
| 154 | "model.visual.": "visual.", |
| 155 | "lm_head.": "language_model.lm_head.", |
| 156 | # OpenSearch-AI: after re-prefix, "language_model.model.*" |
| 157 | # becomes "model.language_model.model.*" — handle this before |
| 158 | # the shorter "model.language_model." rule to avoid double map |
| 159 | "model.language_model.model.": "language_model.model.", |
| 160 | "model.language_model.": "language_model.model.", |
| 161 | } |
| 162 | ) |
| 163 | |
| 164 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): |
| 165 | super().__init__(vllm_config=vllm_config, prefix=prefix) |
| 166 | |
| 167 | config = vllm_config.model_config.hf_config |
| 168 | head_dtype = vllm_config.model_config.head_dtype |
| 169 | |
| 170 | hidden_size = getattr(config, "hidden_size", None) |
| 171 | if hidden_size is None and hasattr(config, "text_config"): |
| 172 | hidden_size = config.text_config.hidden_size |
| 173 | if hidden_size is None: |
| 174 | raise ValueError( |
| 175 | "Unable to determine text hidden size from config. " |
| 176 | "Expected 'hidden_size' or 'text_config.hidden_size'." |
| 177 | ) |
| 178 | self._proj_hidden_size = hidden_size |
| 179 | |
| 180 | # (TomoroAI: embed_dim, OpenSearch: dims, ColPali: dim) |
| 181 | self.embed_dim: int | None = ( |
| 182 | getattr(config, "embed_dim", None) |
| 183 | or getattr(config, "dims", None) |
| 184 | or getattr(config, "dim", None) |
| 185 | or getattr(config, "projection_dim", None) |
| 186 | or getattr(config, "colbert_dim", None) |
| 187 | ) |
| 188 | |
| 189 | # Build the projection layer if embed_dim is known |
| 190 | if self.embed_dim is not None: |
| 191 | self.custom_text_proj = nn.Linear( |
| 192 | hidden_size, |
| 193 | self.embed_dim, |
| 194 | bias=False, |
| 195 | dtype=head_dtype, |
| 196 | ) |
| 197 | else: |
| 198 | # Will be created during load_weights when dim is inferred |
| 199 | self.custom_text_proj = None |
| 200 | |
| 201 | pooler_config = vllm_config.model_config.pooler_config |
| 202 | assert pooler_config is not None |
| 203 | self.pooler = pooler_for_token_embed( |
| 204 | pooler_config, |
| 205 | projector=None, |
| 206 | ) |
| 207 | |
| 208 | def forward( |
| 209 | self, |
| 210 | input_ids: torch.Tensor | None, |
| 211 | positions: torch.Tensor, |
| 212 | intermediate_tensors=None, |
| 213 | inputs_embeds: torch.Tensor | None = None, |
| 214 | **kwargs: object, |
| 215 | ) -> torch.Tensor: |
| 216 | """Run forward pass producing per-token embeddings.""" |
| 217 | hidden_states = super().forward( |
| 218 | input_ids=input_ids, |
| 219 | positions=positions, |
| 220 | intermediate_tensors=intermediate_tensors, |
| 221 | inputs_embeds=inputs_embeds, |
| 222 | **kwargs, |
| 223 | ) |
| 224 | |
| 225 | if not isinstance(hidden_states, torch.Tensor): |
| 226 | return hidden_states # type: ignore |
| 227 | |
| 228 | if self.custom_text_proj is not None: |
| 229 | proj_dtype = self.custom_text_proj.weight.dtype |
| 230 | if hidden_states.dtype != proj_dtype: |
| 231 | hidden_states = hidden_states.to(proj_dtype) |
| 232 | hidden_states = self.custom_text_proj(hidden_states) |
| 233 | |
| 234 | # L2 normalize |
| 235 | return torch.nn.functional.normalize(hidden_states, p=2, dim=-1) |
| 236 | |
| 237 | # Names used for the projection layer across different ColQwen3 variants |
| 238 | _PROJ_LAYER_NAMES = { |
| 239 | "custom_text_proj", # ColPali naming |
| 240 | "embedding_proj_layer", # TomoroAI naming |
| 241 | } |
| 242 | |
| 243 | def _is_proj_weight(self, name: str) -> bool: |
| 244 | """Check if a weight name belongs to the projection layer.""" |
| 245 | return any(proj_name in name for proj_name in self._PROJ_LAYER_NAMES) |
| 246 | |
| 247 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 248 | """Load weights with special handling for ColQwen3 projection layer.""" |
| 249 | weights_list = list(weights) |
| 250 | proj_weights: list[tuple[str, torch.Tensor]] = [] |
| 251 | model_weights: list[tuple[str, torch.Tensor]] = [] |
| 252 | |
| 253 | # Scan all weight names to determine if re-prefixing is needed. |
| 254 | # OpenSearch-AI models have unprefixed weights ("language_model.*", |
| 255 | # "visual.*") that need "model." added so hf_to_vllm_mapper can |
| 256 | # process them. Only re-prefix if ALL backbone weights are |
| 257 | # unprefixed (no "vlm." or "model." prefix found). |
| 258 | has_unprefixed = any( |
| 259 | name.startswith("language_model.") or name.startswith("visual.") |
| 260 | for name, _ in weights_list |
| 261 | ) |
| 262 | has_prefixed = any( |
| 263 | name.startswith("vlm.") or name.startswith("model.") |
| 264 | for name, _ in weights_list |
| 265 | ) |
| 266 | needs_reprefix = has_unprefixed and not has_prefixed |
| 267 | |
| 268 | for name, weight in weights_list: |
| 269 | if self._is_proj_weight(name): |
| 270 | proj_weights.append((name, weight)) |
| 271 | else: |
| 272 | if needs_reprefix and not self._is_proj_weight(name): |
| 273 | name = "model." + name |
| 274 | model_weights.append((name, weight)) |
| 275 | |
| 276 | loader = AutoWeightsLoader(self) |
| 277 | loaded = loader.load_weights(model_weights, mapper=self.hf_to_vllm_mapper) |
| 278 | |
| 279 | if proj_weights: |
| 280 | model_dtype = next(self.language_model.parameters()).dtype |
| 281 | model_device = next(self.language_model.parameters()).device |
| 282 | |
| 283 | for name, weight in proj_weights: |
| 284 | if self.embed_dim is None and "weight" in name: |
| 285 | self.embed_dim = weight.shape[0] |
| 286 | has_bias = any("bias" in n for n, _ in proj_weights) |
| 287 | self.custom_text_proj = nn.Linear( |
| 288 | self._proj_hidden_size, |
| 289 | self.embed_dim, |
| 290 | bias=has_bias, |
| 291 | dtype=model_dtype, |
| 292 | ) |
| 293 | self.custom_text_proj.to(model_device) |
| 294 | |
| 295 | if self.custom_text_proj is not None: |
| 296 | param_name = name.split(".")[-1] |
| 297 | param = getattr(self.custom_text_proj, param_name, None) |
| 298 | if param is not None: |
| 299 | weight = weight.to(device=param.device, dtype=param.dtype) |
| 300 | default_weight_loader(param, weight) |
| 301 | loaded.add(f"custom_text_proj.{param_name}") |
| 302 | |
| 303 | return loaded |