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