@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 197 lines 7.21 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # ruff: noqa: E501
4
5 """Example online usage of Pooling API for ColQwen3 multi-vector retrieval.
6
7 ColQwen3 is a multi-modal late interaction model based on Qwen3-VL that
8 produces per-token embeddings (320-dim, L2-normalized) for both text and
9 image inputs. Similarity is computed via MaxSim scoring.
10
11 This example mirrors the official TomoroAI inference code
12 (https://huggingface.co/TomoroAI/tomoro-colqwen3-embed-4b) but uses the
13 vLLM serving API instead of local HuggingFace model loading.
14
15 Start the server with:
16 vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 4096
17
18 Then run this script:
19 python colqwen3_token_embed_online.py
20 """
21
22 import argparse
23 from io import BytesIO
24
25 import numpy as np
26 import pybase64 as base64
27 import requests
28 from PIL import Image
29
30 # ── Helpers ─────────────────────────────────────────────────
31
32
33 def post_http_request(payload: dict, api_url: str) -> requests.Response:
34 headers = {"User-Agent": "Test Client"}
35 return requests.post(api_url, headers=headers, json=payload)
36
37
38 def load_image(url: str) -> Image.Image:
39 """Download an image from URL (handles Wikimedia 403)."""
40 for hdrs in ({}, {"User-Agent": "Mozilla/5.0 (compatible; ColQwen3-demo/1.0)"}):
41 resp = requests.get(url, headers=hdrs, timeout=10)
42 if resp.status_code == 403:
43 continue
44 resp.raise_for_status()
45 return Image.open(BytesIO(resp.content)).convert("RGB")
46 raise RuntimeError(f"Could not fetch image from {url}")
47
48
49 def encode_image_base64(image: Image.Image) -> str:
50 """Encode a PIL image to a base64 data URI."""
51 buf = BytesIO()
52 image.save(buf, format="PNG")
53 return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
54
55
56 def compute_maxsim(q_emb: np.ndarray, d_emb: np.ndarray) -> float:
57 """Compute ColBERT-style MaxSim score between query and document."""
58 sim = q_emb @ d_emb.T
59 return float(sim.max(axis=-1).sum())
60
61
62 # ── Encode functions ────────────────────────────────────────
63
64
65 def encode_queries(texts: list[str], model: str, api_url: str) -> list[np.ndarray]:
66 """Encode text queries → list of multi-vector embeddings."""
67 resp = post_http_request({"model": model, "input": texts}, api_url)
68 return [np.array(item["data"]) for item in resp.json()["data"]]
69
70
71 def encode_images(image_urls: list[str], model: str, api_url: str) -> list[np.ndarray]:
72 """Encode image documents → list of multi-vector embeddings.
73
74 Images are sent via the chat-style `messages` field so that the
75 vLLM multimodal processor handles them correctly.
76 """
77 embeddings = []
78 for url in image_urls:
79 print(f" Loading: {url.split('/')[-1]}...")
80 image = load_image(url)
81 image_uri = encode_image_base64(image)
82 resp = post_http_request(
83 {
84 "model": model,
85 "messages": [
86 {
87 "role": "user",
88 "content": [
89 {"type": "image_url", "image_url": {"url": image_uri}},
90 {"type": "text", "text": "Describe the image."},
91 ],
92 }
93 ],
94 },
95 api_url,
96 )
97 result = resp.json()
98 if resp.status_code != 200 or "data" not in result:
99 print(f" Error ({resp.status_code}): {str(result)[:200]}")
100 continue
101 embeddings.append(np.array(result["data"][0]["data"]))
102 return embeddings
103
104
105 # ── Main ────────────────────────────────────────────────────
106
107
108 def parse_args():
109 parser = argparse.ArgumentParser()
110 parser.add_argument("--host", type=str, default="localhost")
111 parser.add_argument("--port", type=int, default=8000)
112 parser.add_argument(
113 "--model",
114 type=str,
115 default="TomoroAI/tomoro-colqwen3-embed-4b",
116 )
117 return parser.parse_args()
118
119
120 def main(args):
121 pooling_url = f"http://{args.host}:{args.port}/pooling"
122 score_url = f"http://{args.host}:{args.port}/score"
123 model = args.model
124
125 # Same sample data as the official TomoroAI example
126 queries = [
127 "Retrieve the city of Singapore",
128 "Retrieve the city of Beijing",
129 "Retrieve the city of London",
130 ]
131 image_urls = [
132 "https://upload.wikimedia.org/wikipedia/commons/2/27/Singapore_skyline_2022.jpg",
133 "https://upload.wikimedia.org/wikipedia/commons/6/61/Beijing_skyline_at_night.JPG",
134 "https://upload.wikimedia.org/wikipedia/commons/4/49/London_skyline.jpg",
135 ]
136
137 # ── 1) Text query embeddings ────────────────────────────
138 print("=" * 60)
139 print("1. Encode text queries (multi-vector)")
140 print("=" * 60)
141 query_embeddings = encode_queries(queries, model, pooling_url)
142 for i, emb in enumerate(query_embeddings):
143 norm = float(np.linalg.norm(emb[0]))
144 print(f' Query {i}: {emb.shape} (L2 norm: {norm:.4f}) "{queries[i]}"')
145
146 # ── 2) Image document embeddings ────────────────────────
147 print()
148 print("=" * 60)
149 print("2. Encode image documents (multi-vector)")
150 print("=" * 60)
151 doc_embeddings = encode_images(image_urls, model, pooling_url)
152 for i, emb in enumerate(doc_embeddings):
153 print(f" Doc {i}: {emb.shape} {image_urls[i].split('/')[-1]}")
154
155 # ── 3) Cross-modal MaxSim scoring ───────────────────────
156 if doc_embeddings:
157 print()
158 print("=" * 60)
159 print("3. Cross-modal MaxSim scores (text queries × image docs)")
160 print("=" * 60)
161 # Header
162 print(f"{'':>35s}", end="")
163 for j in range(len(doc_embeddings)):
164 print(f" Doc {j:>2d}", end="")
165 print()
166 # Score matrix
167 for i, q_emb in enumerate(query_embeddings):
168 print(f" {queries[i]:<33s}", end="")
169 for j, d_emb in enumerate(doc_embeddings):
170 score = compute_maxsim(q_emb, d_emb)
171 print(f" {score:6.2f}", end="")
172 print()
173
174 # ── 4) Text-only /score endpoint ────────────────────────
175 print()
176 print("=" * 60)
177 print("4. Text-only late interaction scoring (/score endpoint)")
178 print("=" * 60)
179 text_query = "What is the capital of France?"
180 text_docs = [
181 "The capital of France is Paris.",
182 "Berlin is the capital of Germany.",
183 "Python is a programming language.",
184 ]
185 resp = post_http_request(
186 {"model": model, "text_1": text_query, "text_2": text_docs},
187 score_url,
188 )
189 print(f' Query: "{text_query}"\n')
190 for item in resp.json()["data"]:
191 idx = item["index"]
192 print(f" Doc {idx} (score={item['score']:.4f}): {text_docs[idx]}")
193
194
195 if __name__ == "__main__":
196 args = parse_args()
197 main(args)