| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ColPali late interaction model for multi-modal retrieval and reranking. |
| 4 | |
| 5 | ColPali extends PaliGemma 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: PaliGemma backbone (SigLIP + Gemma) with custom text projection |
| 14 | |
| 15 | Target models: |
| 16 | - vidore/colpali-v1.3-hf |
| 17 | """ |
| 18 | |
| 19 | from collections.abc import Iterable, Mapping |
| 20 | |
| 21 | import torch |
| 22 | import torch.nn as nn |
| 23 | from transformers import PaliGemmaProcessor |
| 24 | |
| 25 | from vllm.config import VllmConfig |
| 26 | from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed |
| 27 | from vllm.model_executor.layers.pooler.tokwise.heads import TokenPoolerHead |
| 28 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 29 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 30 | from vllm.multimodal.parse import MultiModalDataItems |
| 31 | from vllm.multimodal.processing.processor import HFMultiModalInputs |
| 32 | |
| 33 | from .interfaces import SupportsLateInteraction |
| 34 | from .interfaces_base import default_pooling_type |
| 35 | from .paligemma import ( |
| 36 | PaliGemmaDummyInputsBuilder, |
| 37 | PaliGemmaForConditionalGeneration, |
| 38 | PaliGemmaMultiModalProcessor, |
| 39 | PaliGemmaProcessingInfo, |
| 40 | ) |
| 41 | from .utils import AutoWeightsLoader, WeightsMapper |
| 42 | |
| 43 | |
| 44 | class ColPaliProcessingInfo(PaliGemmaProcessingInfo): |
| 45 | """Processing info for ColPali models. |
| 46 | |
| 47 | ColPali models use a custom HuggingFace config (ColPaliConfig) that is |
| 48 | not an instance of PaliGemmaConfig. We override get_hf_config() and |
| 49 | get_hf_processor() to skip the strict type check. |
| 50 | """ |
| 51 | |
| 52 | def get_hf_config(self): |
| 53 | return self.ctx.get_hf_config() |
| 54 | |
| 55 | def get_hf_processor(self, **kwargs: object) -> PaliGemmaProcessor: |
| 56 | # Force standard PaliGemmaProcessor even when trust_remote_code=True. |
| 57 | return self.ctx.get_hf_processor(PaliGemmaProcessor, **kwargs) |
| 58 | |
| 59 | |
| 60 | class ColPaliMultiModalProcessor(PaliGemmaMultiModalProcessor): |
| 61 | """Multimodal processor for ColPali.""" |
| 62 | |
| 63 | def _get_hf_mm_inputs( |
| 64 | self, |
| 65 | mm_items: MultiModalDataItems, |
| 66 | hf_kwargs: Mapping[str, object], |
| 67 | ) -> HFMultiModalInputs: |
| 68 | hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) |
| 69 | |
| 70 | # The ColPali tokenizer_config.json ships with a small default |
| 71 | # max_length (50) that truncates the 1024 image tokens inserted |
| 72 | # by PaliGemmaProcessor, causing a token-count mismatch. |
| 73 | # vLLM enforces its own max_model_len, so we disable HF |
| 74 | # truncation to keep all image + text tokens intact. |
| 75 | return hf_inputs._replace(hf_kwargs=dict(hf_inputs.hf_kwargs, truncation=False)) |
| 76 | |
| 77 | |
| 78 | @default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL") |
| 79 | @MULTIMODAL_REGISTRY.register_processor( |
| 80 | ColPaliMultiModalProcessor, |
| 81 | info=ColPaliProcessingInfo, |
| 82 | dummy_inputs=PaliGemmaDummyInputsBuilder, |
| 83 | ) |
| 84 | class ColPaliModel( |
| 85 | PaliGemmaForConditionalGeneration, |
| 86 | SupportsLateInteraction, |
| 87 | ): |
| 88 | """ColPali late interaction model for multi-modal retrieval/reranking. |
| 89 | |
| 90 | This model extends PaliGemmaForConditionalGeneration with a ColBERT-style |
| 91 | linear projection layer for per-token embeddings. It supports: |
| 92 | - "token_embed" task: Per-token embeddings for late interaction scoring |
| 93 | |
| 94 | The model produces L2-normalized per-token embeddings by: |
| 95 | 1. Running the PaliGemma backbone (vision + language) to get hidden states |
| 96 | 2. Projecting hidden states through a linear layer (hidden_size -> embed_dim) |
| 97 | 3. L2-normalizing the projected embeddings |
| 98 | """ |
| 99 | |
| 100 | # Mark this as a pooling model so vLLM routes to pooler path |
| 101 | is_pooling_model = True |
| 102 | |
| 103 | # Override hf_to_vllm_mapper to handle ColPali weight naming. |
| 104 | hf_to_vllm_mapper = WeightsMapper( |
| 105 | orig_to_new_prefix={ |
| 106 | # HF transformers checkpoint (vidore/colpali-v1.3-hf) |
| 107 | # Weights: vlm.vision_tower.*, vlm.language_model.*, |
| 108 | # vlm.multi_modal_projector.* |
| 109 | "vlm.vision_tower.": "vision_tower.", |
| 110 | "vlm.language_model.": "language_model.", |
| 111 | "vlm.multi_modal_projector.": "multi_modal_projector.", |
| 112 | # colpali-engine checkpoint naming |
| 113 | "model.vision_tower.": "vision_tower.", |
| 114 | "model.language_model.": "language_model.", |
| 115 | "model.multi_modal_projector.": "multi_modal_projector.", |
| 116 | "lm_head.": "language_model.lm_head.", |
| 117 | } |
| 118 | ) |
| 119 | |
| 120 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 121 | super().__init__(vllm_config=vllm_config, prefix=prefix) |
| 122 | |
| 123 | config = vllm_config.model_config.hf_config |
| 124 | head_dtype = vllm_config.model_config.head_dtype |
| 125 | |
| 126 | hidden_size = getattr(config, "hidden_size", None) |
| 127 | if hidden_size is None and hasattr(config, "text_config"): |
| 128 | hidden_size = config.text_config.hidden_size |
| 129 | if hidden_size is None: |
| 130 | raise ValueError( |
| 131 | "Unable to determine text hidden size from config. " |
| 132 | "Expected 'hidden_size' or 'text_config.hidden_size'." |
| 133 | ) |
| 134 | self._proj_hidden_size = hidden_size |
| 135 | |
| 136 | # ColPali uses embedding_dim=128, but also check other naming variants |
| 137 | self.embed_dim: int | None = ( |
| 138 | getattr(config, "embedding_dim", None) |
| 139 | or getattr(config, "embed_dim", None) |
| 140 | or getattr(config, "dim", None) |
| 141 | or getattr(config, "projection_dim", None) |
| 142 | or getattr(config, "colbert_dim", None) |
| 143 | ) |
| 144 | |
| 145 | # Build the projection layer if embed_dim is known |
| 146 | if self.embed_dim is not None: |
| 147 | self.custom_text_proj = nn.Linear( |
| 148 | hidden_size, |
| 149 | self.embed_dim, |
| 150 | bias=False, |
| 151 | dtype=head_dtype, |
| 152 | ) |
| 153 | else: |
| 154 | # Will be created during load_weights when dim is inferred |
| 155 | self.custom_text_proj = None |
| 156 | |
| 157 | pooler_config = vllm_config.model_config.pooler_config |
| 158 | assert pooler_config is not None |
| 159 | self.pooler = pooler_for_token_embed( |
| 160 | pooler_config, |
| 161 | projector=self.custom_text_proj, |
| 162 | ) |
| 163 | |
| 164 | def forward( |
| 165 | self, |
| 166 | input_ids: torch.Tensor | None, |
| 167 | positions: torch.Tensor, |
| 168 | intermediate_tensors=None, |
| 169 | inputs_embeds: torch.Tensor | None = None, |
| 170 | **kwargs: object, |
| 171 | ) -> torch.Tensor: |
| 172 | return super().forward( |
| 173 | input_ids=input_ids, |
| 174 | positions=positions, |
| 175 | intermediate_tensors=intermediate_tensors, |
| 176 | inputs_embeds=inputs_embeds, |
| 177 | **kwargs, |
| 178 | ) |
| 179 | |
| 180 | # Names used for the projection layer across different ColPali variants |
| 181 | _PROJ_LAYER_NAMES = { |
| 182 | "custom_text_proj", # vLLM internal naming |
| 183 | "embedding_proj_layer", # colpali-engine / HF naming |
| 184 | } |
| 185 | |
| 186 | def _is_proj_weight(self, name: str) -> bool: |
| 187 | """Check if a weight name belongs to the projection layer.""" |
| 188 | return any(proj_name in name for proj_name in self._PROJ_LAYER_NAMES) |
| 189 | |
| 190 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 191 | """Load weights with special handling for ColPali projection layer.""" |
| 192 | weights_list = list(weights) |
| 193 | proj_weights: list[tuple[str, torch.Tensor]] = [] |
| 194 | model_weights: list[tuple[str, torch.Tensor]] = [] |
| 195 | |
| 196 | for name, weight in weights_list: |
| 197 | if self._is_proj_weight(name): |
| 198 | proj_weights.append((name, weight)) |
| 199 | else: |
| 200 | model_weights.append((name, weight)) |
| 201 | |
| 202 | loader = AutoWeightsLoader(self) |
| 203 | loaded = loader.load_weights(model_weights, mapper=self.hf_to_vllm_mapper) |
| 204 | |
| 205 | if proj_weights: |
| 206 | model_dtype = next(self.language_model.parameters()).dtype |
| 207 | model_device = next(self.language_model.parameters()).device |
| 208 | |
| 209 | for name, weight in proj_weights: |
| 210 | if self.embed_dim is None and "weight" in name: |
| 211 | self.embed_dim = weight.shape[0] |
| 212 | has_bias = any("bias" in n for n, _ in proj_weights) |
| 213 | self.custom_text_proj = nn.Linear( |
| 214 | self._proj_hidden_size, |
| 215 | self.embed_dim, |
| 216 | bias=has_bias, |
| 217 | dtype=model_dtype, |
| 218 | ) |
| 219 | self.custom_text_proj.to(model_device) |
| 220 | |
| 221 | if self.custom_text_proj is not None: |
| 222 | param_name = name.split(".")[-1] |
| 223 | param = getattr(self.custom_text_proj, param_name, None) |
| 224 | if param is not None: |
| 225 | weight = weight.to(device=param.device, dtype=param.dtype) |
| 226 | default_weight_loader(param, weight) |
| 227 | loaded.add(f"custom_text_proj.{param_name}") |
| 228 | |
| 229 | # Update pooler projector for the lazy-creation path |
| 230 | head = self.pooler.head |
| 231 | assert isinstance(head, TokenPoolerHead) |
| 232 | head.projector = self.custom_text_proj |
| 233 | |
| 234 | # Mark pooler projector params as loaded |
| 235 | if hasattr(self, "pooler") and hasattr(self.pooler, "head"): |
| 236 | head = self.pooler.head |
| 237 | projector = getattr(head, "projector", None) |
| 238 | if projector is not None and isinstance(projector, nn.Module): |
| 239 | for pname, _ in projector.named_parameters(): |
| 240 | loaded.add(f"pooler.head.projector.{pname}") |
| 241 | |
| 242 | return loaded |