@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 165 lines 5.03 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Example of using ColModernVBERT late interaction model for reranking.
4
5 ColModernVBERT is a multi-modal ColBERT-style model combining a SigLIP
6 vision encoder with a ModernBERT text encoder. It produces per-token
7 embeddings and uses MaxSim scoring for retrieval and reranking.
8 Supports both text and image inputs.
9
10 Start the server with:
11 vllm serve ModernVBERT/colmodernvbert-merged --max-model-len 8192
12
13 Then run this script:
14 python colmodernvbert_rerank_online.py
15 """
16
17 import requests
18
19 MODEL = "ModernVBERT/colmodernvbert-merged"
20 BASE_URL = "http://127.0.0.1:8000"
21
22 headers = {"accept": "application/json", "Content-Type": "application/json"}
23
24 IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/thumb/4/47/PNG_transparency_demonstration_1.png/300px-PNG_transparency_demonstration_1.png" # noqa: E501
25
26
27 def rerank_text():
28 """Text-only reranking via /rerank endpoint."""
29 print("=" * 60)
30 print("1. Text reranking (/rerank)")
31 print("=" * 60)
32
33 data = {
34 "model": MODEL,
35 "query": "What is machine learning?",
36 "documents": [
37 "Machine learning is a subset of artificial intelligence.",
38 "Python is a programming language.",
39 "Deep learning uses neural networks for complex tasks.",
40 "The weather today is sunny.",
41 ],
42 }
43
44 response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
45
46 if response.status_code == 200:
47 result = response.json()
48 print("\n Ranked documents (most relevant first):")
49 for item in result["results"]:
50 doc_idx = item["index"]
51 score = item["relevance_score"]
52 print(f" [{score:.4f}] {data['documents'][doc_idx]}")
53 else:
54 print(f" Request failed: {response.status_code}")
55 print(f" {response.text[:300]}")
56
57
58 def score_text():
59 """Text-only scoring via /score endpoint."""
60 print()
61 print("=" * 60)
62 print("2. Text scoring (/score)")
63 print("=" * 60)
64
65 query = "What is the capital of France?"
66 documents = [
67 "The capital of France is Paris.",
68 "Berlin is the capital of Germany.",
69 "Python is a programming language.",
70 ]
71
72 data = {
73 "model": MODEL,
74 "text_1": query,
75 "text_2": documents,
76 }
77
78 response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
79
80 if response.status_code == 200:
81 result = response.json()
82 print(f"\n Query: {query}\n")
83 for item in result["data"]:
84 idx = item["index"]
85 score = item["score"]
86 print(f" Doc {idx} (score={score:.4f}): {documents[idx]}")
87 else:
88 print(f" Request failed: {response.status_code}")
89 print(f" {response.text[:300]}")
90
91
92 def score_text_top_n():
93 """Text reranking with top_n filtering via /rerank endpoint."""
94 print()
95 print("=" * 60)
96 print("3. Text reranking with top_n=2 (/rerank)")
97 print("=" * 60)
98
99 data = {
100 "model": MODEL,
101 "query": "What is the capital of France?",
102 "documents": [
103 "The capital of France is Paris.",
104 "Berlin is the capital of Germany.",
105 "Python is a programming language.",
106 "The Eiffel Tower is in Paris.",
107 ],
108 "top_n": 2,
109 }
110
111 response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
112
113 if response.status_code == 200:
114 result = response.json()
115 print(f"\n Top {data['top_n']} results:")
116 for item in result["results"]:
117 doc_idx = item["index"]
118 score = item["relevance_score"]
119 print(f" [{score:.4f}] {data['documents'][doc_idx]}")
120 else:
121 print(f" Request failed: {response.status_code}")
122 print(f" {response.text[:300]}")
123
124
125 def rerank_multimodal():
126 """Multimodal reranking with text and image documents via /rerank."""
127 print()
128 print("=" * 60)
129 print("4. Multimodal reranking: text query vs image document (/rerank)")
130 print("=" * 60)
131
132 data = {
133 "model": MODEL,
134 "query": "A colorful logo with transparency",
135 "documents": [
136 {"content": [{"type": "image_url", "image_url": {"url": IMAGE_URL}}]},
137 "Python is a programming language.",
138 "The weather today is sunny.",
139 ],
140 }
141
142 response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
143
144 if response.status_code == 200:
145 result = response.json()
146 print("\n Ranked documents (most relevant first):")
147 labels = ["[image]", "Python doc", "Weather doc"]
148 for item in result["results"]:
149 doc_idx = item["index"]
150 score = item["relevance_score"]
151 print(f" [{score:.4f}] {labels[doc_idx]}")
152 else:
153 print(f" Request failed: {response.status_code}")
154 print(f" {response.text[:300]}")
155
156
157 def main():
158 rerank_text()
159 score_text()
160 score_text_top_n()
161 rerank_multimodal()
162
163
164 if __name__ == "__main__":
165 main()