| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for ColBERT late interaction scoring. |
| 4 | |
| 5 | Tests are parametrized across multiple ColBERT backbones to ensure the |
| 6 | generic ColBERT support works with different encoder architectures. |
| 7 | """ |
| 8 | |
| 9 | from contextlib import contextmanager |
| 10 | |
| 11 | import pytest |
| 12 | import torch |
| 13 | |
| 14 | from tests.utils import wait_for_memory_to_settle |
| 15 | from vllm.distributed import cleanup_dist_env_and_memory |
| 16 | from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score |
| 17 | |
| 18 | # ----------------------------------------------------------------------- |
| 19 | # Model definitions: (model_name, colbert_dim, extra vllm_runner kwargs) |
| 20 | # ----------------------------------------------------------------------- |
| 21 | COLBERT_MODELS = { |
| 22 | "bert": { |
| 23 | "model": "answerdotai/answerai-colbert-small-v1", |
| 24 | "colbert_dim": 96, |
| 25 | "max_model_len": 512, |
| 26 | "extra_kwargs": {}, |
| 27 | "hf_comparison": { |
| 28 | "weights_file": "model.safetensors", |
| 29 | "weights_key": "linear.weight", |
| 30 | "trust_remote_code": False, |
| 31 | "model_cls": "BertModel", |
| 32 | }, |
| 33 | }, |
| 34 | "modernbert": { |
| 35 | "model": "lightonai/GTE-ModernColBERT-v1", |
| 36 | "colbert_dim": 128, |
| 37 | "max_model_len": 299, |
| 38 | "extra_kwargs": { |
| 39 | "hf_overrides": { |
| 40 | "architectures": ["ColBERTModernBertModel"], |
| 41 | }, |
| 42 | }, |
| 43 | "hf_comparison": { |
| 44 | "weights_file": "1_Dense/model.safetensors", |
| 45 | "weights_key": "linear.weight", |
| 46 | "trust_remote_code": False, |
| 47 | "model_cls": "AutoModel", |
| 48 | }, |
| 49 | }, |
| 50 | "jina": { |
| 51 | "model": "jinaai/jina-colbert-v2", |
| 52 | "colbert_dim": 128, |
| 53 | "max_model_len": 8192, |
| 54 | "extra_kwargs": { |
| 55 | "hf_overrides": { |
| 56 | "architectures": ["ColBERTJinaRobertaModel"], |
| 57 | }, |
| 58 | }, |
| 59 | "hf_comparison": { |
| 60 | "weights_file": "model.safetensors", |
| 61 | "weights_key": "linear.weight", |
| 62 | "trust_remote_code": True, |
| 63 | "model_cls": "AutoModel", |
| 64 | }, |
| 65 | }, |
| 66 | "lfm2": { |
| 67 | "model": "LiquidAI/LFM2-ColBERT-350M", |
| 68 | "colbert_dim": 128, |
| 69 | "max_model_len": 511, |
| 70 | "extra_kwargs": { |
| 71 | "hf_overrides": { |
| 72 | "architectures": ["ColBERTLfm2Model"], |
| 73 | }, |
| 74 | }, |
| 75 | "hf_comparison": { |
| 76 | "weights_file": "1_Dense/model.safetensors", |
| 77 | "weights_key": "linear.weight", |
| 78 | "trust_remote_code": False, |
| 79 | "model_cls": "AutoModel", |
| 80 | }, |
| 81 | }, |
| 82 | } |
| 83 | |
| 84 | |
| 85 | TEXTS_1 = [ |
| 86 | "What is the capital of France?", |
| 87 | "What is the capital of Germany?", |
| 88 | ] |
| 89 | |
| 90 | TEXTS_2 = [ |
| 91 | "The capital of France is Paris.", |
| 92 | "The capital of Germany is Berlin.", |
| 93 | ] |
| 94 | |
| 95 | DTYPE = "half" |
| 96 | |
| 97 | |
| 98 | def _load_hf_model(model_name: str, hf_spec: dict, device: torch.device): |
| 99 | """Load HF model on the given device with a compatible attention impl.""" |
| 100 | from transformers import AutoModel, BertModel |
| 101 | |
| 102 | cls = BertModel if hf_spec["model_cls"] == "BertModel" else AutoModel |
| 103 | trust = hf_spec.get("trust_remote_code", False) |
| 104 | |
| 105 | # Flash / Triton kernels require GPU tensors; fall back to eager on CPU. |
| 106 | extra = {} |
| 107 | if device.type == "cpu": |
| 108 | extra["attn_implementation"] = "eager" |
| 109 | |
| 110 | model = cls.from_pretrained( |
| 111 | model_name, |
| 112 | trust_remote_code=trust, |
| 113 | **extra, |
| 114 | ).to(device) |
| 115 | model.eval() |
| 116 | |
| 117 | # Transformers 5.0 weight materialization can clear non-persistent |
| 118 | # buffers (e.g. rotary inv_freq) that were registered with |
| 119 | # persistent=False. Re-compute them so the model produces valid output. |
| 120 | for mod in model.modules(): |
| 121 | if hasattr(mod, "_compute_inv_freq") and hasattr(mod, "inv_freq"): |
| 122 | mod.inv_freq = mod._compute_inv_freq(device=device) |
| 123 | |
| 124 | return model |
| 125 | |
| 126 | |
| 127 | def _load_projection_weight(model_name: str, hf_spec: dict, device: torch.device): |
| 128 | """Download and return the ColBERT linear projection weight.""" |
| 129 | from safetensors.torch import load_file |
| 130 | |
| 131 | from vllm.transformers_utils.repo_utils import hf_api, with_retry |
| 132 | |
| 133 | path = with_retry( |
| 134 | lambda: hf_api().hf_hub_download(model_name, filename=hf_spec["weights_file"]), |
| 135 | f"Error downloading {hf_spec['weights_file']} from {model_name}", |
| 136 | ) |
| 137 | weights = load_file(path) |
| 138 | return weights[hf_spec["weights_key"]].to(device) |
| 139 | |
| 140 | |
| 141 | def _compute_hf_colbert_embeddings(model, tokenizer, linear_weight, texts, device): |
| 142 | """Run HF model + projection and return L2-normalised token embeddings.""" |
| 143 | import torch.nn.functional as F |
| 144 | |
| 145 | embeddings = [] |
| 146 | for text in texts: |
| 147 | inputs = tokenizer(text, return_tensors="pt").to(device) |
| 148 | with torch.no_grad(): |
| 149 | hidden = model(**inputs).last_hidden_state.float() |
| 150 | projected = F.linear(hidden, linear_weight.float()) |
| 151 | normalised = F.normalize(projected, p=2, dim=-1) |
| 152 | embeddings.append(normalised.squeeze(0).cpu()) |
| 153 | return embeddings |
| 154 | |
| 155 | |
| 156 | @contextmanager |
| 157 | def _hf_colbert_model(model_name: str, hf_spec: dict, device: torch.device): |
| 158 | """Load the HF backbone + ColBERT projection, freeing the GPU on exit. |
| 159 | |
| 160 | These live outside any runner context, so without explicit cleanup ROCm |
| 161 | keeps the VRAM resident and the next backend parametrization (or test) |
| 162 | OOMs on startup. |
| 163 | """ |
| 164 | hf_model = _load_hf_model(model_name, hf_spec, device) |
| 165 | linear_weight = _load_projection_weight(model_name, hf_spec, device) |
| 166 | try: |
| 167 | yield hf_model, linear_weight |
| 168 | finally: |
| 169 | del hf_model, linear_weight |
| 170 | cleanup_dist_env_and_memory() |
| 171 | wait_for_memory_to_settle() |
| 172 | |
| 173 | |
| 174 | def _assert_embeddings_close(vllm_outputs, hf_embeddings): |
| 175 | """Assert that vLLM and HuggingFace embeddings match.""" |
| 176 | for i, (hf_emb, vllm_out) in enumerate(zip(hf_embeddings, vllm_outputs)): |
| 177 | vllm_emb = torch.as_tensor(vllm_out).float() |
| 178 | |
| 179 | assert hf_emb.shape == vllm_emb.shape, ( |
| 180 | f"Shape mismatch for text {i}: HF {hf_emb.shape} vs vLLM {vllm_emb.shape}" |
| 181 | ) |
| 182 | |
| 183 | torch.testing.assert_close( |
| 184 | vllm_emb, |
| 185 | hf_emb, |
| 186 | rtol=1e-2, |
| 187 | atol=1e-2, |
| 188 | msg=f"Embedding mismatch for text {i}", |
| 189 | ) |
| 190 | |
| 191 | |
| 192 | @pytest.fixture(params=list(COLBERT_MODELS.keys()), scope="class") |
| 193 | def colbert_spec(request): |
| 194 | """Return the model spec dict for the current parametrization.""" |
| 195 | return COLBERT_MODELS[request.param] |
| 196 | |
| 197 | |
| 198 | @pytest.fixture(scope="class") |
| 199 | def colbert_model_name(colbert_spec): |
| 200 | return colbert_spec["model"] |
| 201 | |
| 202 | |
| 203 | @pytest.fixture(scope="class") |
| 204 | def colbert_dim(colbert_spec): |
| 205 | return colbert_spec["colbert_dim"] |
| 206 | |
| 207 | |
| 208 | @pytest.fixture(scope="class") |
| 209 | def colbert_max_model_len(colbert_spec): |
| 210 | return colbert_spec["max_model_len"] |
| 211 | |
| 212 | |
| 213 | @pytest.fixture(scope="class") |
| 214 | def colbert_extra_kwargs(colbert_spec): |
| 215 | return colbert_spec["extra_kwargs"] |
| 216 | |
| 217 | |
| 218 | @pytest.fixture(scope="class") |
| 219 | def colbert_model( |
| 220 | vllm_runner, |
| 221 | colbert_model_name, |
| 222 | colbert_max_model_len, |
| 223 | colbert_extra_kwargs, |
| 224 | ): |
| 225 | with vllm_runner( |
| 226 | colbert_model_name, |
| 227 | runner="pooling", |
| 228 | dtype=DTYPE, |
| 229 | max_model_len=colbert_max_model_len, |
| 230 | enforce_eager=True, |
| 231 | **colbert_extra_kwargs, |
| 232 | ) as vllm_model: |
| 233 | yield vllm_model |
| 234 | |
| 235 | |
| 236 | class TestColbertSharedEngine: |
| 237 | """Tests sharing one engine per model. |
| 238 | |
| 239 | Class-scoped so the engine is released before `test_colbert_hf_comparison`, |
| 240 | which needs a runner of its own and cannot start while this one holds VRAM. |
| 241 | """ |
| 242 | |
| 243 | def test_colbert_token_embed(self, colbert_model, colbert_dim): |
| 244 | """Test that ColBERT model produces token embeddings.""" |
| 245 | outputs = colbert_model.token_embed([TEXTS_1[0]]) |
| 246 | |
| 247 | assert len(outputs) == 1 |
| 248 | emb = torch.as_tensor(outputs[0]) |
| 249 | assert emb.dim() == 2 |
| 250 | assert emb.shape[1] == colbert_dim |
| 251 | assert emb.shape[0] > 1 |
| 252 | |
| 253 | def test_colbert_late_interaction_1_to_1(self, colbert_model): |
| 254 | """Test ColBERT late interaction scoring with 1:1 query-document pair.""" |
| 255 | q_outputs = colbert_model.token_embed([TEXTS_1[0]]) |
| 256 | d_outputs = colbert_model.token_embed([TEXTS_2[0]]) |
| 257 | |
| 258 | q_emb = torch.as_tensor(q_outputs[0]) |
| 259 | d_emb = torch.as_tensor(d_outputs[0]) |
| 260 | |
| 261 | manual_score = compute_maxsim_score(q_emb, d_emb).item() |
| 262 | |
| 263 | vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2[0]) |
| 264 | |
| 265 | assert len(vllm_scores) == 1 |
| 266 | assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) |
| 267 | |
| 268 | def test_colbert_late_interaction_1_to_N(self, colbert_model): |
| 269 | """Test ColBERT late interaction scoring with 1:N query-documents.""" |
| 270 | q_outputs = colbert_model.token_embed([TEXTS_1[0]]) |
| 271 | d_outputs = colbert_model.token_embed(TEXTS_2) |
| 272 | |
| 273 | q_emb = torch.as_tensor(q_outputs[0]) |
| 274 | |
| 275 | manual_scores = [] |
| 276 | for d_out in d_outputs: |
| 277 | d_emb = torch.as_tensor(d_out) |
| 278 | manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) |
| 279 | |
| 280 | vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2) |
| 281 | |
| 282 | assert len(vllm_scores) == 2 |
| 283 | for i in range(2): |
| 284 | assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01) |
| 285 | |
| 286 | def test_colbert_late_interaction_N_to_N(self, colbert_model): |
| 287 | """Test ColBERT late interaction scoring with N:N query-documents.""" |
| 288 | q_outputs = colbert_model.token_embed(TEXTS_1) |
| 289 | d_outputs = colbert_model.token_embed(TEXTS_2) |
| 290 | |
| 291 | manual_scores = [] |
| 292 | for q_out, d_out in zip(q_outputs, d_outputs): |
| 293 | q_emb = torch.as_tensor(q_out) |
| 294 | d_emb = torch.as_tensor(d_out) |
| 295 | manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) |
| 296 | |
| 297 | vllm_scores = colbert_model.score(TEXTS_1, TEXTS_2) |
| 298 | |
| 299 | assert len(vllm_scores) == 2 |
| 300 | for i in range(2): |
| 301 | assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01) |
| 302 | |
| 303 | def test_colbert_relevance_ordering(self, colbert_model): |
| 304 | """Test that ColBERT scores relevant documents higher than irrelevant.""" |
| 305 | query = "What is machine learning?" |
| 306 | documents = [ |
| 307 | "Machine learning is a subset of artificial intelligence.", |
| 308 | "Python is a programming language.", |
| 309 | "Deep learning uses neural networks.", |
| 310 | ] |
| 311 | |
| 312 | scores = colbert_model.score(query, documents) |
| 313 | |
| 314 | assert len(scores) == 3 |
| 315 | assert scores[0] > scores[1], "ML doc should score higher than Python doc" |
| 316 | assert scores[2] > scores[1], "DL doc should score higher than Python doc" |
| 317 | |
| 318 | def test_colbert_embed_not_supported(self, colbert_model): |
| 319 | """Test that ColBERT model does not support the embed task.""" |
| 320 | with pytest.raises(ValueError, match="Embedding API is not supported"): |
| 321 | colbert_model.embed([TEXTS_1[0]]) |
| 322 | |
| 323 | |
| 324 | @pytest.mark.parametrize( |
| 325 | ("backend", "use_v2"), |
| 326 | [ |
| 327 | pytest.param("bert", True, id="bert-v2"), |
| 328 | pytest.param("modernbert", True, id="modernbert-v2"), |
| 329 | pytest.param("jina", True, id="jina-v2"), |
| 330 | pytest.param("lfm2", True, id="lfm2-v2"), |
| 331 | ], |
| 332 | ) |
| 333 | def test_colbert_hf_comparison(vllm_runner, monkeypatch, backend, use_v2): |
| 334 | """Test that vLLM ColBERT embeddings match HuggingFace for each backend.""" |
| 335 | from transformers import AutoTokenizer |
| 336 | |
| 337 | spec = COLBERT_MODELS[backend] |
| 338 | hf_spec = spec["hf_comparison"] |
| 339 | extra_kwargs = spec["extra_kwargs"] |
| 340 | model_name = spec["model"] |
| 341 | assert isinstance(model_name, str) |
| 342 | assert isinstance(hf_spec, dict) |
| 343 | assert isinstance(extra_kwargs, dict) |
| 344 | test_texts = [TEXTS_1[0], TEXTS_2[0]] |
| 345 | |
| 346 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2 else "0") |
| 347 | |
| 348 | with vllm_runner( |
| 349 | model_name, |
| 350 | runner="pooling", |
| 351 | dtype="float32", |
| 352 | max_model_len=spec["max_model_len"], |
| 353 | enforce_eager=True, |
| 354 | **extra_kwargs, |
| 355 | ) as vllm_model: |
| 356 | assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner == use_v2 |
| 357 | vllm_outputs = vllm_model.token_embed(test_texts) |
| 358 | |
| 359 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 360 | |
| 361 | hf_tokenizer = AutoTokenizer.from_pretrained( |
| 362 | model_name, |
| 363 | trust_remote_code=hf_spec.get("trust_remote_code", False), |
| 364 | ) |
| 365 | with _hf_colbert_model(model_name, hf_spec, device) as (hf_model, linear_weight): |
| 366 | hf_embeddings = _compute_hf_colbert_embeddings( |
| 367 | hf_model, |
| 368 | hf_tokenizer, |
| 369 | linear_weight, |
| 370 | test_texts, |
| 371 | device, |
| 372 | ) |
| 373 | |
| 374 | _assert_embeddings_close(vllm_outputs, hf_embeddings) |