@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 257 lines 7.97 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # ruff: noqa: E501
4 """Example of using ColQwen3 late interaction model for reranking and scoring.
5
6 ColQwen3 is a multi-modal ColBERT-style model based on Qwen3-VL.
7 It produces per-token embeddings and uses MaxSim scoring for retrieval
8 and reranking. Supports both text and image inputs.
9
10 Start the server with:
11 vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 50000
12
13 Then run this script:
14 python colqwen3_rerank_online.py
15 """
16
17 from io import BytesIO
18
19 import pybase64 as base64
20 import requests
21 from PIL import Image
22
23 MODEL = "TomoroAI/tomoro-colqwen3-embed-4b"
24 BASE_URL = "http://127.0.0.1:8000"
25
26 headers = {"accept": "application/json", "Content-Type": "application/json"}
27
28 # ── Image helpers ──────────────────────────────────────────
29
30
31 def load_image(url: str) -> Image.Image:
32 """Download an image from URL (handles Wikimedia 403)."""
33 for hdrs in (
34 {},
35 {"User-Agent": "Mozilla/5.0 (compatible; ColQwen3-demo/1.0)"},
36 ):
37 resp = requests.get(url, headers=hdrs, timeout=15)
38 if resp.status_code == 403:
39 continue
40 resp.raise_for_status()
41 return Image.open(BytesIO(resp.content)).convert("RGB")
42 raise RuntimeError(f"Could not fetch image from {url}")
43
44
45 def encode_image_base64(image: Image.Image) -> str:
46 """Encode a PIL image to a base64 data URI."""
47 buf = BytesIO()
48 image.save(buf, format="PNG")
49 return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
50
51
52 def make_image_content(image_url: str, text: str = "Describe the image.") -> dict:
53 """Build a ScoreMultiModalParam dict from an image URL."""
54 image = load_image(image_url)
55 return {
56 "content": [
57 {
58 "type": "image_url",
59 "image_url": {"url": encode_image_base64(image)},
60 },
61 {"type": "text", "text": text},
62 ]
63 }
64
65
66 # ── Sample image URLs ─────────────────────────────────────
67
68 IMAGE_URLS = {
69 "beijing": "https://upload.wikimedia.org/wikipedia/commons/6/61/Beijing_skyline_at_night.JPG",
70 "london": "https://upload.wikimedia.org/wikipedia/commons/4/49/London_skyline.jpg",
71 "singapore": "https://upload.wikimedia.org/wikipedia/commons/2/27/Singapore_skyline_2022.jpg",
72 }
73
74 # ── Text-only examples ────────────────────────────────────
75
76
77 def rerank_text():
78 """Text-only reranking via /rerank endpoint."""
79 print("=" * 60)
80 print("1. Text reranking (/rerank)")
81 print("=" * 60)
82
83 data = {
84 "model": MODEL,
85 "query": "What is machine learning?",
86 "documents": [
87 "Machine learning is a subset of artificial intelligence.",
88 "Python is a programming language.",
89 "Deep learning uses neural networks for complex tasks.",
90 "The weather today is sunny.",
91 ],
92 }
93
94 response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
95
96 if response.status_code == 200:
97 result = response.json()
98 print("\n Ranked documents (most relevant first):")
99 for item in result["results"]:
100 doc_idx = item["index"]
101 score = item["relevance_score"]
102 print(f" [{score:.4f}] {data['documents'][doc_idx]}")
103 else:
104 print(f" Request failed: {response.status_code}")
105 print(f" {response.text[:300]}")
106
107
108 def score_text():
109 """Text-only scoring via /score endpoint."""
110 print()
111 print("=" * 60)
112 print("2. Text scoring (/score)")
113 print("=" * 60)
114
115 query = "What is the capital of France?"
116 documents = [
117 "The capital of France is Paris.",
118 "Berlin is the capital of Germany.",
119 "Python is a programming language.",
120 ]
121
122 data = {
123 "model": MODEL,
124 "text_1": query,
125 "text_2": documents,
126 }
127
128 response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
129
130 if response.status_code == 200:
131 result = response.json()
132 print(f"\n Query: {query}\n")
133 for item in result["data"]:
134 idx = item["index"]
135 score = item["score"]
136 print(f" Doc {idx} (score={score:.4f}): {documents[idx]}")
137 else:
138 print(f" Request failed: {response.status_code}")
139 print(f" {response.text[:300]}")
140
141
142 def score_text_top_n():
143 """Text reranking with top_n filtering via /rerank endpoint."""
144 print()
145 print("=" * 60)
146 print("3. Text reranking with top_n=2 (/rerank)")
147 print("=" * 60)
148
149 data = {
150 "model": MODEL,
151 "query": "What is the capital of France?",
152 "documents": [
153 "The capital of France is Paris.",
154 "Berlin is the capital of Germany.",
155 "Python is a programming language.",
156 "The Eiffel Tower is in Paris.",
157 ],
158 "top_n": 2,
159 }
160
161 response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
162
163 if response.status_code == 200:
164 result = response.json()
165 print(f"\n Top {data['top_n']} results:")
166 for item in result["results"]:
167 doc_idx = item["index"]
168 score = item["relevance_score"]
169 print(f" [{score:.4f}] {data['documents'][doc_idx]}")
170 else:
171 print(f" Request failed: {response.status_code}")
172 print(f" {response.text[:300]}")
173
174
175 # ── Multi-modal examples (text query × image documents) ──
176
177
178 def score_text_vs_images():
179 """Score a text query against image documents via /score."""
180 print()
181 print("=" * 60)
182 print("4. Multi-modal scoring: text query vs image docs (/score)")
183 print("=" * 60)
184
185 query = "Retrieve the city of Beijing"
186 labels = list(IMAGE_URLS.keys())
187 print(f"\n Loading {len(labels)} images...")
188 image_contents = [make_image_content(IMAGE_URLS[name]) for name in labels]
189
190 data = {
191 "model": MODEL,
192 "data_1": query,
193 "data_2": image_contents,
194 }
195
196 response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
197
198 if response.status_code == 200:
199 result = response.json()
200 print(f'\n Query: "{query}"\n')
201 for item in result["data"]:
202 idx = item["index"]
203 print(f" Doc {idx} [{labels[idx]}] score={item['score']:.4f}")
204 else:
205 print(f" Request failed: {response.status_code}")
206 print(f" {response.text[:300]}")
207
208
209 def rerank_text_vs_images():
210 """Rerank image documents by a text query via /rerank."""
211 print()
212 print("=" * 60)
213 print("5. Multi-modal reranking: text query vs image docs (/rerank)")
214 print("=" * 60)
215
216 query = "Retrieve the city of London"
217 labels = list(IMAGE_URLS.keys())
218 print(f"\n Loading {len(labels)} images...")
219 image_contents = [make_image_content(IMAGE_URLS[name]) for name in labels]
220
221 data = {
222 "model": MODEL,
223 "query": query,
224 "documents": image_contents,
225 "top_n": 2,
226 }
227
228 response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
229
230 if response.status_code == 200:
231 result = response.json()
232 print(f'\n Query: "{query}"')
233 print(f" Top {data['top_n']} results:\n")
234 for item in result["results"]:
235 idx = item["index"]
236 print(f" [{item['relevance_score']:.4f}] {labels[idx]}")
237 else:
238 print(f" Request failed: {response.status_code}")
239 print(f" {response.text[:300]}")
240
241
242 # ── Main ──────────────────────────────────────────────────
243
244
245 def main():
246 # Text-only
247 rerank_text()
248 score_text()
249 score_text_top_n()
250
251 # Multi-modal (text query × image documents)
252 score_text_vs_images()
253 rerank_text_vs_images()
254
255
256 if __name__ == "__main__":
257 main()