| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ColModernVBERT: multimodal late-interaction retrieval model. |
| 4 | |
| 5 | Combines SigLIP vision encoder + ModernBERT text encoder with a pixel |
| 6 | shuffle connector and ColBERT-style 128-dim per-token embeddings. |
| 7 | |
| 8 | Reference: https://huggingface.co/ModernVBERT/colmodernvbert-merged |
| 9 | """ |
| 10 | |
| 11 | from collections.abc import Iterable, Mapping, Sequence |
| 12 | |
| 13 | import torch |
| 14 | from torch import nn |
| 15 | from transformers import BatchFeature |
| 16 | |
| 17 | from vllm.config import VllmConfig |
| 18 | from vllm.config.multimodal import MultiModalDummyOptions |
| 19 | from vllm.inputs import MultiModalDataDict |
| 20 | from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed |
| 21 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 22 | from vllm.multimodal.inputs import ( |
| 23 | MultiModalFieldConfig, |
| 24 | MultiModalKwargsItems, |
| 25 | ) |
| 26 | from vllm.multimodal.parse import ImageSize, MultiModalDataItems |
| 27 | from vllm.multimodal.processing import ( |
| 28 | BaseDummyInputsBuilder, |
| 29 | BaseMultiModalProcessor, |
| 30 | BaseProcessingInfo, |
| 31 | PromptIndexTargets, |
| 32 | PromptReplacement, |
| 33 | PromptUpdate, |
| 34 | ) |
| 35 | from vllm.sequence import IntermediateTensors |
| 36 | from vllm.tokenizers.hf import HfTokenizer |
| 37 | from vllm.transformers_utils.configs.colmodernvbert import ColModernVBertConfig |
| 38 | |
| 39 | from .interfaces import ( |
| 40 | MultiModalEmbeddings, |
| 41 | SupportsLateInteraction, |
| 42 | SupportsMultiModal, |
| 43 | ) |
| 44 | from .interfaces_base import default_pooling_type |
| 45 | from .modernbert import ModernBertEmbeddings, ModernBertLayer |
| 46 | from .siglip import SiglipVisionModel |
| 47 | from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix |
| 48 | |
| 49 | # --------------------------------------------------------------------------- |
| 50 | # Connector: pixel shuffle + simple linear projection |
| 51 | # --------------------------------------------------------------------------- |
| 52 | |
| 53 | |
| 54 | class ColModernVBertConnector(nn.Module): |
| 55 | """Pixel shuffle spatial reduction followed by a linear projection. |
| 56 | |
| 57 | Reduces the vision encoder's token count by ``factor^2`` via pixel-shuffle |
| 58 | spatial rearrangement, then projects the concatenated channels to the text |
| 59 | encoder's hidden size with a single bias-free linear layer. |
| 60 | """ |
| 61 | |
| 62 | def __init__(self, config: ColModernVBertConfig): |
| 63 | super().__init__() |
| 64 | self.pixel_shuffle_factor = config.pixel_shuffle_factor |
| 65 | vision_hidden_size = config.vision_config.hidden_size |
| 66 | input_size = vision_hidden_size * (self.pixel_shuffle_factor**2) |
| 67 | output_size = config.hidden_size |
| 68 | self.proj = nn.Linear(input_size, output_size, bias=False) |
| 69 | |
| 70 | def pixel_shuffle(self, features: torch.Tensor) -> torch.Tensor: |
| 71 | """Spatial rearrangement that reduces seq length by factor^2.""" |
| 72 | batch_size, seq_length, hidden_size = features.shape |
| 73 | height = width = int(seq_length**0.5) |
| 74 | factor = self.pixel_shuffle_factor |
| 75 | |
| 76 | # Reshape to (B, H, W, C) |
| 77 | features = features.view(batch_size, height, width, hidden_size) |
| 78 | |
| 79 | # Reshape to (B, H/f, f, W/f, f, C) |
| 80 | features = features.view( |
| 81 | batch_size, height // factor, factor, width // factor, factor, hidden_size |
| 82 | ) |
| 83 | |
| 84 | # Permute to (B, H/f, W/f, f, f, C) |
| 85 | features = features.permute(0, 1, 3, 2, 4, 5) |
| 86 | |
| 87 | # Reshape to (B, H/f, W/f, C * f^2) |
| 88 | new_hidden_size = hidden_size * (factor**2) |
| 89 | features = features.reshape( |
| 90 | batch_size, height // factor, width // factor, new_hidden_size |
| 91 | ) |
| 92 | |
| 93 | return features |
| 94 | |
| 95 | def forward(self, features: torch.Tensor) -> torch.Tensor: |
| 96 | features = self.pixel_shuffle(features) |
| 97 | batch_size = features.shape[0] |
| 98 | features = features.reshape(batch_size, -1, features.shape[-1]) |
| 99 | return self.proj(features) |
| 100 | |
| 101 | |
| 102 | # --------------------------------------------------------------------------- |
| 103 | # Multimodal processing |
| 104 | # --------------------------------------------------------------------------- |
| 105 | |
| 106 | |
| 107 | class ColModernVBertProcessingInfo(BaseProcessingInfo): |
| 108 | def get_hf_config(self) -> ColModernVBertConfig: |
| 109 | return self.ctx.get_hf_config(ColModernVBertConfig) |
| 110 | |
| 111 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 112 | return {"image": None} |
| 113 | |
| 114 | def get_image_size_with_most_features(self) -> ImageSize: |
| 115 | config = self.get_hf_config() |
| 116 | size = config.vision_config.image_size |
| 117 | return ImageSize(width=size, height=size) |
| 118 | |
| 119 | def get_num_image_tokens( |
| 120 | self, |
| 121 | *, |
| 122 | image_width: int, |
| 123 | image_height: int, |
| 124 | ) -> int: |
| 125 | return self.get_hf_config().image_seq_len |
| 126 | |
| 127 | |
| 128 | class ColModernVBertDummyInputsBuilder( |
| 129 | BaseDummyInputsBuilder[ColModernVBertProcessingInfo], |
| 130 | ): |
| 131 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 132 | return "" |
| 133 | |
| 134 | def get_dummy_mm_data( |
| 135 | self, |
| 136 | seq_len: int, |
| 137 | mm_counts: Mapping[str, int], |
| 138 | mm_options: MultiModalDummyOptions, |
| 139 | ) -> MultiModalDataDict: |
| 140 | target_width, target_height = self.info.get_image_size_with_most_features() |
| 141 | |
| 142 | return { |
| 143 | "image": self._get_dummy_images( |
| 144 | width=target_width, |
| 145 | height=target_height, |
| 146 | num_images=mm_counts.get("image", 0), |
| 147 | overrides=mm_options.get("image"), |
| 148 | ) |
| 149 | } |
| 150 | |
| 151 | |
| 152 | class ColModernVBertMultiModalProcessor( |
| 153 | BaseMultiModalProcessor[ColModernVBertProcessingInfo], |
| 154 | ): |
| 155 | def _apply_hf_processor_main( |
| 156 | self, |
| 157 | mm_items: MultiModalDataItems, |
| 158 | hf_kwargs: Mapping[str, object], |
| 159 | ) -> BatchFeature: |
| 160 | mm_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( |
| 161 | mm_items, hf_kwargs |
| 162 | ) |
| 163 | |
| 164 | if not mm_data: |
| 165 | return self._finalize_hf_mm_data(mm_data, hf_kwargs, passthrough_data) |
| 166 | |
| 167 | prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) |
| 168 | |
| 169 | tokenizer = self.info.get_tokenizer() |
| 170 | assert isinstance(tokenizer, HfTokenizer) |
| 171 | text_encoding = tokenizer( |
| 172 | prompt_text, |
| 173 | return_tensors="pt", |
| 174 | ) |
| 175 | result = BatchFeature(data=dict(text_encoding)) |
| 176 | |
| 177 | images = mm_data.get("images") |
| 178 | if images: |
| 179 | from transformers import Idefics3ImageProcessor |
| 180 | |
| 181 | image_processor = Idefics3ImageProcessor.from_pretrained( |
| 182 | self.info.ctx.model_config.model, |
| 183 | revision=self.info.ctx.model_config.revision, |
| 184 | ) |
| 185 | image_outputs = image_processor( |
| 186 | images=images, |
| 187 | do_image_splitting=False, |
| 188 | return_tensors="pt", |
| 189 | ) |
| 190 | result.update(image_outputs) |
| 191 | |
| 192 | processed_data = result |
| 193 | return self._finalize_hf_mm_data( |
| 194 | mm_data, hf_kwargs, passthrough_data, processed_data |
| 195 | ) |
| 196 | |
| 197 | def _get_mm_fields_config( |
| 198 | self, |
| 199 | hf_inputs: BatchFeature, |
| 200 | hf_processor_mm_kwargs: Mapping[str, object], |
| 201 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 202 | return dict( |
| 203 | pixel_values=MultiModalFieldConfig.batched("image"), |
| 204 | ) |
| 205 | |
| 206 | def _get_prompt_updates( |
| 207 | self, |
| 208 | mm_items: MultiModalDataItems, |
| 209 | hf_processor_mm_kwargs: Mapping[str, object], |
| 210 | out_mm_kwargs: MultiModalKwargsItems, |
| 211 | ) -> Sequence[PromptUpdate]: |
| 212 | config = self.info.get_hf_config() |
| 213 | image_token_id = config.image_token_id |
| 214 | num_tokens = config.image_seq_len |
| 215 | |
| 216 | def get_replacement(item_idx: int): |
| 217 | return [image_token_id] * num_tokens |
| 218 | |
| 219 | return [ |
| 220 | PromptReplacement( |
| 221 | modality="image", |
| 222 | target=PromptIndexTargets.start(), |
| 223 | replacement=get_replacement, |
| 224 | ), |
| 225 | ] |
| 226 | |
| 227 | |
| 228 | # --------------------------------------------------------------------------- |
| 229 | # Model |
| 230 | # --------------------------------------------------------------------------- |
| 231 | |
| 232 | |
| 233 | @MULTIMODAL_REGISTRY.register_processor( |
| 234 | ColModernVBertMultiModalProcessor, |
| 235 | info=ColModernVBertProcessingInfo, |
| 236 | dummy_inputs=ColModernVBertDummyInputsBuilder, |
| 237 | ) |
| 238 | @default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL") |
| 239 | class ColModernVBertForRetrieval( |
| 240 | nn.Module, SupportsMultiModal, SupportsLateInteraction |
| 241 | ): |
| 242 | """ColModernVBERT multimodal late-interaction retrieval model. |
| 243 | |
| 244 | Architecture: |
| 245 | Image -> SiglipVisionModel -> ColModernVBertConnector |
| 246 | ↓ |
| 247 | Text -> ModernBertEmbeddings → [merge] → ModernBertLayers → norm |
| 248 | ↓ |
| 249 | custom_text_proj → L2 norm |
| 250 | ↓ |
| 251 | per-token 128-d embeddings |
| 252 | """ |
| 253 | |
| 254 | is_pooling_model = True |
| 255 | |
| 256 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 257 | super().__init__() |
| 258 | config: ColModernVBertConfig = vllm_config.model_config.hf_config |
| 259 | self.config = config |
| 260 | text_config = config.text_config |
| 261 | quant_config = vllm_config.quant_config |
| 262 | |
| 263 | # --- Vision encoder (reuses SiglipVisionModel from siglip.py) --- |
| 264 | self.vision_model = SiglipVisionModel( |
| 265 | config.vision_config, |
| 266 | quant_config, |
| 267 | prefix=maybe_prefix(prefix, "vision_model"), |
| 268 | ) |
| 269 | |
| 270 | # --- Connector (pixel shuffle + linear projection) --- |
| 271 | self.connector = ColModernVBertConnector(config) |
| 272 | |
| 273 | # --- Text encoder (built from ModernBERT components directly) --- |
| 274 | # We build the components individually rather than wrapping |
| 275 | # ``ModernBertModel`` because ``ModernBertEncoderLayer`` reads |
| 276 | # ``vllm_config.model_config.hf_config`` which would be |
| 277 | # ``ColModernVBertConfig``, not ``ModernBertConfig``. |
| 278 | self.text_embeddings = ModernBertEmbeddings(text_config) |
| 279 | self.text_layers = nn.ModuleList( |
| 280 | [ |
| 281 | ModernBertLayer( |
| 282 | config=text_config, |
| 283 | layer_id=i, |
| 284 | prefix=f"{prefix}.text_layers.{i}", |
| 285 | ) |
| 286 | for i in range(text_config.num_hidden_layers) |
| 287 | ] |
| 288 | ) |
| 289 | self.text_final_norm = nn.LayerNorm( |
| 290 | text_config.hidden_size, |
| 291 | eps=text_config.norm_eps, |
| 292 | bias=text_config.norm_bias, |
| 293 | ) |
| 294 | |
| 295 | # --- ColBERT projection (768 -> 128, with bias) --- |
| 296 | self.custom_text_proj = nn.Linear( |
| 297 | text_config.hidden_size, |
| 298 | config.embedding_dim, |
| 299 | bias=True, |
| 300 | dtype=vllm_config.model_config.head_dtype, |
| 301 | ) |
| 302 | |
| 303 | # --- Pooler (applies projection + L2 normalize) --- |
| 304 | pooler_config = vllm_config.model_config.pooler_config |
| 305 | assert pooler_config is not None |
| 306 | self.pooler = pooler_for_token_embed( |
| 307 | pooler_config, |
| 308 | projector=self.custom_text_proj, |
| 309 | ) |
| 310 | |
| 311 | # ---- multimodal --------------------------------------------------------- |
| 312 | |
| 313 | def _get_image_features( |
| 314 | self, |
| 315 | pixel_values: torch.Tensor, |
| 316 | ) -> torch.Tensor: |
| 317 | # Idefics3ImageProcessor may return (batch, tiles, C, H, W); |
| 318 | # flatten to (batch*tiles, C, H, W) for SiglipVisionModel. |
| 319 | if pixel_values.dim() == 5: |
| 320 | b, t, c, h, w = pixel_values.shape |
| 321 | pixel_values = pixel_values.reshape(b * t, c, h, w) |
| 322 | vision_outputs = self.vision_model( |
| 323 | pixel_values.to(dtype=self.vision_model.dtype), |
| 324 | ) |
| 325 | return self.connector(vision_outputs) |
| 326 | |
| 327 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 328 | pixel_values = kwargs.pop("pixel_values", None) |
| 329 | if pixel_values is None: |
| 330 | return [] |
| 331 | assert isinstance(pixel_values, torch.Tensor) |
| 332 | image_features = self._get_image_features(pixel_values) |
| 333 | return list(image_features) |
| 334 | |
| 335 | # ---- forward ------------------------------------------------------------ |
| 336 | |
| 337 | def forward( |
| 338 | self, |
| 339 | input_ids: torch.Tensor, |
| 340 | positions: torch.Tensor, |
| 341 | intermediate_tensors: IntermediateTensors | None = None, |
| 342 | inputs_embeds: torch.Tensor | None = None, |
| 343 | ) -> torch.Tensor: |
| 344 | hidden_states = self.text_embeddings(input_ids, inputs_embeds=inputs_embeds) |
| 345 | |
| 346 | for layer in self.text_layers: |
| 347 | hidden_states = layer(hidden_states, positions) |
| 348 | |
| 349 | return self.text_final_norm(hidden_states) |
| 350 | |
| 351 | # ---- weight loading ----------------------------------------------------- |
| 352 | |
| 353 | # Checkpoint prefix → vLLM param prefix. |
| 354 | # More-specific prefixes must appear before shorter ones. |
| 355 | hf_to_vllm_mapper = WeightsMapper( |
| 356 | orig_to_new_prefix={ |
| 357 | "model.text_model.layers.": "text_layers.", |
| 358 | "model.text_model.embeddings.": "text_embeddings.", |
| 359 | "model.text_model.final_norm.": "text_final_norm.", |
| 360 | "model.connector.modality_projection.": "connector.proj.", |
| 361 | "model.custom_text_proj.": "custom_text_proj.", |
| 362 | "model.vision_model.vision_model.": "vision_model.vision_model.", |
| 363 | "model.": "", |
| 364 | }, |
| 365 | ) |
| 366 | |
| 367 | def load_weights( |
| 368 | self, |
| 369 | weights: Iterable[tuple[str, torch.Tensor]], |
| 370 | ) -> set[str]: |
| 371 | loader = AutoWeightsLoader(self) |
| 372 | loaded_params = loader.load_weights( |
| 373 | weights, |
| 374 | mapper=self.hf_to_vllm_mapper, |
| 375 | ) |
| 376 | |
| 377 | # The pooler wraps ``custom_text_proj`` as its head projector. |
| 378 | # Mark those params as loaded under the pooler path too. |
| 379 | if hasattr(self, "pooler") and hasattr(self.pooler, "head"): |
| 380 | head = self.pooler.head |
| 381 | projector = getattr(head, "projector", None) |
| 382 | if projector is not None and isinstance(projector, nn.Module): |
| 383 | for pname, _ in projector.named_parameters(): |
| 384 | loaded_params.add(f"pooler.head.projector.{pname}") |
| 385 | |
| 386 | return loaded_params |