[CI/Build] Reduce more duplicate runner startup in tests (#52570)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Isotr0py committed
Aug 17, 2026 at 13:19 UTC
53e211d2923192ea1f7442de93d59cfbb150e089
10 files changed
+586
-855
tests/models/language/pooling/test_colbert.py
+58
-108
@@ -212,14 +212,13 @@ def colbert_extra_kwargs(colbert_spec):
212
return colbert_spec["extra_kwargs"]
213
214
215
-def test_colbert_token_embed(
215
+@pytest.fixture(scope="module")
216
+def colbert_model(
217
vllm_runner,
218
colbert_model_name,
218
- colbert_dim,
219
colbert_max_model_len,
220
colbert_extra_kwargs,
221
):
222
- """Test that ColBERT model produces token embeddings."""
222
with vllm_runner(
223
colbert_model_name,
224
runner="pooling",
@@ -228,112 +227,84 @@ def test_colbert_token_embed(
227
enforce_eager=True,
228
**colbert_extra_kwargs,
229
) as vllm_model:
231
- outputs = vllm_model.token_embed([TEXTS_1[0]])
230
+ yield vllm_model
231
+
232
+
233
+def test_colbert_token_embed(
234
+ colbert_model,
235
+ colbert_dim,
236
+):
237
+ """Test that ColBERT model produces token embeddings."""
238
+ outputs = colbert_model.token_embed([TEXTS_1[0]])
239
233
- assert len(outputs) == 1
234
- emb = torch.as_tensor(outputs[0])
235
- assert emb.dim() == 2
236
- assert emb.shape[1] == colbert_dim
237
- assert emb.shape[0] > 1
240
+ assert len(outputs) == 1
241
+ emb = torch.as_tensor(outputs[0])
242
+ assert emb.dim() == 2
243
+ assert emb.shape[1] == colbert_dim
244
+ assert emb.shape[0] > 1
245
246
247
def test_colbert_late_interaction_1_to_1(
241
- vllm_runner,
242
- colbert_model_name,
243
- colbert_max_model_len,
244
- colbert_extra_kwargs,
248
+ colbert_model,
249
):
250
"""Test ColBERT late interaction scoring with 1:1 query-document pair."""
247
- with vllm_runner(
248
- colbert_model_name,
249
- runner="pooling",
250
- dtype=DTYPE,
251
- max_model_len=colbert_max_model_len,
252
- enforce_eager=True,
253
- **colbert_extra_kwargs,
254
- ) as vllm_model:
255
- q_outputs = vllm_model.token_embed([TEXTS_1[0]])
256
- d_outputs = vllm_model.token_embed([TEXTS_2[0]])
251
+ q_outputs = colbert_model.token_embed([TEXTS_1[0]])
252
+ d_outputs = colbert_model.token_embed([TEXTS_2[0]])
253
258
- q_emb = torch.as_tensor(q_outputs[0])
259
- d_emb = torch.as_tensor(d_outputs[0])
254
+ q_emb = torch.as_tensor(q_outputs[0])
255
+ d_emb = torch.as_tensor(d_outputs[0])
256
261
- manual_score = compute_maxsim_score(q_emb, d_emb).item()
257
+ manual_score = compute_maxsim_score(q_emb, d_emb).item()
258
263
- vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2[0])
259
+ vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2[0])
260
265
- assert len(vllm_scores) == 1
266
- assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
261
+ assert len(vllm_scores) == 1
262
+ assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
263
264
265
def test_colbert_late_interaction_1_to_N(
270
- vllm_runner,
271
- colbert_model_name,
272
- colbert_max_model_len,
273
- colbert_extra_kwargs,
266
+ colbert_model,
267
):
268
"""Test ColBERT late interaction scoring with 1:N query-documents."""
276
- with vllm_runner(
277
- colbert_model_name,
278
- runner="pooling",
279
- dtype=DTYPE,
280
- max_model_len=colbert_max_model_len,
281
- enforce_eager=True,
282
- **colbert_extra_kwargs,
283
- ) as vllm_model:
284
- q_outputs = vllm_model.token_embed([TEXTS_1[0]])
285
- d_outputs = vllm_model.token_embed(TEXTS_2)
269
+ q_outputs = colbert_model.token_embed([TEXTS_1[0]])
270
+ d_outputs = colbert_model.token_embed(TEXTS_2)
271
287
- q_emb = torch.as_tensor(q_outputs[0])
272
+ q_emb = torch.as_tensor(q_outputs[0])
273
289
- manual_scores = []
290
- for d_out in d_outputs:
291
- d_emb = torch.as_tensor(d_out)
292
- manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
274
+ manual_scores = []
275
+ for d_out in d_outputs:
276
+ d_emb = torch.as_tensor(d_out)
277
+ manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
278
294
- vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2)
279
+ vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2)
280
296
- assert len(vllm_scores) == 2
297
- for i in range(2):
298
- assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
281
+ assert len(vllm_scores) == 2
282
+ for i in range(2):
283
+ assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
284
285
286
def test_colbert_late_interaction_N_to_N(
302
- vllm_runner,
303
- colbert_model_name,
304
- colbert_max_model_len,
305
- colbert_extra_kwargs,
287
+ colbert_model,
288
):
289
"""Test ColBERT late interaction scoring with N:N query-documents."""
308
- with vllm_runner(
309
- colbert_model_name,
310
- runner="pooling",
311
- dtype=DTYPE,
312
- max_model_len=colbert_max_model_len,
313
- enforce_eager=True,
314
- **colbert_extra_kwargs,
315
- ) as vllm_model:
316
- q_outputs = vllm_model.token_embed(TEXTS_1)
317
- d_outputs = vllm_model.token_embed(TEXTS_2)
290
+ q_outputs = colbert_model.token_embed(TEXTS_1)
291
+ d_outputs = colbert_model.token_embed(TEXTS_2)
292
319
- manual_scores = []
320
- for q_out, d_out in zip(q_outputs, d_outputs):
321
- q_emb = torch.as_tensor(q_out)
322
- d_emb = torch.as_tensor(d_out)
323
- manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
293
+ manual_scores = []
294
+ for q_out, d_out in zip(q_outputs, d_outputs):
295
+ q_emb = torch.as_tensor(q_out)
296
+ d_emb = torch.as_tensor(d_out)
297
+ manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
298
325
- vllm_scores = vllm_model.score(TEXTS_1, TEXTS_2)
299
+ vllm_scores = colbert_model.score(TEXTS_1, TEXTS_2)
300
327
- assert len(vllm_scores) == 2
328
- for i in range(2):
329
- assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
301
+ assert len(vllm_scores) == 2
302
+ for i in range(2):
303
+ assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
304
305
306
def test_colbert_relevance_ordering(
333
- vllm_runner,
334
- colbert_model_name,
335
- colbert_max_model_len,
336
- colbert_extra_kwargs,
307
+ colbert_model,
308
):
309
"""Test that ColBERT scores relevant documents higher than irrelevant."""
310
query = "What is machine learning?"
@@ -343,40 +314,19 @@ def test_colbert_relevance_ordering(
314
"Deep learning uses neural networks.",
315
]
316
346
- with vllm_runner(
347
- colbert_model_name,
348
- runner="pooling",
349
- dtype=DTYPE,
350
- max_model_len=colbert_max_model_len,
351
- enforce_eager=True,
352
- **colbert_extra_kwargs,
353
- ) as vllm_model:
354
- scores = vllm_model.score(query, documents)
317
+ scores = colbert_model.score(query, documents)
318
356
- assert len(scores) == 3
357
- assert scores[0] > scores[1], "ML doc should score higher than Python doc"
358
- assert scores[2] > scores[1], "DL doc should score higher than Python doc"
319
+ assert len(scores) == 3
320
+ assert scores[0] > scores[1], "ML doc should score higher than Python doc"
321
+ assert scores[2] > scores[1], "DL doc should score higher than Python doc"
322
323
324
def test_colbert_embed_not_supported(
362
- vllm_runner,
363
- colbert_model_name,
364
- colbert_max_model_len,
365
- colbert_extra_kwargs,
325
+ colbert_model,
326
):
367
- """Test that ColBERT model does not support 'embed' task."""
368
- with (
369
- vllm_runner(
370
- colbert_model_name,
371
- runner="pooling",
372
- dtype=DTYPE,
373
- max_model_len=colbert_max_model_len,
374
- enforce_eager=True,
375
- **colbert_extra_kwargs,
376
- ) as vllm_model,
377
- pytest.raises(ValueError, match="Embedding API is not supported"),
378
- ):
379
- vllm_model.embed([TEXTS_1[0]])
327
+ """Test that ColBERT model does not support the embed task."""
328
+ with pytest.raises(ValueError, match="Embedding API is not supported"):
329
+ colbert_model.embed([TEXTS_1[0]])
330
331
332
@pytest.mark.parametrize(
tests/models/language/pooling/test_truncation_control.py
+21
-34
@@ -22,60 +22,47 @@ calculus, each contributing unique perspectives that would shape this new
22
field."""
23
24
25
+@pytest.fixture(scope="module")
26
+def vllm_model(vllm_runner):
27
+ with vllm_runner(
28
+ MODEL_NAME, runner="pooling", max_model_len=max_model_len
29
+ ) as model:
30
+ yield model
31
+
32
+
33
def test_smaller_truncation_size(
26
- vllm_runner, model_name=MODEL_NAME, input_str=input_str
34
+ vllm_model,
35
):
36
truncate_prompt_tokens = 10
37
30
- with vllm_runner(
31
- model_name, runner="pooling", max_model_len=max_model_len
32
- ) as vllm_model:
33
- vllm_output = vllm_model.llm.embed(
34
- input_str,
35
- tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens),
36
- )
38
+ vllm_output = vllm_model.llm.embed(
39
+ input_str,
40
+ tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens),
41
+ )
42
43
prompt_tokens = vllm_output[0].prompt_token_ids
44
45
assert len(prompt_tokens) == truncate_prompt_tokens
46
47
43
-def test_max_truncation_size(vllm_runner, model_name=MODEL_NAME, input_str=input_str):
48
+def test_max_truncation_size(vllm_model):
49
truncate_prompt_tokens = -1
50
46
- with vllm_runner(
47
- model_name, runner="pooling", max_model_len=max_model_len
48
- ) as vllm_model:
49
- vllm_output = vllm_model.llm.embed(
50
- input_str,
51
- tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens),
52
- )
51
+ vllm_output = vllm_model.llm.embed(
52
+ input_str,
53
+ tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens),
54
+ )
55
56
prompt_tokens = vllm_output[0].prompt_token_ids
57
58
assert len(prompt_tokens) == max_model_len
59
60
59
-def test_bigger_truncation_size(
60
- vllm_runner, model_name=MODEL_NAME, input_str=input_str
61
-):
61
+def test_bigger_truncation_size(vllm_model):
62
truncate_prompt_tokens = max_model_len + 1
63
64
- with (
65
- pytest.raises(VLLMValidationError),
66
- vllm_runner(
67
- model_name, runner="pooling", max_model_len=max_model_len
68
- ) as vllm_model,
69
- ):
70
- llm_output = vllm_model.llm.embed(
64
+ with pytest.raises(VLLMValidationError):
65
+ vllm_model.llm.embed(
66
input_str,
67
tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens),
68
)
74
-
75
- assert (
76
- llm_output
77
- == f"""truncate_prompt_tokens value
78
- ({truncate_prompt_tokens}) is greater than
79
- max_model_len ({max_model_len}). Please, select
80
- a smaller truncation size."""
81
- )
tests/models/multimodal/generation/test_whisper.py
+49
-39
@@ -20,6 +20,7 @@ VLLM_PROMPT = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>"
20
HF_PROMPT = ""
21
# Whisper expects 16kHz audio
22
WHISPER_SAMPLE_RATE = 16000
23
+BEAM_WIDTHS = (1, 2)
24
25
26
@pytest.fixture(autouse=True)
@@ -127,13 +128,11 @@ def check_model_available(model: str) -> None:
128
129
@pytest.mark.parametrize("dtype", ["half"])
130
@pytest.mark.parametrize("max_tokens", [64])
130
-@pytest.mark.parametrize("beam_width", [1, 2])
131
def test_beam_search_encoder_decoder(
132
hf_runner,
133
vllm_runner,
134
dtype: str,
135
max_tokens: int,
136
- beam_width: int,
136
resampled_assets,
137
) -> None:
138
"""Test beam search with encoder-decoder models (Whisper)."""
@@ -146,12 +145,15 @@ def test_beam_search_encoder_decoder(
145
]
146
147
with hf_runner(model, dtype=dtype, auto_cls=AutoModelForSpeechSeq2Seq) as hf_model:
149
- hf_outputs = hf_model.generate_beam_search(
150
- hf_prompts,
151
- beam_width=beam_width,
152
- max_tokens=max_tokens,
153
- audios=resampled_assets,
154
- )
148
+ hf_outputs_by_beam_width = [
149
+ hf_model.generate_beam_search(
150
+ hf_prompts,
151
+ beam_width=beam_width,
152
+ max_tokens=max_tokens,
153
+ audios=resampled_assets,
154
+ )
155
+ for beam_width in BEAM_WIDTHS
156
+ ]
157
158
# Test both explicit encoder/decoder prompts
159
vllm_prompts = [
@@ -179,38 +181,46 @@ def test_beam_search_encoder_decoder(
181
limit_mm_per_prompt={"audio": 2},
182
enforce_eager=True,
183
) as vllm_model:
182
- vllm_outputs = vllm_model.generate_beam_search(
183
- vllm_prompts,
184
- beam_width=beam_width,
185
- max_tokens=max_tokens,
186
- )
187
-
188
- for i in range(len(vllm_prompts)):
189
- hf_output_ids, hf_output_texts = hf_outputs[i]
190
- vllm_output_ids, vllm_output_texts = vllm_outputs[i]
191
-
192
- for j, (hf_text, vllm_text) in enumerate(
193
- zip(hf_output_texts, vllm_output_texts)
194
- ):
195
- print(f">>>{j}-th hf output [NOTE: special tokens are filtered]:")
196
- print(hf_text)
197
- print(f">>>{j}-th vllm output:")
198
- print(vllm_text)
199
-
200
- # Check that we got the same number of beams
201
- assert len(hf_output_ids) == len(vllm_output_ids)
202
-
203
- # For encoder-decoder models, we primarily want to verify that:
204
- # 1. Beam search completes without errors
205
- # 2. We get the expected number of beams
206
- # 3. Outputs are reasonable (non-empty, diverse beams)
207
- for j in range(len(vllm_output_ids)):
208
- # Check that outputs are not empty
209
- assert len(vllm_output_ids[j]) > 0, f"Prompt {i}, beam {j}: empty output"
210
- # Check that decoded text is not empty
211
- assert len(vllm_output_texts[j].strip()) > 0, (
212
- f"Prompt {i}, beam {j}: empty text output"
184
+ vllm_outputs_by_beam_width = [
185
+ vllm_model.generate_beam_search(
186
+ vllm_prompts,
187
+ beam_width=beam_width,
188
+ max_tokens=max_tokens,
189
)
190
+ for beam_width in BEAM_WIDTHS
191
+ ]
192
+
193
+ for beam_width, hf_outputs, vllm_outputs in zip(
194
+ BEAM_WIDTHS, hf_outputs_by_beam_width, vllm_outputs_by_beam_width
195
+ ):
196
+ for i in range(len(vllm_prompts)):
197
+ hf_output_ids, hf_output_texts = hf_outputs[i]
198
+ vllm_output_ids, vllm_output_texts = vllm_outputs[i]
199
+
200
+ for j, (hf_text, vllm_text) in enumerate(
201
+ zip(hf_output_texts, vllm_output_texts)
202
+ ):
203
+ print(f">>>{j}-th hf output [NOTE: special tokens are filtered]:")
204
+ print(hf_text)
205
+ print(f">>>{j}-th vllm output:")
206
+ print(vllm_text)
207
+
208
+ # Check that we got the same number of beams
209
+ assert len(hf_output_ids) == len(vllm_output_ids) == beam_width
210
+
211
+ # For encoder-decoder models, we primarily want to verify that:
212
+ # 1. Beam search completes without errors
213
+ # 2. We get the expected number of beams
214
+ # 3. Outputs are reasonable (non-empty, diverse beams)
215
+ for j in range(len(vllm_output_ids)):
216
+ # Check that outputs are not empty
217
+ assert len(vllm_output_ids[j]) > 0, (
218
+ f"Prompt {i}, beam {j}: empty output"
219
+ )
220
+ # Check that decoded text is not empty
221
+ assert len(vllm_output_texts[j].strip()) > 0, (
222
+ f"Prompt {i}, beam {j}: empty text output"
223
+ )
224
225
226
def test_parse_language_detection_output():
tests/models/multimodal/pooling/test_clip.py
+52
-83
@@ -26,8 +26,7 @@ MODELS = ["openai/clip-vit-base-patch32"]
26
def _run_test(
27
hf_runner: type[HfRunner],
28
vllm_runner: type[VllmRunner],
29
- input_texts: list[str],
30
- input_images: PromptImageInput,
29
+ input_cases: list[tuple[list[str], PromptImageInput]],
30
model: str,
31
*,
32
dtype: str,
@@ -39,105 +38,75 @@ def _run_test(
38
with vllm_runner(
39
model, runner="pooling", dtype=dtype, enforce_eager=True, max_model_len=77
40
) as vllm_model:
42
- vllm_outputs = vllm_model.embed(input_texts, images=input_images)
43
-
44
- with hf_runner(model, dtype=dtype, auto_cls=CLIPModel) as hf_model:
45
- all_inputs = hf_model.get_inputs(input_texts, images=input_images)
46
-
47
- all_outputs = []
48
- for inputs in all_inputs:
49
- inputs = hf_model.wrap_device(inputs)
50
-
51
- if "pixel_values" in inputs:
52
- pooled_output = hf_model.model.get_image_features(
53
- pixel_values=inputs.pixel_values,
54
- )
55
- else:
56
- pooled_output = hf_model.model.get_text_features(
57
- input_ids=inputs.input_ids,
58
- attention_mask=inputs.attention_mask,
59
- )
60
-
61
- if not isinstance(pooled_output, torch.Tensor):
62
- pooled_output = pooled_output.pooler_output
63
- pooled_output = pooled_output.squeeze(0)
64
- all_outputs.append(pooled_output.tolist())
65
-
66
- hf_outputs = all_outputs
67
-
68
- check_embeddings_close(
69
- embeddings_0_lst=hf_outputs,
70
- embeddings_1_lst=vllm_outputs,
71
- name_0="hf",
72
- name_1="vllm",
73
- )
41
+ vllm_outputs_per_case = [
42
+ vllm_model.embed(input_texts, images=input_images)
43
+ for input_texts, input_images in input_cases
44
+ ]
45
46
+ texts = [HF_TEXT_PROMPTS[0]]
47
+ images = [input_cases[1][1][0]]
48
+ with pytest.raises(ValueError, match="not both"):
49
+ vllm_model.embed(texts, images=images)
50
76
-@pytest.mark.parametrize("model", MODELS)
77
-@pytest.mark.parametrize("dtype", ["float"])
78
-def test_models_text(
79
- hf_runner,
80
- vllm_runner,
81
- image_assets,
82
- model: str,
83
- dtype: str,
84
-) -> None:
85
- input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS]
86
- input_texts = [text for text, _ in input_texts_images]
87
- input_images = [image for _, image in input_texts_images]
51
+ # Should still be able to run subsequent requests
52
+ vllm_model.embed(texts)
53
+ vllm_model.embed([""], images=images)
54
89
- _run_test(
90
- hf_runner,
91
- vllm_runner,
92
- input_texts,
93
- input_images, # type: ignore
94
- model,
95
- dtype=dtype,
96
- )
55
+ with hf_runner(model, dtype=dtype, auto_cls=CLIPModel) as hf_model:
56
+ hf_outputs_per_case = []
57
+ for input_texts, input_images in input_cases:
58
+ all_inputs = hf_model.get_inputs(input_texts, images=input_images)
59
+
60
+ hf_outputs = []
61
+ for inputs in all_inputs:
62
+ inputs = hf_model.wrap_device(inputs)
63
+
64
+ if "pixel_values" in inputs:
65
+ pooled_output = hf_model.model.get_image_features(
66
+ pixel_values=inputs.pixel_values,
67
+ )
68
+ else:
69
+ pooled_output = hf_model.model.get_text_features(
70
+ input_ids=inputs.input_ids,
71
+ attention_mask=inputs.attention_mask,
72
+ )
73
+
74
+ if not isinstance(pooled_output, torch.Tensor):
75
+ pooled_output = pooled_output.pooler_output
76
+ pooled_output = pooled_output.squeeze(0)
77
+ hf_outputs.append(pooled_output.tolist())
78
+
79
+ hf_outputs_per_case.append(hf_outputs)
80
+
81
+ for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
82
+ check_embeddings_close(
83
+ embeddings_0_lst=hf_outputs,
84
+ embeddings_1_lst=vllm_outputs,
85
+ name_0="hf",
86
+ name_1="vllm",
87
+ )
88
89
90
@pytest.mark.parametrize("model", MODELS)
91
@pytest.mark.parametrize("dtype", ["float"])
101
-def test_models_image(
92
+def test_models(
93
hf_runner,
94
vllm_runner,
95
image_assets,
96
model: str,
97
dtype: str,
98
) -> None:
108
- input_texts_images = [
109
- (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
99
+ text_images = [None] * len(HF_TEXT_PROMPTS)
100
+ images = [asset.pil_image for asset in image_assets]
101
+ input_cases = [
102
+ (HF_TEXT_PROMPTS, text_images),
103
+ (HF_IMAGE_PROMPTS, images),
104
]
111
- input_texts = [text for text, _ in input_texts_images]
112
- input_images = [image for _, image in input_texts_images]
105
106
_run_test(
107
hf_runner,
108
vllm_runner,
117
- input_texts,
118
- input_images,
109
+ input_cases, # type: ignore[arg-type]
110
model,
111
dtype=dtype,
112
)
122
-
123
-
124
-@pytest.mark.parametrize("model", MODELS)
125
-@pytest.mark.parametrize("dtype", ["float"])
126
-def test_models_text_image_no_crash(
127
- vllm_runner,
128
- image_assets,
129
- model: str,
130
- dtype: str,
131
-) -> None:
132
- texts = [HF_TEXT_PROMPTS[0]]
133
- images = [image_assets[0].pil_image]
134
-
135
- with vllm_runner(
136
- model, runner="pooling", dtype=dtype, enforce_eager=True, max_model_len=77
137
- ) as vllm_model:
138
- with pytest.raises(ValueError, match="not both"):
139
- vllm_model.embed(texts, images=images)
140
-
141
- # Should still be able to run subsequent requests
142
- vllm_model.embed(texts)
143
- vllm_model.embed([""], images=images)
tests/models/multimodal/pooling/test_colmodernvbert.py
+49
-62
@@ -17,29 +17,34 @@ COLBERT_DIM = 128
17
DTYPE = "half"
18
19
20
-# -----------------------------------------------------------------------
21
-# Text-only tests
22
-# -----------------------------------------------------------------------
23
-
24
-
25
-def test_colmodernvbert_text_token_embed(vllm_runner):
26
- """Text query produces per-token embeddings with shape (seq_len, 128)."""
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:
33
- outputs = vllm_model.token_embed(["What is machine learning?"])
28
+ yield vllm_model
29
+
30
+
31
+# -----------------------------------------------------------------------
32
+# Text-only tests
33
+# -----------------------------------------------------------------------
34
+
35
35
- assert len(outputs) == 1
36
- emb = torch.tensor(outputs[0])
37
- assert emb.dim() == 2
38
- assert emb.shape[1] == COLBERT_DIM
39
- assert emb.shape[0] > 1
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
42
-def test_colmodernvbert_text_relevance_ordering(vllm_runner):
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 = [
@@ -47,40 +52,28 @@ def test_colmodernvbert_text_relevance_ordering(vllm_runner):
52
"The weather in Paris is mild in spring.",
53
]
54
50
- with vllm_runner(
51
- MODEL_NAME,
52
- runner="pooling",
53
- dtype=DTYPE,
54
- enforce_eager=True,
55
- ) as vllm_model:
56
- scores = vllm_model.score(query, documents)
55
+ scores = colmodernvbert_model.score(query, documents)
56
58
- assert len(scores) == 2
59
- assert scores[0] > scores[1], "ML doc should score higher than weather doc"
57
+ assert len(scores) == 2
58
+ assert scores[0] > scores[1], "ML doc should score higher than weather doc"
59
60
62
-def test_colmodernvbert_text_late_interaction(vllm_runner):
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
67
- with vllm_runner(
68
- MODEL_NAME,
69
- runner="pooling",
70
- dtype=DTYPE,
71
- enforce_eager=True,
72
- ) as vllm_model:
73
- q_out = vllm_model.token_embed([query])
74
- d_out = vllm_model.token_embed([doc])
66
+ q_out = colmodernvbert_model.token_embed([query])
67
+ d_out = colmodernvbert_model.token_embed([doc])
68
76
- q_emb = torch.tensor(q_out[0])
77
- d_emb = torch.tensor(d_out[0])
78
- manual_score = compute_maxsim_score(q_emb, d_emb).item()
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
80
- vllm_scores = vllm_model.score(query, doc)
73
+ vllm_scores = colmodernvbert_model.score(query, doc)
74
82
- assert len(vllm_scores) == 1
83
- assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
75
+ assert len(vllm_scores) == 1
76
+ assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
77
78
79
# -----------------------------------------------------------------------
@@ -88,28 +81,22 @@ def test_colmodernvbert_text_late_interaction(vllm_runner):
81
# -----------------------------------------------------------------------
82
83
91
-def test_colmodernvbert_image_token_embed(vllm_runner, image_assets):
84
+def test_colmodernvbert_image_token_embed(colmodernvbert_model, image_assets):
85
"""Image input produces per-token embeddings including vision tokens."""
93
- with vllm_runner(
94
- MODEL_NAME,
95
- runner="pooling",
96
- dtype=DTYPE,
97
- enforce_eager=True,
98
- ) as vllm_model:
99
- image = image_assets[0].pil_image
100
- inputs = vllm_model.get_inputs(
101
- [""],
102
- images=[image],
103
- )
104
- req_outputs = vllm_model.llm.encode(
105
- inputs,
106
- pooling_task="token_embed",
107
- )
108
- outputs = [req_output.outputs.data for req_output in req_outputs]
109
-
110
- assert len(outputs) == 1
111
- emb = torch.tensor(outputs[0])
112
- assert emb.dim() == 2
113
- assert emb.shape[1] == COLBERT_DIM
114
- # Should have at least the image tokens (64 after pixel shuffle)
115
- assert emb.shape[0] >= 64
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
tests/models/multimodal/pooling/test_colpali.py
+77
-162
@@ -75,75 +75,51 @@ def _make_image_mm_param(
75
76
77
def _run_token_embed_test(
78
- vllm_runner: type[VllmRunner],
78
+ vllm_model: VllmRunner,
79
model: str,
80
- *,
81
- dtype: str,
80
) -> None:
81
"""Verify per-token embedding shape and L2 normalization."""
84
- with vllm_runner(
85
- model,
86
- runner="pooling",
87
- dtype=dtype,
88
- max_model_len=4096,
89
- enforce_eager=True,
90
- gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
91
- ) as vllm_model:
92
- outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
93
-
94
- assert len(outputs) == 1
95
- emb = torch.tensor(outputs[0])
96
- # Token embeddings should be 2D: [num_tokens, embed_dim]
97
- assert emb.dim() == 2
98
- assert emb.shape[1] == EMBED_DIMS[model]
99
- assert emb.shape[0] > 1
100
-
101
- # Verify L2 normalization
102
- norms = torch.norm(emb, p=2, dim=-1)
103
- torch.testing.assert_close(
104
- norms,
105
- torch.ones_like(norms),
106
- rtol=1e-2,
107
- atol=1e-2,
108
- )
82
+ outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
83
+
84
+ assert len(outputs) == 1
85
+ emb = torch.tensor(outputs[0])
86
+ # Token embeddings should be 2D: [num_tokens, embed_dim]
87
+ assert emb.dim() == 2
88
+ assert emb.shape[1] == EMBED_DIMS[model]
89
+ assert emb.shape[0] > 1
90
+
91
+ # Verify L2 normalization
92
+ norms = torch.norm(emb, p=2, dim=-1)
93
+ torch.testing.assert_close(
94
+ norms,
95
+ torch.ones_like(norms),
96
+ rtol=1e-2,
97
+ atol=1e-2,
98
+ )
99
100
101
def _run_late_interaction_test(
112
- vllm_runner: type[VllmRunner],
113
- model: str,
114
- *,
115
- dtype: str,
102
+ vllm_model: VllmRunner,
103
) -> None:
104
"""Verify MaxSim scoring matches manual computation."""
105
from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
106
120
- with vllm_runner(
121
- model,
122
- runner="pooling",
123
- dtype=dtype,
124
- max_model_len=4096,
125
- enforce_eager=True,
126
- gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
127
- ) as vllm_model:
128
- q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
129
- d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
107
+ q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
108
+ d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
109
131
- q_emb = torch.tensor(q_outputs[0])
132
- d_emb = torch.tensor(d_outputs[0])
110
+ q_emb = torch.tensor(q_outputs[0])
111
+ d_emb = torch.tensor(d_outputs[0])
112
134
- manual_score = compute_maxsim_score(q_emb, d_emb).item()
113
+ manual_score = compute_maxsim_score(q_emb, d_emb).item()
114
136
- vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0])
115
+ vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0])
116
138
- assert len(vllm_scores) == 1
139
- assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
117
+ assert len(vllm_scores) == 1
118
+ assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
119
120
121
def _run_relevance_test(
143
- vllm_runner: type[VllmRunner],
144
- model: str,
145
- *,
146
- dtype: str,
122
+ vllm_model: VllmRunner,
123
) -> None:
124
"""Verify that relevant documents score higher than irrelevant ones."""
125
query = "What is machine learning?"
@@ -153,59 +129,18 @@ def _run_relevance_test(
129
"Deep learning uses neural networks for complex tasks.",
130
]
131
156
- with vllm_runner(
157
- model,
158
- runner="pooling",
159
- dtype=dtype,
160
- max_model_len=4096,
161
- enforce_eager=True,
162
- gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
163
- ) as vllm_model:
164
- scores = vllm_model.score(query, documents)
132
+ scores = vllm_model.score(query, documents)
133
166
- assert len(scores) == 3
167
- assert scores[0] > scores[1], "ML doc should score higher than weather doc"
168
- assert scores[2] > scores[1], "DL doc should score higher than weather doc"
169
-
170
-
171
-@pytest.mark.parametrize("model", MODELS)
172
-@pytest.mark.parametrize("dtype", [DTYPE])
173
-def test_colpali_token_embed(
174
- vllm_runner,
175
- model: str,
176
- dtype: str,
177
-) -> None:
178
- _run_token_embed_test(vllm_runner, model, dtype=dtype)
179
-
180
-
181
-@pytest.mark.parametrize("model", MODELS)
182
-@pytest.mark.parametrize("dtype", [DTYPE])
183
-def test_colpali_late_interaction_scoring(
184
- vllm_runner,
185
- model: str,
186
- dtype: str,
187
-) -> None:
188
- _run_late_interaction_test(vllm_runner, model, dtype=dtype)
189
-
190
-
191
-@pytest.mark.parametrize("model", MODELS)
192
-@pytest.mark.parametrize("dtype", [DTYPE])
193
-def test_colpali_relevance_ordering(
194
- vllm_runner,
195
- model: str,
196
- dtype: str,
197
-) -> None:
198
- _run_relevance_test(vllm_runner, model, dtype=dtype)
134
+ assert len(scores) == 3
135
+ assert scores[0] > scores[1], "ML doc should score higher than weather doc"
136
+ assert scores[2] > scores[1], "DL doc should score higher than weather doc"
137
138
139
# ── Multimodal scoring tests ────────────────────────────────
140
141
142
def _run_multimodal_text_query_image_docs_test(
205
- vllm_runner: type[VllmRunner],
206
- model: str,
207
- *,
208
- dtype: str,
143
+ vllm_model: VllmRunner,
144
) -> None:
145
"""Score a text query against image documents via the multimodal path."""
146
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
@@ -218,29 +153,15 @@ def _run_multimodal_text_query_image_docs_test(
153
]
154
attention_backend = "FLASH_ATTN" if current_platform.is_cuda() else None
155
221
- with vllm_runner(
222
- model,
223
- runner="pooling",
224
- dtype=dtype,
225
- max_model_len=4096,
226
- enforce_eager=True,
227
- gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
228
- attention_backend=attention_backend,
229
- kernel_config={"enable_flashinfer_autotune": False},
230
- ) as vllm_model:
231
- assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
232
- scores = vllm_model.llm.score(query, image_docs)
156
+ scores = vllm_model.llm.score(query, image_docs)
157
234
- assert len(scores) == 2
235
- for s in scores:
236
- assert isinstance(s.outputs.score, float)
158
+ assert len(scores) == 2
159
+ for s in scores:
160
+ assert isinstance(s.outputs.score, float)
161
162
163
def _run_multimodal_mixed_docs_test(
240
- vllm_runner: type[VllmRunner],
241
- model: str,
242
- *,
243
- dtype: str,
164
+ vllm_model: VllmRunner,
165
) -> None:
166
"""Score a text query against a mix of text and image documents."""
167
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
@@ -251,28 +172,17 @@ def _run_multimodal_mixed_docs_test(
172
_make_image_mm_param(red_image),
173
]
174
254
- with vllm_runner(
255
- model,
256
- runner="pooling",
257
- dtype=dtype,
258
- max_model_len=4096,
259
- enforce_eager=True,
260
- gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
261
- ) as vllm_model:
262
- scores = vllm_model.llm.score(query, documents)
175
+ scores = vllm_model.llm.score(query, documents)
176
264
- assert len(scores) == 2
265
- for s in scores:
266
- assert isinstance(s.outputs.score, float)
267
- # Text document about France should score higher than a random image
268
- assert scores[0].outputs.score > scores[1].outputs.score
177
+ assert len(scores) == 2
178
+ for s in scores:
179
+ assert isinstance(s.outputs.score, float)
180
+ # Text document about France should score higher than a random image
181
+ assert scores[0].outputs.score > scores[1].outputs.score
182
183
184
def _run_multimodal_image_query_text_docs_test(
272
- vllm_runner: type[VllmRunner],
273
- model: str,
274
- *,
275
- dtype: str,
185
+ vllm_model: VllmRunner,
186
) -> None:
187
"""Score an image query against text documents."""
188
red_image = _make_base64_image(64, 64, color=(255, 0, 0))
@@ -283,6 +193,20 @@ def _run_multimodal_image_query_text_docs_test(
193
"The weather forecast shows rain tomorrow.",
194
]
195
196
+ scores = vllm_model.llm.score(image_query, documents)
197
+
198
+ assert len(scores) == 2
199
+ for s in scores:
200
+ assert isinstance(s.outputs.score, float)
201
+
202
+
203
+@pytest.mark.parametrize("model", MODELS)
204
+@pytest.mark.parametrize("dtype", [DTYPE])
205
+def test_colpali_default_runner(
206
+ vllm_runner,
207
+ model: str,
208
+ dtype: str,
209
+) -> None:
210
with vllm_runner(
211
model,
212
runner="pooling",
@@ -291,40 +215,31 @@ def _run_multimodal_image_query_text_docs_test(
215
enforce_eager=True,
216
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
217
) as vllm_model:
294
- scores = vllm_model.llm.score(image_query, documents)
295
-
296
- assert len(scores) == 2
297
- for s in scores:
298
- assert isinstance(s.outputs.score, float)
218
+ _run_token_embed_test(vllm_model, model)
219
+ _run_late_interaction_test(vllm_model)
220
+ _run_relevance_test(vllm_model)
221
+ _run_multimodal_mixed_docs_test(vllm_model)
222
+ _run_multimodal_image_query_text_docs_test(vllm_model)
223
224
225
@pytest.mark.parametrize("model", MODELS)
226
@pytest.mark.parametrize("dtype", [DTYPE])
303
-def test_colpali_multimodal_text_query_image_docs(
227
+def test_colpali_v2_multimodal_text_query_image_docs(
228
vllm_runner,
229
monkeypatch: pytest.MonkeyPatch,
230
model: str,
231
dtype: str,
232
) -> None:
233
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
310
- _run_multimodal_text_query_image_docs_test(vllm_runner, model, dtype=dtype)
311
-
312
-
313
-@pytest.mark.parametrize("model", MODELS)
314
-@pytest.mark.parametrize("dtype", [DTYPE])
315
-def test_colpali_multimodal_mixed_docs(
316
- vllm_runner,
317
- model: str,
318
- dtype: str,
319
-) -> None:
320
- _run_multimodal_mixed_docs_test(vllm_runner, model, dtype=dtype)
321
-
322
-
323
-@pytest.mark.parametrize("model", MODELS)
324
-@pytest.mark.parametrize("dtype", [DTYPE])
325
-def test_colpali_multimodal_image_query_text_docs(
326
- vllm_runner,
327
- model: str,
328
- dtype: str,
329
-) -> None:
330
- _run_multimodal_image_query_text_docs_test(vllm_runner, model, dtype=dtype)
234
+ with vllm_runner(
235
+ model,
236
+ runner="pooling",
237
+ dtype=dtype,
238
+ max_model_len=4096,
239
+ enforce_eager=True,
240
+ gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
241
+ attention_backend=attention_backend,
242
+ kernel_config={"enable_flashinfer_autotune": False},
243
+ ) as vllm_model:
244
+ assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
245
+ _run_multimodal_text_query_image_docs_test(vllm_model)
tests/models/multimodal/pooling/test_colqwen3_5.py
+50
-81
@@ -35,74 +35,65 @@ TEXT_DOCUMENTS = [
35
DTYPE = "half"
36
37
38
-def _run_token_embed_test(
39
- vllm_runner: type[VllmRunner],
40
- model: str,
41
- *,
42
- dtype: str,
43
-) -> None:
44
- """Verify per-token embedding shape and L2 normalization."""
38
+@pytest.fixture(scope="module", params=MODELS)
39
+def colqwen3_5_model(request, vllm_runner):
40
+ model = request.param
41
with vllm_runner(
42
model,
43
runner="pooling",
48
- dtype=dtype,
44
+ dtype=DTYPE,
45
max_model_len=4096,
46
enforce_eager=True,
47
) as vllm_model:
52
- outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
48
+ yield model, vllm_model
49
54
- assert len(outputs) == 1
55
- emb = torch.tensor(outputs[0])
56
- # Token embeddings should be 2D: [num_tokens, embed_dim]
57
- assert emb.dim() == 2
58
- assert emb.shape[1] == EMBED_DIMS[model]
59
- assert emb.shape[0] > 1
50
61
- # Verify L2 normalization
62
- norms = torch.norm(emb, p=2, dim=-1)
63
- torch.testing.assert_close(
64
- norms,
65
- torch.ones_like(norms),
66
- rtol=1e-2,
67
- atol=1e-2,
68
- )
51
+def _run_token_embed_test(
52
+ vllm_model: VllmRunner,
53
+ model: str,
54
+) -> None:
55
+ """Verify per-token embedding shape and L2 normalization."""
56
+ outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
57
+
58
+ assert len(outputs) == 1
59
+ emb = torch.tensor(outputs[0])
60
+ # Token embeddings should be 2D: [num_tokens, embed_dim]
61
+ assert emb.dim() == 2
62
+ assert emb.shape[1] == EMBED_DIMS[model]
63
+ assert emb.shape[0] > 1
64
+
65
+ # Verify L2 normalization
66
+ norms = torch.norm(emb, p=2, dim=-1)
67
+ torch.testing.assert_close(
68
+ norms,
69
+ torch.ones_like(norms),
70
+ rtol=1e-2,
71
+ atol=1e-2,
72
+ )
73
74
75
def _run_late_interaction_test(
72
- vllm_runner: type[VllmRunner],
73
- model: str,
74
- *,
75
- dtype: str,
76
+ vllm_model: VllmRunner,
77
) -> None:
78
"""Verify MaxSim scoring matches manual computation."""
79
from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
80
80
- with vllm_runner(
81
- model,
82
- runner="pooling",
83
- dtype=dtype,
84
- max_model_len=4096,
85
- enforce_eager=True,
86
- ) as vllm_model:
87
- q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
88
- d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
81
+ q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
82
+ d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
83
90
- q_emb = torch.tensor(q_outputs[0])
91
- d_emb = torch.tensor(d_outputs[0])
84
+ q_emb = torch.tensor(q_outputs[0])
85
+ d_emb = torch.tensor(d_outputs[0])
86
93
- manual_score = compute_maxsim_score(q_emb, d_emb).item()
87
+ manual_score = compute_maxsim_score(q_emb, d_emb).item()
88
95
- vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0])
89
+ vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0])
90
97
- assert len(vllm_scores) == 1
98
- assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
91
+ assert len(vllm_scores) == 1
92
+ assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
93
94
95
def _run_relevance_test(
102
- vllm_runner: type[VllmRunner],
103
- model: str,
104
- *,
105
- dtype: str,
96
+ vllm_model: VllmRunner,
97
) -> None:
98
"""Verify that relevant documents score higher than irrelevant ones."""
99
query = "What is machine learning?"
@@ -112,48 +103,26 @@ def _run_relevance_test(
103
"Deep learning uses neural networks for complex tasks.",
104
]
105
115
- with vllm_runner(
116
- model,
117
- runner="pooling",
118
- dtype=dtype,
119
- max_model_len=4096,
120
- enforce_eager=True,
121
- ) as vllm_model:
122
- scores = vllm_model.score(query, documents)
106
+ scores = vllm_model.score(query, documents)
107
124
- assert len(scores) == 3
125
- assert scores[0] > scores[1], "ML doc should score higher than weather doc"
126
- assert scores[2] > scores[1], "DL doc should score higher than weather doc"
108
+ assert len(scores) == 3
109
+ assert scores[0] > scores[1], "ML doc should score higher than weather doc"
110
+ assert scores[2] > scores[1], "DL doc should score higher than weather doc"
111
112
129
-@pytest.mark.parametrize("model", MODELS)
130
-@pytest.mark.parametrize("dtype", [DTYPE])
131
-def test_colqwen3_5_token_embed(
132
- vllm_runner,
133
- model: str,
134
- dtype: str,
135
-) -> None:
136
- _run_token_embed_test(vllm_runner, model, dtype=dtype)
113
+def test_colqwen3_5_token_embed(colqwen3_5_model) -> None:
114
+ model, vllm_model = colqwen3_5_model
115
+ _run_token_embed_test(vllm_model, model)
116
117
139
-@pytest.mark.parametrize("model", MODELS)
140
-@pytest.mark.parametrize("dtype", [DTYPE])
141
-def test_colqwen3_5_late_interaction_scoring(
142
- vllm_runner,
143
- model: str,
144
- dtype: str,
145
-) -> None:
146
- _run_late_interaction_test(vllm_runner, model, dtype=dtype)
118
+def test_colqwen3_5_late_interaction_scoring(colqwen3_5_model) -> None:
119
+ _, vllm_model = colqwen3_5_model
120
+ _run_late_interaction_test(vllm_model)
121
122
149
-@pytest.mark.parametrize("model", MODELS)
150
-@pytest.mark.parametrize("dtype", [DTYPE])
151
-def test_colqwen3_5_relevance_ordering(
152
- vllm_runner,
153
- model: str,
154
- dtype: str,
155
-) -> None:
156
- _run_relevance_test(vllm_runner, model, dtype=dtype)
123
+def test_colqwen3_5_relevance_ordering(colqwen3_5_model) -> None:
124
+ _, vllm_model = colqwen3_5_model
125
+ _run_relevance_test(vllm_model)
126
127
128
@pytest.mark.parametrize(
tests/models/multimodal/pooling/test_llama_nemotron_vl.py
+150
-171
@@ -9,12 +9,14 @@ Tests for the LlamaNemotronVL model family:
9
Both variants share a SigLIP vision encoder with a bidirectional LLaMA backbone.
10
"""
11
12
+from collections.abc import Sequence
13
from io import BytesIO
14
from pathlib import Path
15
16
import pybase64 as base64
17
import pytest
18
import torch
19
+from PIL import Image
20
from transformers import AutoModel, AutoModelForSequenceClassification, AutoProcessor
21
22
from vllm.entrypoints.chat_utils import (
@@ -54,17 +56,15 @@ MODELS = ["nvidia/llama-nemotron-embed-vl-1b-v2"]
56
def _run_test(
57
hf_runner: type[HfRunner],
58
vllm_runner: type[VllmRunner],
57
- input_texts: list[str],
58
- input_images: PromptImageInput,
59
+ input_cases: list[tuple[list[str], PromptImageInput]],
60
model: str,
61
*,
62
dtype: str,
63
) -> None:
63
- """Run embedding comparison test between HF and vLLM.
64
+ """Compare HF and vLLM embeddings for all input cases.
65
66
NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing.
67
"""
67
- # Run vLLM inference first
68
with vllm_runner(
69
model,
70
runner="pooling",
@@ -74,91 +74,70 @@ def _run_test(
74
trust_remote_code=True,
75
**ROCM_ENGINE_KWARGS,
76
) as vllm_model:
77
- vllm_outputs = vllm_model.embed(input_texts, images=input_images)
77
+ vllm_outputs_per_case = [
78
+ vllm_model.embed(input_texts, images=input_images)
79
+ for input_texts, input_images in input_cases
80
+ ]
81
79
- # Run HF inference using the model's encode_queries/encode_documents API
82
with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
81
- hf_outputs = []
82
- for text, image in zip(input_texts, input_images):
83
- with torch.inference_mode():
84
- if text.startswith(QUERY_PREFIX):
85
- # Strip prefix and use encode_queries for query texts
86
- query_text = text[len(QUERY_PREFIX) :]
87
- embedding = hf_model.model.encode_queries([query_text])
88
- elif text.startswith(PASSAGE_PREFIX):
89
- # Strip prefix and use encode_documents for passages/images
90
- passage_text = text[len(PASSAGE_PREFIX) :]
91
- if image is not None:
92
- # Image document - pass image to encode_documents
93
- embedding = hf_model.model.encode_documents(
94
- images=[image],
95
- texts=[passage_text],
96
- )
83
+ hf_outputs_per_case = []
84
+ for input_texts, input_images in input_cases:
85
+ hf_outputs = []
86
+ for text, image in zip(input_texts, input_images):
87
+ with torch.inference_mode():
88
+ if text.startswith(QUERY_PREFIX):
89
+ query_text = text[len(QUERY_PREFIX) :]
90
+ embedding = hf_model.model.encode_queries([query_text])
91
+ elif text.startswith(PASSAGE_PREFIX):
92
+ passage_text = text[len(PASSAGE_PREFIX) :]
93
+ if image is not None:
94
+ embedding = hf_model.model.encode_documents(
95
+ images=[image],
96
+ texts=[passage_text],
97
+ )
98
+ else:
99
+ embedding = hf_model.model.encode_documents(
100
+ texts=[passage_text]
101
+ )
102
else:
98
- # Text-only document
99
- embedding = hf_model.model.encode_documents(
100
- texts=[passage_text]
103
+ raise ValueError(
104
+ f"Text must start with {QUERY_PREFIX!r} "
105
+ f"or {PASSAGE_PREFIX!r}"
106
)
102
- else:
103
- raise ValueError(
104
- f"Text must start with '{QUERY_PREFIX}' or '{PASSAGE_PREFIX}'"
105
- )
106
-
107
- hf_outputs.append(embedding[0].tolist())
108
-
109
- check_embeddings_close(
110
- embeddings_0_lst=hf_outputs,
111
- embeddings_1_lst=vllm_outputs,
112
- name_0="hf",
113
- name_1="vllm",
114
- )
115
-
107
117
-@pytest.mark.parametrize("model", MODELS)
118
-@pytest.mark.parametrize("dtype", ["half"])
119
-def test_models_text(
120
- hf_runner,
121
- vllm_runner,
122
- image_assets,
123
- model: str,
124
- dtype: str,
125
-) -> None:
126
- """Test text-only embedding."""
127
- input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS]
128
- input_texts = [text for text, _ in input_texts_images]
129
- input_images = [image for _, image in input_texts_images]
108
+ hf_outputs.append(embedding[0].tolist())
109
+ hf_outputs_per_case.append(hf_outputs)
110
131
- _run_test(
132
- hf_runner,
133
- vllm_runner,
134
- input_texts,
135
- input_images, # type: ignore
136
- model,
137
- dtype=dtype,
138
- )
111
+ for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
112
+ check_embeddings_close(
113
+ embeddings_0_lst=hf_outputs,
114
+ embeddings_1_lst=vllm_outputs,
115
+ name_0="hf",
116
+ name_1="vllm",
117
+ )
118
119
120
@pytest.mark.parametrize("model", MODELS)
121
@pytest.mark.parametrize("dtype", ["half"])
143
-def test_models_image(
122
+def test_models(
123
hf_runner,
124
vllm_runner,
125
image_assets,
126
model: str,
127
dtype: str,
128
) -> None:
150
- """Test image embedding."""
151
- input_texts_images = [
152
- (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
129
+ """Test text and image embedding."""
130
+ text_images = [None] * len(HF_TEXT_PROMPTS)
131
+ images = [asset.pil_image for asset in image_assets]
132
+ input_cases = [
133
+ (HF_TEXT_PROMPTS, text_images),
134
+ (HF_IMAGE_PROMPTS, images),
135
]
154
- input_texts = [text for text, _ in input_texts_images]
155
- input_images = [image for _, image in input_texts_images]
136
137
_run_test(
138
hf_runner,
139
vllm_runner,
160
- input_texts,
161
- input_images,
140
+ input_cases, # type: ignore[arg-type]
141
model,
142
dtype=dtype,
143
)
@@ -188,8 +167,11 @@ RERANKER_TEXT_DOCS = [
167
168
RERANKER_IMAGE_QUERY = "photo of a red stop sign on a street"
169
170
+RerankerDocument = tuple[str | None, Image.Image | None]
171
+RerankerCase = tuple[str, Sequence[RerankerDocument]]
172
+
173
192
-def _pil_to_data_uri(image) -> str:
174
+def _pil_to_data_uri(image: Image.Image) -> str:
175
buf = BytesIO()
176
image.save(buf, format="PNG")
177
b64 = base64.b64encode(buf.getvalue()).decode()
@@ -200,10 +182,9 @@ def _run_hf_reranker(
182
hf_runner: type[HfRunner],
183
model: str,
184
dtype: str,
203
- query: str,
204
- docs: list,
205
-) -> list[float]:
206
- """Run HF reranker inference; docs is a list of (doc_text, doc_image|None)."""
185
+ input_cases: Sequence[RerankerCase],
186
+) -> list[list[float]]:
187
+ """Run all HF reranker cases in one model lifecycle."""
188
with hf_runner(
189
model,
190
dtype=dtype,
@@ -217,35 +198,37 @@ def _run_hf_reranker(
198
use_thumbnail=True,
199
rerank_max_length=2048,
200
)
220
- examples = [
221
- {
222
- "question": query,
223
- "doc_text": doc_text if doc_text is not None else "",
224
- "doc_image": doc_image if doc_image is not None else "",
201
+ scores_per_case = []
202
+ for query, docs in input_cases:
203
+ examples = [
204
+ {
205
+ "question": query,
206
+ "doc_text": doc_text if doc_text is not None else "",
207
+ "doc_image": doc_image if doc_image is not None else "",
208
+ }
209
+ for doc_text, doc_image in docs
210
+ ]
211
+ batch_dict = processor.process_queries_documents_crossencoder(examples)
212
+ batch_dict = {
213
+ k: v.to(hf_model.model.device) if isinstance(v, torch.Tensor) else v
214
+ for k, v in batch_dict.items()
215
}
226
- for doc_text, doc_image in docs
227
- ]
228
- batch_dict = processor.process_queries_documents_crossencoder(examples)
229
- batch_dict = {
230
- k: v.to(hf_model.model.device) if isinstance(v, torch.Tensor) else v
231
- for k, v in batch_dict.items()
232
- }
233
- with torch.inference_mode():
234
- logits = hf_model.model(**batch_dict, return_dict=True).logits
235
- # vLLM applies sigmoid activation to the raw logits before returning
236
- # scores; apply the same here so both sides are comparable.
237
- scores = torch.sigmoid(logits.squeeze(-1).float())
238
- return scores.detach().cpu().tolist()
216
+ with torch.inference_mode():
217
+ logits = hf_model.model(**batch_dict, return_dict=True).logits
218
+ # vLLM applies sigmoid activation to raw logits before returning scores.
219
+ scores = torch.sigmoid(logits.squeeze(-1).float())
220
+ scores_per_case.append(scores.detach().cpu().tolist())
221
+
222
+ return scores_per_case
223
224
225
def _run_vllm_reranker(
226
vllm_runner: type[VllmRunner],
227
model: str,
228
dtype: str,
245
- query: str,
246
- docs: list,
247
-) -> list[float]:
248
- """Run vLLM reranker inference; docs is a list of (doc_text, doc_image|None)."""
229
+ input_cases: Sequence[RerankerCase],
230
+) -> list[list[float]]:
231
+ """Run all vLLM reranker cases in one model lifecycle."""
232
with vllm_runner(
233
model,
234
runner="pooling",
@@ -255,57 +238,59 @@ def _run_vllm_reranker(
238
trust_remote_code=True,
239
**ROCM_ENGINE_KWARGS,
240
) as vllm_model:
258
- has_images = any(img is not None for _, img in docs)
259
-
260
- if not has_images:
261
- # Text-only path: use the simple string score API.
262
- queries = [query] * len(docs)
263
- doc_texts = [doc_text for doc_text, _ in docs]
264
- outputs = vllm_model.score(
265
- queries,
266
- doc_texts,
267
- chat_template=_RERANKER_SCORE_TEMPLATE,
268
- )
269
- else:
270
- # Multimodal path: build ScoreMultiModalParam for each pair.
271
- query_params = [
272
- ScoreMultiModalParam(
273
- content=[
274
- ChatCompletionContentPartTextParam(
275
- type="text",
276
- text=query,
277
- )
278
- ]
241
+ scores_per_case = []
242
+ for query, docs in input_cases:
243
+ has_images = any(img is not None for _, img in docs)
244
+
245
+ if not has_images:
246
+ queries = [query] * len(docs)
247
+ doc_texts = [doc_text for doc_text, _ in docs]
248
+ outputs = vllm_model.score(
249
+ queries,
250
+ doc_texts,
251
+ chat_template=_RERANKER_SCORE_TEMPLATE,
252
)
280
- ] * len(docs)
281
-
282
- doc_params = []
283
- for doc_text, doc_image in docs:
284
- content: list = []
285
- if doc_image is not None:
286
- content.append(
287
- ChatCompletionContentPartImageParam(
288
- type="image_url",
289
- image_url={"url": _pil_to_data_uri(doc_image)},
290
- )
253
+ else:
254
+ query_params = [
255
+ ScoreMultiModalParam(
256
+ content=[
257
+ ChatCompletionContentPartTextParam(
258
+ type="text",
259
+ text=query,
260
+ )
261
+ ]
262
)
292
- if doc_text:
293
- content.append(
294
- ChatCompletionContentPartTextParam(
295
- type="text",
296
- text=doc_text,
263
+ ] * len(docs)
264
+
265
+ doc_params = []
266
+ for doc_text, doc_image in docs:
267
+ content: list = []
268
+ if doc_image is not None:
269
+ content.append(
270
+ ChatCompletionContentPartImageParam(
271
+ type="image_url",
272
+ image_url={"url": _pil_to_data_uri(doc_image)},
273
+ )
274
)
298
- )
299
- doc_params.append(ScoreMultiModalParam(content=content))
275
+ if doc_text:
276
+ content.append(
277
+ ChatCompletionContentPartTextParam(
278
+ type="text",
279
+ text=doc_text,
280
+ )
281
+ )
282
+ doc_params.append(ScoreMultiModalParam(content=content))
283
301
- raw_outputs = vllm_model.llm.score(
302
- query_params,
303
- doc_params,
304
- chat_template=_RERANKER_SCORE_TEMPLATE,
305
- )
306
- outputs = [o.outputs.score for o in raw_outputs]
284
+ raw_outputs = vllm_model.llm.score(
285
+ query_params,
286
+ doc_params,
287
+ chat_template=_RERANKER_SCORE_TEMPLATE,
288
+ )
289
+ outputs = [output.outputs.score for output in raw_outputs]
290
+
291
+ scores_per_case.append(outputs)
292
308
- return outputs
293
+ return scores_per_case
294
295
296
def _run_reranker_test(
@@ -313,50 +298,44 @@ def _run_reranker_test(
298
vllm_runner: type[VllmRunner],
299
model: str,
300
dtype: str,
316
- query: str,
317
- docs: list,
301
+ input_cases: Sequence[RerankerCase],
302
) -> None:
319
- """Compare HF and vLLM reranker scores.
303
+ """Compare HF and vLLM reranker scores for all input cases.
304
305
NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing.
306
"""
323
- vllm_scores = _run_vllm_reranker(vllm_runner, model, dtype, query, docs)
324
- hf_scores = _run_hf_reranker(hf_runner, model, dtype, query, docs)
307
+ vllm_scores_per_case = _run_vllm_reranker(vllm_runner, model, dtype, input_cases)
308
+ hf_scores_per_case = _run_hf_reranker(hf_runner, model, dtype, input_cases)
309
326
- assert len(hf_scores) == len(vllm_scores), (
327
- f"Output length mismatch: HF={len(hf_scores)}, vLLM={len(vllm_scores)}"
328
- )
329
- # NOTE: ROCm shows slightly higher numerical variance dues to different attention
330
- # backend between vLLM and HF; use a marginally looser tolerance
310
+ # ROCm has slightly higher variance because vLLM and HF use different
311
+ # attention backends.
312
rel_tol = 0.022 if current_platform.is_rocm() else 0.02
332
- for i, (hf_score, vllm_score) in enumerate(zip(hf_scores, vllm_scores)):
333
- assert hf_score == pytest.approx(vllm_score, rel=rel_tol), (
334
- f"Score mismatch at index {i}: HF={hf_score:.4f}, vLLM={vllm_score:.4f}"
313
+ for hf_scores, vllm_scores in zip(hf_scores_per_case, vllm_scores_per_case):
314
+ assert len(hf_scores) == len(vllm_scores), (
315
+ f"Output length mismatch: HF={len(hf_scores)}, vLLM={len(vllm_scores)}"
316
)
317
+ for i, (hf_score, vllm_score) in enumerate(zip(hf_scores, vllm_scores)):
318
+ assert hf_score == pytest.approx(vllm_score, rel=rel_tol), (
319
+ f"Score mismatch at index {i}: HF={hf_score:.4f}, vLLM={vllm_score:.4f}"
320
+ )
321
322
323
@pytest.mark.parametrize("model", RERANKER_MODELS)
324
@pytest.mark.parametrize("dtype", ["half"])
340
-def test_reranker_text(
341
- hf_runner,
342
- vllm_runner,
343
- model: str,
344
- dtype: str,
345
-) -> None:
346
- """Test reranking with text-only query and text documents."""
347
- docs = [(text, None) for text in RERANKER_TEXT_DOCS]
348
- _run_reranker_test(hf_runner, vllm_runner, model, dtype, RERANKER_TEXT_QUERY, docs)
349
-
350
-
351
-@pytest.mark.parametrize("model", RERANKER_MODELS)
352
-@pytest.mark.parametrize("dtype", ["half"])
353
-def test_reranker_image_doc(
325
+def test_reranker(
326
hf_runner,
327
vllm_runner,
328
image_assets,
329
model: str,
330
dtype: str,
331
) -> None:
360
- """Test reranking with text query against image documents."""
361
- docs = [(None, asset.pil_image) for asset in image_assets]
362
- _run_reranker_test(hf_runner, vllm_runner, model, dtype, RERANKER_IMAGE_QUERY, docs)
332
+ """Test reranking with text and image documents."""
333
+ text_docs: list[RerankerDocument] = [(text, None) for text in RERANKER_TEXT_DOCS]
334
+ image_docs: list[RerankerDocument] = [
335
+ (None, asset.pil_image) for asset in image_assets
336
+ ]
337
+ input_cases: list[RerankerCase] = [
338
+ (RERANKER_TEXT_QUERY, text_docs),
339
+ (RERANKER_IMAGE_QUERY, image_docs),
340
+ ]
341
+ _run_reranker_test(hf_runner, vllm_runner, model, dtype, input_cases)
tests/models/multimodal/pooling/test_siglip.py
+65
-98
@@ -33,16 +33,11 @@ MODELS = [
33
def _run_test(
34
hf_runner: type[HfRunner],
35
vllm_runner: type[VllmRunner],
36
- input_texts: list[str],
37
- input_images: PromptImageInput,
36
+ input_cases: list[tuple[list[str], PromptImageInput, dict[str, Any]]],
37
model: str,
38
*,
39
dtype: str,
41
- tokenization_kwargs: dict[str, Any] | None = None,
40
) -> None:
43
- if tokenization_kwargs is None:
44
- tokenization_kwargs = {}
45
-
41
with vllm_runner(
42
model,
43
runner="pooling",
@@ -51,116 +46,88 @@ def _run_test(
46
max_model_len=64,
47
gpu_memory_utilization=0.7,
48
) as vllm_model:
54
- vllm_outputs = vllm_model.embed(
55
- input_texts, images=input_images, tokenization_kwargs=tokenization_kwargs
56
- )
49
+ vllm_outputs_per_case = [
50
+ vllm_model.embed(
51
+ input_texts,
52
+ images=input_images,
53
+ tokenization_kwargs=tokenization_kwargs,
54
+ )
55
+ for input_texts, input_images, tokenization_kwargs in input_cases
56
+ ]
57
+
58
+ texts = [HF_TEXT_PROMPTS[0]]
59
+ images = [input_cases[1][1][0]]
60
+ with pytest.raises(ValueError, match="not both"):
61
+ vllm_model.embed(texts, images=images)
62
+
63
+ vllm_model.embed(texts)
64
+ vllm_model.embed([""], images=images)
65
66
with hf_runner(model, dtype=dtype, auto_cls=SiglipModel) as hf_model:
59
- all_inputs = hf_model.get_inputs(
60
- input_texts, images=input_images, tokenization_kwargs=tokenization_kwargs
67
+ hf_outputs_per_case = []
68
+ for input_texts, input_images, tokenization_kwargs in input_cases:
69
+ all_inputs = hf_model.get_inputs(
70
+ input_texts,
71
+ images=input_images,
72
+ tokenization_kwargs=tokenization_kwargs,
73
+ )
74
+
75
+ hf_outputs = []
76
+ for inputs in all_inputs:
77
+ inputs = hf_model.wrap_device(inputs)
78
+
79
+ if "pixel_values" in inputs:
80
+ pooled_output = hf_model.model.get_image_features(
81
+ pixel_values=inputs.pixel_values,
82
+ )
83
+ else:
84
+ pooled_output = hf_model.model.get_text_features(
85
+ input_ids=inputs.input_ids,
86
+ )
87
+
88
+ if not isinstance(pooled_output, torch.Tensor):
89
+ pooled_output = pooled_output.pooler_output
90
+ pooled_output = pooled_output.squeeze(0)
91
+ hf_outputs.append(pooled_output.tolist())
92
+
93
+ hf_outputs_per_case.append(hf_outputs)
94
+
95
+ for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case):
96
+ check_embeddings_close(
97
+ embeddings_0_lst=hf_outputs,
98
+ embeddings_1_lst=vllm_outputs,
99
+ name_0="hf",
100
+ name_1="vllm",
101
)
102
63
- all_outputs = []
64
- for inputs in all_inputs:
65
- inputs = hf_model.wrap_device(inputs)
66
-
67
- if "pixel_values" in inputs:
68
- pooled_output = hf_model.model.get_image_features(
69
- pixel_values=inputs.pixel_values,
70
- )
71
- else:
72
- pooled_output = hf_model.model.get_text_features(
73
- input_ids=inputs.input_ids,
74
- )
75
-
76
- if not isinstance(pooled_output, torch.Tensor):
77
- pooled_output = pooled_output.pooler_output
78
- pooled_output = pooled_output.squeeze(0)
79
- all_outputs.append(pooled_output.tolist())
80
-
81
- hf_outputs = all_outputs
82
-
83
- check_embeddings_close(
84
- embeddings_0_lst=hf_outputs,
85
- embeddings_1_lst=vllm_outputs,
86
- name_0="hf",
87
- name_1="vllm",
88
- )
89
-
103
104
@pytest.mark.parametrize("model", MODELS)
105
@pytest.mark.parametrize("dtype", ["float"])
93
-def test_models_text(
106
+def test_models(
107
hf_runner,
108
vllm_runner,
109
image_assets,
110
model: str,
111
dtype: str,
112
) -> None:
100
- input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS]
101
- input_texts = [text for text, _ in input_texts_images]
102
- input_images = [image for _, image in input_texts_images]
103
-
104
- _run_test(
105
- hf_runner,
106
- vllm_runner,
107
- input_texts,
108
- input_images, # type: ignore
109
- model,
110
- dtype=dtype,
111
- tokenization_kwargs={
112
- "padding": "max_length",
113
- "max_length": 64,
114
- }, # siglip2 was trained with this padding setting.
115
- )
116
-
117
-
118
-@pytest.mark.parametrize("model", MODELS)
119
-@pytest.mark.parametrize("dtype", ["float"])
120
-def test_models_image(
121
- hf_runner,
122
- vllm_runner,
123
- image_assets,
124
- model: str,
125
- dtype: str,
126
-) -> None:
127
- input_texts_images = [
128
- (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
113
+ text_images = [None] * len(HF_TEXT_PROMPTS)
114
+ images = [asset.pil_image for asset in image_assets]
115
+ input_cases = [
116
+ (
117
+ HF_TEXT_PROMPTS,
118
+ text_images,
119
+ {
120
+ "padding": "max_length",
121
+ "max_length": 64,
122
+ },
123
+ ),
124
+ (HF_IMAGE_PROMPTS, images, {}),
125
]
130
- input_texts = [text for text, _ in input_texts_images]
131
- input_images = [image for _, image in input_texts_images]
126
127
_run_test(
128
hf_runner,
129
vllm_runner,
136
- input_texts,
137
- input_images,
130
+ input_cases, # type: ignore[arg-type]
131
model,
132
dtype=dtype,
133
)
141
-
142
-
143
-@pytest.mark.parametrize("model", MODELS)
144
-@pytest.mark.parametrize("dtype", ["float"])
145
-def test_models_text_image_no_crash(
146
- vllm_runner,
147
- image_assets,
148
- model: str,
149
- dtype: str,
150
-) -> None:
151
- texts = [HF_TEXT_PROMPTS[0]]
152
- images = [image_assets[0].pil_image]
153
-
154
- with vllm_runner(
155
- model,
156
- runner="pooling",
157
- dtype=dtype,
158
- enforce_eager=True,
159
- max_model_len=64,
160
- gpu_memory_utilization=0.7,
161
- ) as vllm_model:
162
- with pytest.raises(ValueError, match="not both"):
163
- vllm_model.embed(texts, images=images)
164
-
165
- vllm_model.embed(texts)
166
- vllm_model.embed([""], images=images)
tests/models/quantization/test_awq.py
+15
-17
@@ -17,6 +17,15 @@ HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
17
}
18
)
19
20
+IMAGE_SIZE_FACTOR_GROUPS = (
21
+ # Single-scale
22
+ (1.0,),
23
+ # Single-scale, batched
24
+ (1.0, 1.0, 1.0),
25
+ # Multi-scale
26
+ (0.25, 0.5, 1.0),
27
+)
28
+
29
30
def run_awq_test(
31
vllm_runner: type[VllmRunner],
@@ -24,7 +33,7 @@ def run_awq_test(
33
source_model: str,
34
quant_model: str,
35
*,
27
- size_factors: list[float],
36
+ size_factor_groups: tuple[tuple[float, ...], ...],
37
dtype: str,
38
max_tokens: int,
39
num_logprobs: int,
@@ -33,11 +42,12 @@ def run_awq_test(
42
):
43
images = [asset.pil_image for asset in image_assets]
44
36
- inputs_per_image = [
45
+ inputs_per_image_and_size_group = [
46
(
47
[prompt for _ in size_factors],
48
[rescale_image_size(image, factor) for factor in size_factors],
49
)
50
+ for size_factors in size_factor_groups
51
for image, prompt in zip(images, HF_IMAGE_PROMPTS)
52
]
53
@@ -60,7 +70,7 @@ def run_awq_test(
70
vllm_model.generate_greedy_logprobs(
71
prompts, max_tokens, num_logprobs=num_logprobs, images=images
72
)
63
- for prompts, images in inputs_per_image
73
+ for prompts, images in inputs_per_image_and_size_group
74
]
75
76
with vllm_runner(
@@ -77,7 +87,7 @@ def run_awq_test(
87
vllm_model.generate_greedy_logprobs(
88
prompts, max_tokens, num_logprobs=num_logprobs, images=images
89
)
80
- for prompts, images in inputs_per_image
90
+ for prompts, images in inputs_per_image_and_size_group
91
]
92
93
for source_outputs, quant_outputs in zip(
@@ -128,17 +138,6 @@ def test_awq_load(
138
("source_model", "quant_model"),
139
[("OpenGVLab/InternVL2-2B", "OpenGVLab/InternVL2-2B-AWQ")],
140
)
131
-@pytest.mark.parametrize(
132
- "size_factors",
133
- [
134
- # Single-scale
135
- [1.0],
136
- # Single-scale, batched
137
- [1.0, 1.0, 1.0],
138
- # Multi-scale
139
- [0.25, 0.5, 1.0],
140
- ],
141
-)
141
@pytest.mark.parametrize("dtype", ["half"])
142
@pytest.mark.parametrize("max_tokens", [128])
143
@pytest.mark.parametrize("num_logprobs", [5])
@@ -148,7 +147,6 @@ def test_awq_models(
147
image_assets,
148
source_model,
149
quant_model,
151
- size_factors,
150
dtype,
151
max_tokens,
152
num_logprobs,
@@ -158,7 +156,7 @@ def test_awq_models(
156
image_assets,
157
source_model,
158
quant_model,
161
- size_factors=size_factors,
159
+ size_factor_groups=IMAGE_SIZE_FACTOR_GROUPS,
160
dtype=dtype,
161
max_tokens=max_tokens,
162
num_logprobs=num_logprobs,