A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 21:30 UTC
[Bugfix][Frontend] Fix Jina reranker multimodal input compatibility (#31445)
Signed-off-by: tianwenjing <tianwenjing@jfgenius.com> Signed-off-by: twj <151701930+twjww@users.noreply.github.com> Co-authored-by: tianwenjing <tianwenjing@jfgenius.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
twj committed
Dec 29, 2025 at 17:13 UTC
bf73a3e4d7e16d03c89ac068de9cf80664f8195f
2 files changed
+316
-138
tests/models/multimodal/pooling/test_jinavl_reranker.py
+313
-137
@@ -1,194 +1,370 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import cast
4
5
import pytest
6
from transformers import AutoModel
7
7
-from vllm.entrypoints.chat_utils import ChatCompletionContentPartImageParam
8
+from vllm.entrypoints.chat_utils import (
9
+ ChatCompletionContentPartImageEmbedsParam,
10
+ ChatCompletionContentPartImageParam,
11
+ ChatCompletionContentPartTextParam,
12
+)
13
from vllm.entrypoints.score_utils import ScoreMultiModalParam
14
15
from ....conftest import HfRunner, VllmRunner
16
12
-model_name = "jinaai/jina-reranker-m0"
17
+MODELS = ["jinaai/jina-reranker-m0"]
18
14
-mm_processor_kwargs = {
19
+MM_PROCESSOR_KWARGS = {
20
"min_pixels": 3136,
21
"max_pixels": 602112,
22
}
23
19
-limit_mm_per_prompt = {"image": 2}
24
+LIMIT_MM_PER_PROMPT = {"image": 2}
25
26
+CHECKPOINT_TO_HF_MAPPER = {
27
+ "visual.": "model.visual.",
28
+ "model.": "model.language_model.",
29
+}
30
+
31
+# Shared long text for test data
32
+LONG_TEXT_DOC = """We present ReaderLM-v2, a compact 1.5 billion parameter language model designed for efficient
33
+web content extraction. Our model processes documents up to 512K tokens, transforming messy HTML
34
+into clean Markdown or JSON formats with high accuracy -- making it an ideal tool for grounding
35
+large language models. The models effectiveness results from two key innovations: (1) a three-stage
36
+data synthesis pipeline that generates high quality, diverse training data by iteratively drafting,
37
+refining, and critiquing web content extraction; and (2) a unified training framework combining
38
+continuous pre-training with multi-objective optimization. Intensive evaluation demonstrates that
39
+ReaderLM-v2 outperforms GPT-4o-2024-08-06 and other larger models by 15-20% on carefully curated
40
+benchmarks, particularly excelling at documents exceeding 100K tokens, while maintaining significantly
41
+lower computational requirements.""" # noqa: E501
42
+
43
+# Test data for different scenarios
44
+TEXT_IMAGE_TEST_DATA = {
45
+ "query": [{"text": "slm markdown"}],
46
+ "documents": [
47
+ {
48
+ "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
49
+ },
50
+ {
51
+ "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
52
+ },
53
+ ],
54
+}
55
+
56
+TEXT_TEXT_TEST_DATA = {
57
+ "query": [{"text": "slm markdown"}],
58
+ "documents": [
59
+ {"text": LONG_TEXT_DOC},
60
+ {"text": "数据提取么？为什么不用正则啊,你用正则不就全解决了么?"},
61
+ ],
62
+}
63
+
64
+IMAGE_TEXT_TEST_DATA = {
65
+ "query": [
66
+ {
67
+ "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
68
+ }
69
+ ],
70
+ "documents": [
71
+ {"text": LONG_TEXT_DOC},
72
+ {"text": "数据提取么?为什么不用正则啊,你用正则不就全解决了么?"},
73
+ ],
74
+}
75
+
76
+IMAGE_IMAGE_TEST_DATA = {
77
+ "query": [
78
+ {
79
+ "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
80
+ }
81
+ ],
82
+ "documents": [
83
+ {
84
+ "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
85
+ },
86
+ {
87
+ "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
88
+ },
89
+ ],
90
+}
91
22
-def vllm_reranker(
92
+TEXT_MIXED_DOCS_TEST_DATA = {
93
+ "query": [{"text": "slm markdown"}],
94
+ "documents": [
95
+ {"text": LONG_TEXT_DOC},
96
+ {
97
+ "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
98
+ },
99
+ {"text": "数据提取么？为什么不用正则啊,你用正则不就全解决了么?"},
100
+ {
101
+ "image": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
102
+ },
103
+ ],
104
+}
105
+
106
+
107
+def _normalize_image(image_val: str) -> str:
108
+ """Normalize image value to proper format for HF model."""
109
+ return (
110
+ image_val
111
+ if image_val.startswith(("http://", "https://"))
112
+ else f"data:image/png;base64,{image_val}"
113
+ )
114
+
115
+
116
+def create_score_multimodal_param(
117
+ content_parts: list[dict],
118
+) -> ScoreMultiModalParam:
119
+ """
120
+ Create a ScoreMultiModalParam from a list of content dictionaries.
121
+
122
+ Each dict supports the following formats:
123
+ - Text: {'text': 'content'}
124
+ - Image URL: {'image': 'https://...'}
125
+ - Image Base64: {'image': 'base64_str'}
126
+ """
127
+ formatted_content = []
128
+
129
+ for part in content_parts:
130
+ if "text" in part:
131
+ formatted_content.append(
132
+ ChatCompletionContentPartTextParam(
133
+ type="text",
134
+ text=part["text"],
135
+ )
136
+ )
137
+ elif "image" in part:
138
+ image_val = part["image"]
139
+ if image_val.startswith(("http://", "https://")):
140
+ formatted_content.append(
141
+ ChatCompletionContentPartImageParam(
142
+ type="image_url",
143
+ image_url={"url": image_val},
144
+ )
145
+ )
146
+ else:
147
+ formatted_content.append(
148
+ ChatCompletionContentPartImageEmbedsParam(
149
+ type="image_embeds", image_embeds=image_val
150
+ )
151
+ )
152
+
153
+ return ScoreMultiModalParam(content=formatted_content)
154
+
155
+
156
+def _run_vllm(
157
vllm_runner: type[VllmRunner],
24
- model_name: str,
158
+ model: str,
159
dtype: str,
26
- query_strs: list[str],
27
- document_strs: list[str],
28
- query_type: str = "text",
29
- doc_type: str = "text",
30
-):
31
- def create_image_param(url: str) -> ChatCompletionContentPartImageParam:
32
- return {"type": "image_url", "image_url": {"url": f"{url}"}}
33
-
34
- query: list[str] | ScoreMultiModalParam
35
- if query_type == "text":
36
- query = query_strs
37
- elif query_type == "image":
38
- query = ScoreMultiModalParam(
39
- content=[create_image_param(url) for url in query_strs]
40
- )
41
-
42
- documents: list[str] | ScoreMultiModalParam
43
- if doc_type == "text":
44
- documents = document_strs
45
- elif doc_type == "image":
46
- documents = ScoreMultiModalParam(
47
- content=[create_image_param(url) for url in document_strs]
48
- )
160
+ query_strs: list[dict[str, str]],
161
+ document_strs: list[dict[str, str]],
162
+) -> list[float]:
163
+ """Run vLLM reranker and return scores."""
164
+ query = create_score_multimodal_param(query_strs)
165
+ documents = create_score_multimodal_param(document_strs)
166
167
with vllm_runner(
51
- model_name,
168
+ model,
169
runner="pooling",
170
dtype=dtype,
171
max_num_seqs=2,
172
max_model_len=2048,
56
- mm_processor_kwargs=mm_processor_kwargs,
57
- limit_mm_per_prompt=limit_mm_per_prompt,
173
+ mm_processor_kwargs=MM_PROCESSOR_KWARGS,
174
+ limit_mm_per_prompt=LIMIT_MM_PER_PROMPT,
175
) as vllm_model:
176
outputs = vllm_model.llm.score(query, documents)
177
178
return [output.outputs.score for output in outputs]
179
180
64
-def hf_reranker(
181
+def _run_hf(
182
hf_runner: type[HfRunner],
66
- model_name: str,
183
+ model: str,
184
dtype: str,
68
- query_strs: list[str],
69
- document_strs: list[str],
70
- query_type: str = "text",
71
- doc_type: str = "text",
72
-):
73
- checkpoint_to_hf_mapper = {
74
- "visual.": "model.visual.",
75
- "model.": "model.language_model.",
76
- }
77
-
78
- data_pairs = [[query_strs[0], d] for d in document_strs]
185
+ query_strs: list[dict[str, str]],
186
+ document_strs: list[dict[str, str]],
187
+) -> list[float]:
188
+ """Run HuggingFace reranker and return scores."""
189
+ query = query_strs[0]
190
+ if "text" in query:
191
+ query_type = "text"
192
+ query_data = query["text"]
193
+ elif "image" in query:
194
+ query_type = "image"
195
+ query_data = _normalize_image(query["image"])
196
+ else:
197
+ raise ValueError("Unsupported query format")
198
+
199
+ # Separate documents by type
200
+ text_docs: list[str] = []
201
+ image_docs: list[str] = []
202
+ text_indices: list[int] = []
203
+ image_indices: list[int] = []
204
+
205
+ for idx, doc in enumerate(document_strs):
206
+ if "text" in doc:
207
+ text_docs.append(doc["text"])
208
+ text_indices.append(idx)
209
+ elif "image" in doc:
210
+ image_docs.append(_normalize_image(doc["image"]))
211
+ image_indices.append(idx)
212
+ else:
213
+ raise ValueError(f"Unsupported document format at index {idx}")
214
+
215
+ scores: list[None | float] = [None] * len(document_strs)
216
217
with hf_runner(
81
- model_name,
218
+ model,
219
dtype=dtype,
220
trust_remote_code=True,
221
auto_cls=AutoModel,
85
- model_kwargs={"key_mapping": checkpoint_to_hf_mapper},
222
+ model_kwargs={"key_mapping": CHECKPOINT_TO_HF_MAPPER},
223
) as hf_model:
87
- return hf_model.model.compute_score(
88
- data_pairs, max_length=2048, query_type=query_type, doc_type=doc_type
89
- )
224
+ # Score text documents
225
+ if text_docs:
226
+ text_scores = hf_model.model.compute_score(
227
+ [[query_data, d] for d in text_docs],
228
+ max_length=2048,
229
+ query_type=query_type,
230
+ doc_type="text",
231
+ )
232
+ for i, s in zip(text_indices, text_scores):
233
+ scores[i] = s
234
235
+ # Score image documents
236
+ if image_docs:
237
+ image_scores = hf_model.model.compute_score(
238
+ [[query_data, d] for d in image_docs],
239
+ max_length=2048,
240
+ query_type=query_type,
241
+ doc_type="image",
242
+ )
243
+ for i, s in zip(image_indices, image_scores):
244
+ scores[i] = s
245
92
-# Visual Documents Reranking
93
-@pytest.mark.parametrize("model_name", [model_name])
94
-@pytest.mark.parametrize("dtype", ["half"])
95
-def test_model_text_image(hf_runner, vllm_runner, model_name, dtype):
96
- query = ["slm markdown"]
97
- documents = [
98
- "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png",
99
- "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png",
100
- ]
101
-
102
- hf_outputs = hf_reranker(
103
- hf_runner, model_name, dtype, query, documents, "text", "image"
104
- )
105
- vllm_outputs = vllm_reranker(
106
- vllm_runner, model_name, dtype, query, documents, "text", "image"
107
- )
246
+ assert all(s is not None for s in scores)
247
+ return cast(list[float], scores)
248
109
- assert hf_outputs[0] == pytest.approx(vllm_outputs[0], rel=0.02)
110
- assert hf_outputs[1] == pytest.approx(vllm_outputs[1], rel=0.02)
249
250
+def _run_test(
251
+ hf_runner: type[HfRunner],
252
+ vllm_runner: type[VllmRunner],
253
+ model: str,
254
+ dtype: str,
255
+ query_strs: list[dict[str, str]],
256
+ document_strs: list[dict[str, str]],
257
+) -> None:
258
+ """Run comparison test between vLLM and HuggingFace implementations."""
259
+ # NOTE: take care of the order. run vLLM first, and then run HF.
260
+ # vLLM needs a fresh new process without cuda initialization.
261
+ # if we run HF first, the cuda initialization will be done and it
262
+ # will hurt multiprocessing backend with fork method (the default method).
263
113
-# Textual Documents Reranking
114
-@pytest.mark.parametrize("model_name", [model_name])
115
-@pytest.mark.parametrize("dtype", ["half"])
116
-def test_model_text_text(hf_runner, vllm_runner, model_name, dtype):
117
- query = ["slm markdown"]
118
- documents = [
119
- """We present ReaderLM-v2, a compact 1.5 billion parameter language model designed for efficient
120
- web content extraction. Our model processes documents up to 512K tokens, transforming messy HTML
121
- into clean Markdown or JSON formats with high accuracy -- making it an ideal tool for grounding
122
- large language models. The models effectiveness results from two key innovations: (1) a three-stage
123
- data synthesis pipeline that generates high quality, diverse training data by iteratively drafting,
124
- refining, and critiquing web content extraction; and (2) a unified training framework combining
125
- continuous pre-training with multi-objective optimization. Intensive evaluation demonstrates that
126
- ReaderLM-v2 outperforms GPT-4o-2024-08-06 and other larger models by 15-20% on carefully curated
127
- benchmarks, particularly excelling at documents exceeding 100K tokens, while maintaining significantly
128
- lower computational requirements.""", # noqa: E501
129
- "数据提取么？为什么不用正则啊，你用正则不就全解决了么？",
130
- ]
131
- hf_outputs = hf_reranker(
132
- hf_runner, model_name, dtype, query, documents, "text", "text"
133
- )
134
- vllm_outputs = vllm_reranker(
135
- vllm_runner, model_name, dtype, query, documents, "text", "text"
264
+ vllm_outputs = _run_vllm(vllm_runner, model, dtype, query_strs, document_strs)
265
+ hf_outputs = _run_hf(hf_runner, model, dtype, query_strs, document_strs)
266
+
267
+ # Compare outputs
268
+ assert len(hf_outputs) == len(vllm_outputs), (
269
+ f"Output length mismatch: HF={len(hf_outputs)}, vLLM={len(vllm_outputs)}"
270
)
271
138
- assert hf_outputs[0] == pytest.approx(vllm_outputs[0], rel=0.02)
139
- assert hf_outputs[1] == pytest.approx(vllm_outputs[1], rel=0.02)
272
+ for i, (hf_score, vllm_score) in enumerate(zip(hf_outputs, vllm_outputs)):
273
+ assert hf_score == pytest.approx(vllm_score, rel=0.02), (
274
+ f"Score mismatch at index {i}: HF={hf_score}, vLLM={vllm_score}"
275
+ )
276
277
142
-# Image Querying for Textual Documents
143
-@pytest.mark.parametrize("model_name", [model_name])
278
+@pytest.mark.parametrize("model", MODELS)
279
@pytest.mark.parametrize("dtype", ["half"])
145
-def test_model_image_text(hf_runner, vllm_runner, model_name, dtype):
146
- query = [
147
- "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
148
- ]
149
- documents = [
150
- """We present ReaderLM-v2, a compact 1.5 billion parameter language model designed for efficient
151
- web content extraction. Our model processes documents up to 512K tokens, transforming messy HTML
152
- into clean Markdown or JSON formats with high accuracy -- making it an ideal tool for grounding
153
- large language models. The models effectiveness results from two key innovations: (1) a three-stage
154
- data synthesis pipeline that generates high quality, diverse training data by iteratively drafting,
155
- refining, and critiquing web content extraction; and (2) a unified training framework combining
156
- continuous pre-training with multi-objective optimization. Intensive evaluation demonstrates that
157
- ReaderLM-v2 outperforms GPT-4o-2024-08-06 and other larger models by 15-20% on carefully curated
158
- benchmarks, particularly excelling at documents exceeding 100K tokens, while maintaining significantly
159
- lower computational requirements.""", # noqa: E501
160
- "数据提取么？为什么不用正则啊，你用正则不就全解决了么？",
161
- ]
162
-
163
- hf_outputs = hf_reranker(
164
- hf_runner, model_name, dtype, query, documents, "image", "text"
165
- )
166
- vllm_outputs = vllm_reranker(
167
- vllm_runner, model_name, dtype, query, documents, "image", "text"
280
+def test_model_text_image(
281
+ hf_runner,
282
+ vllm_runner,
283
+ model: str,
284
+ dtype: str,
285
+) -> None:
286
+ """Visual Documents Reranking"""
287
+ _run_test(
288
+ hf_runner,
289
+ vllm_runner,
290
+ model,
291
+ dtype,
292
+ TEXT_IMAGE_TEST_DATA["query"],
293
+ TEXT_IMAGE_TEST_DATA["documents"],
294
)
295
170
- assert hf_outputs[0] == pytest.approx(vllm_outputs[0], rel=0.02)
171
- assert hf_outputs[1] == pytest.approx(vllm_outputs[1], rel=0.02)
296
+
297
+@pytest.mark.parametrize("model", MODELS)
298
+@pytest.mark.parametrize("dtype", ["half"])
299
+def test_model_text_text(
300
+ hf_runner,
301
+ vllm_runner,
302
+ model: str,
303
+ dtype: str,
304
+) -> None:
305
+ """Textual Documents Reranking"""
306
+ _run_test(
307
+ hf_runner,
308
+ vllm_runner,
309
+ model,
310
+ dtype,
311
+ TEXT_TEXT_TEST_DATA["query"],
312
+ TEXT_TEXT_TEST_DATA["documents"],
313
+ )
314
315
174
-# Image Querying for Image Documents
175
-@pytest.mark.parametrize("model_name", [model_name])
316
+@pytest.mark.parametrize("model", MODELS)
317
@pytest.mark.parametrize("dtype", ["half"])
177
-def test_model_image_image(hf_runner, vllm_runner, model_name, dtype):
178
- query = [
179
- "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
180
- ]
181
- documents = [
182
- "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png",
183
- "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png",
184
- ]
185
-
186
- hf_outputs = hf_reranker(
187
- hf_runner, model_name, dtype, query, documents, "image", "image"
318
+def test_model_image_text(
319
+ hf_runner,
320
+ vllm_runner,
321
+ model: str,
322
+ dtype: str,
323
+) -> None:
324
+ """Image Querying for Textual Documents"""
325
+ _run_test(
326
+ hf_runner,
327
+ vllm_runner,
328
+ model,
329
+ dtype,
330
+ IMAGE_TEXT_TEST_DATA["query"],
331
+ IMAGE_TEXT_TEST_DATA["documents"],
332
)
189
- vllm_outputs = vllm_reranker(
190
- vllm_runner, model_name, dtype, query, documents, "image", "image"
333
+
334
+
335
+@pytest.mark.parametrize("model", MODELS)
336
+@pytest.mark.parametrize("dtype", ["half"])
337
+def test_model_image_image(
338
+ hf_runner,
339
+ vllm_runner,
340
+ model: str,
341
+ dtype: str,
342
+) -> None:
343
+ """Image Querying for Image Documents"""
344
+ _run_test(
345
+ hf_runner,
346
+ vllm_runner,
347
+ model,
348
+ dtype,
349
+ IMAGE_IMAGE_TEST_DATA["query"],
350
+ IMAGE_IMAGE_TEST_DATA["documents"],
351
)
352
193
- assert hf_outputs[0] == pytest.approx(vllm_outputs[0], rel=0.02)
194
- assert hf_outputs[1] == pytest.approx(vllm_outputs[1], rel=0.02)
353
+
354
+@pytest.mark.parametrize("model", MODELS)
355
+@pytest.mark.parametrize("dtype", ["half"])
356
+def test_model_text_mixed_documents(
357
+ hf_runner,
358
+ vllm_runner,
359
+ model: str,
360
+ dtype: str,
361
+) -> None:
362
+ """Text Query for Mixed Text and Image Documents"""
363
+ _run_test(
364
+ hf_runner,
365
+ vllm_runner,
366
+ model,
367
+ dtype,
368
+ TEXT_MIXED_DOCS_TEST_DATA["query"],
369
+ TEXT_MIXED_DOCS_TEST_DATA["documents"],
370
+ )
vllm/entrypoints/score_utils.py
+3
-1
@@ -24,7 +24,9 @@ from vllm.outputs import PoolingRequestOutput
24
from vllm.tokenizers import TokenizerLike
25
26
ScoreContentPartParam: TypeAlias = (
27
- ChatCompletionContentPartImageParam | ChatCompletionContentPartImageEmbedsParam
27
+ ChatCompletionContentPartImageParam
28
+ | ChatCompletionContentPartImageEmbedsParam
29
+ | ChatCompletionContentPartTextParam
30
)
31
32