| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for ColModernVBERT multimodal late-interaction model. |
| 4 | |
| 5 | ColModernVBERT combines SigLIP vision encoder + ModernBERT text encoder |
| 6 | with a pixel shuffle connector and ColBERT-style 128-dim per-token |
| 7 | embeddings for visual document retrieval. |
| 8 | """ |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | |
| 13 | from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score |
| 14 | |
| 15 | MODEL_NAME = "ModernVBERT/colmodernvbert-merged" |
| 16 | COLBERT_DIM = 128 |
| 17 | DTYPE = "half" |
| 18 | |
| 19 | |
| 20 | @pytest.fixture(scope="module") |
| 21 | def colmodernvbert_model(vllm_runner): |
| 22 | with vllm_runner( |
| 23 | MODEL_NAME, |
| 24 | runner="pooling", |
| 25 | dtype=DTYPE, |
| 26 | enforce_eager=True, |
| 27 | ) as vllm_model: |
| 28 | yield vllm_model |
| 29 | |
| 30 | |
| 31 | # ----------------------------------------------------------------------- |
| 32 | # Text-only tests |
| 33 | # ----------------------------------------------------------------------- |
| 34 | |
| 35 | |
| 36 | def test_colmodernvbert_text_token_embed(colmodernvbert_model): |
| 37 | """Text query produces per-token embeddings with shape (seq_len, 128).""" |
| 38 | outputs = colmodernvbert_model.token_embed(["What is machine learning?"]) |
| 39 | |
| 40 | assert len(outputs) == 1 |
| 41 | emb = torch.tensor(outputs[0]) |
| 42 | assert emb.dim() == 2 |
| 43 | assert emb.shape[1] == COLBERT_DIM |
| 44 | assert emb.shape[0] > 1 |
| 45 | |
| 46 | |
| 47 | def test_colmodernvbert_text_relevance_ordering(colmodernvbert_model): |
| 48 | """Relevant documents score higher than irrelevant ones.""" |
| 49 | query = "What is machine learning?" |
| 50 | documents = [ |
| 51 | "Machine learning is a subset of artificial intelligence.", |
| 52 | "The weather in Paris is mild in spring.", |
| 53 | ] |
| 54 | |
| 55 | scores = colmodernvbert_model.score(query, documents) |
| 56 | |
| 57 | assert len(scores) == 2 |
| 58 | assert scores[0] > scores[1], "ML doc should score higher than weather doc" |
| 59 | |
| 60 | |
| 61 | def test_colmodernvbert_text_late_interaction(colmodernvbert_model): |
| 62 | """MaxSim scoring via vLLM matches manual computation.""" |
| 63 | query = "What is the capital of France?" |
| 64 | doc = "The capital of France is Paris." |
| 65 | |
| 66 | q_out = colmodernvbert_model.token_embed([query]) |
| 67 | d_out = colmodernvbert_model.token_embed([doc]) |
| 68 | |
| 69 | q_emb = torch.tensor(q_out[0]) |
| 70 | d_emb = torch.tensor(d_out[0]) |
| 71 | manual_score = compute_maxsim_score(q_emb, d_emb).item() |
| 72 | |
| 73 | vllm_scores = colmodernvbert_model.score(query, doc) |
| 74 | |
| 75 | assert len(vllm_scores) == 1 |
| 76 | assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) |
| 77 | |
| 78 | |
| 79 | # ----------------------------------------------------------------------- |
| 80 | # Image tests |
| 81 | # ----------------------------------------------------------------------- |
| 82 | |
| 83 | |
| 84 | def test_colmodernvbert_image_token_embed(colmodernvbert_model, image_assets): |
| 85 | """Image input produces per-token embeddings including vision tokens.""" |
| 86 | image = image_assets[0].pil_image |
| 87 | inputs = colmodernvbert_model.get_inputs( |
| 88 | [""], |
| 89 | images=[image], |
| 90 | ) |
| 91 | req_outputs = colmodernvbert_model.llm.encode( |
| 92 | inputs, |
| 93 | pooling_task="token_embed", |
| 94 | ) |
| 95 | outputs = [req_output.outputs.data for req_output in req_outputs] |
| 96 | |
| 97 | assert len(outputs) == 1 |
| 98 | emb = torch.tensor(outputs[0]) |
| 99 | assert emb.dim() == 2 |
| 100 | assert emb.shape[1] == COLBERT_DIM |
| 101 | # Should have at least the image tokens (64 after pixel shuffle) |
| 102 | assert emb.shape[0] >= 64 |