@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 183 lines 5.98 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """This example shows how to use vLLM for running offline inference with
4 vision language reranker models for multimodal scoring tasks.
5
6 Vision language rerankers score the relevance between a text query and
7 multimodal documents (text + images/videos).
8 """
9
10 from argparse import Namespace
11 from collections.abc import Callable
12 from pathlib import Path
13 from typing import NamedTuple
14
15 from vllm import LLM, EngineArgs
16 from vllm.multimodal.utils import encode_image_url, fetch_image
17 from vllm.utils.argparse_utils import FlexibleArgumentParser
18
19 TEMPLATE_HOME = Path(__file__).parent / "template"
20
21
22 query = "A woman playing with her dog on a beach at sunset."
23 document = (
24 "A woman shares a joyful moment with her golden retriever on a sun-drenched "
25 "beach at sunset, as the dog offers its paw in a heartwarming display of "
26 "companionship and trust."
27 )
28 image_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
29 video_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/demo/draw.mp4"
30 documents = [
31 {
32 "type": "text",
33 "text": document,
34 },
35 {
36 "type": "image_url",
37 "image_url": {"url": image_url},
38 },
39 {
40 "type": "image_url",
41 "image_url": {"url": encode_image_url(fetch_image(image_url))},
42 },
43 {
44 "type": "video_url",
45 "video_url": {"url": video_url},
46 },
47 ]
48
49
50 class RerankModelData(NamedTuple):
51 engine_args: EngineArgs
52 chat_template: str | None = None
53 modality: set[str] = {}
54
55
56 def run_jinavl_reranker() -> RerankModelData:
57 engine_args = EngineArgs(
58 model="jinaai/jina-reranker-m0",
59 runner="pooling",
60 max_model_len=32768,
61 trust_remote_code=True,
62 mm_processor_kwargs={
63 "min_pixels": 3136,
64 "max_pixels": 602112,
65 },
66 )
67 return RerankModelData(engine_args=engine_args, modality={"image"})
68
69
70 def run_qwen3_vl_reranker() -> RerankModelData:
71 engine_args = EngineArgs(
72 model="Qwen/Qwen3-VL-Reranker-2B",
73 runner="pooling",
74 max_model_len=16384,
75 # HuggingFace model configuration overrides required for compatibility
76 hf_overrides={
77 # Manually route to sequence classification architecture
78 # This tells vLLM to use Qwen3VLForSequenceClassification instead of
79 # the default Qwen3VLForConditionalGeneration
80 "architectures": ["Qwen3VLForSequenceClassification"],
81 # Specify which token logits to extract from the language model head
82 # The original reranker uses "no" and "yes" token logits for scoring
83 "classifier_from_token": ["no", "yes"],
84 # Enable special handling for original Qwen3-Reranker models
85 # This flag triggers conversion logic that transforms the two token
86 # vectors into a single classification vector
87 "is_original_qwen3_reranker": True,
88 },
89 )
90 chat_template_path = "qwen3_vl_reranker.jinja"
91 chat_template = (TEMPLATE_HOME / chat_template_path).read_text()
92 return RerankModelData(
93 engine_args=engine_args,
94 chat_template=chat_template,
95 modality={"image", "video"},
96 )
97
98
99 model_example_map: dict[str, Callable[[], RerankModelData]] = {
100 "jinavl_reranker": run_jinavl_reranker,
101 "qwen3_vl_reranker": run_qwen3_vl_reranker,
102 }
103
104
105 def parse_args():
106 parser = FlexibleArgumentParser(
107 description="Demo on using vLLM for offline inference with "
108 "vision language reranker models for multimodal scoring tasks."
109 )
110 parser.add_argument(
111 "--model-name",
112 "-m",
113 type=str,
114 default="jinavl_reranker",
115 choices=model_example_map.keys(),
116 help="The name of the reranker model.",
117 )
118 return parser.parse_args()
119
120
121 def main(args: Namespace):
122 # Run the selected reranker model
123 model_request = model_example_map[args.model_name]()
124 engine_args = model_request.engine_args
125
126 llm = LLM.from_engine_args(engine_args)
127
128 print("Query: string & Document: string")
129 outputs = llm.score(query, document)
130 print("Relevance scores:", [output.outputs.score for output in outputs])
131
132 print("Query: string & Document: text")
133 outputs = llm.score(
134 query, {"content": [documents[0]]}, chat_template=model_request.chat_template
135 )
136 print("Relevance scores:", [output.outputs.score for output in outputs])
137
138 print("Query: string & Document: image url")
139 outputs = llm.score(
140 query, {"content": [documents[1]]}, chat_template=model_request.chat_template
141 )
142 print("Relevance scores:", [output.outputs.score for output in outputs])
143
144 print("Query: string & Document: image base64")
145 outputs = llm.score(
146 query, {"content": [documents[2]]}, chat_template=model_request.chat_template
147 )
148 print("Relevance scores:", [output.outputs.score for output in outputs])
149
150 if "video" in model_request.modality:
151 print("Query: string & Document: video url")
152 outputs = llm.score(
153 query,
154 {"content": [documents[3]]},
155 chat_template=model_request.chat_template,
156 )
157 print("Relevance scores:", [output.outputs.score for output in outputs])
158
159 print("Query: string & Document: text + image url")
160 outputs = llm.score(
161 query,
162 {"content": [documents[0], documents[1]]},
163 chat_template=model_request.chat_template,
164 )
165 print("Relevance scores:", [output.outputs.score for output in outputs])
166
167 print("Query: string & Document: list")
168 outputs = llm.score(
169 query,
170 [
171 document,
172 {"content": [documents[0]]},
173 {"content": [documents[1]]},
174 {"content": [documents[0], documents[1]]},
175 ],
176 chat_template=model_request.chat_template,
177 )
178 print("Relevance scores:", [output.outputs.score for output in outputs])
179
180
181 if __name__ == "__main__":
182 args = parse_args()
183 main(args)