@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 145 lines 4.71 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Example of using ColQwen3.5 late interaction model for reranking.
4
5 ColQwen3.5 is a multi-modal ColBERT-style model based on Qwen3.5.
6 It produces per-token embeddings and uses MaxSim scoring for retrieval
7 and reranking. Supports both text and image inputs.
8
9 Works for any ColQwen3.5 checkpoint, e.g. `athrael-soju/colqwen3.5-4.5B-v3`
10 or `vultr/VultronRetrieverPrime-Qwen3.5-8B`.
11
12 Start the server with:
13 vllm serve athrael-soju/colqwen3.5-4.5B-v3 --max-model-len 4096 \
14 --mm-processor-kwargs '{"min_pixels": 65536, "max_pixels": 1835008}'
15
16 Then run this script:
17 python colqwen3_5_rerank_online.py
18
19 Parity note (matching the native colpali ColQwen3_5Processor pipeline):
20 - Visual-token budget: ColQwen3_5Processor uses max_num_visual_tokens=1792,
21 i.e. max_pixels = 1792 * (patch_size*merge_size)^2 = 1792 * 32^2 = 1835008
22 (with min_pixels = shortest_edge = 65536). Pass these via --mm-processor-kwargs
23 as above; the default budget gives fewer visual tokens and lower retrieval ndcg.
24 - When you build prompts yourself (token_embed), reproduce the processor exactly:
25 image (document): wrap in the instruction template
26 "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
27 "Describe the image.<|im_end|><|endoftext|>"
28 query: append the augmentation suffix <text> + "<|endoftext|>" * 10
29 Omitting these reproduces a silent ~2.5 ndcg@10 drop vs the native pipeline.
30 """
31
32 import requests
33
34 MODEL = "athrael-soju/colqwen3.5-4.5B"
35 BASE_URL = "http://127.0.0.1:8000"
36
37 headers = {"accept": "application/json", "Content-Type": "application/json"}
38
39
40 def rerank_text():
41 """Text-only reranking via /rerank endpoint."""
42 print("=" * 60)
43 print("1. Text reranking (/rerank)")
44 print("=" * 60)
45
46 data = {
47 "model": MODEL,
48 "query": "What is machine learning?",
49 "documents": [
50 "Machine learning is a subset of artificial intelligence.",
51 "Python is a programming language.",
52 "Deep learning uses neural networks for complex tasks.",
53 "The weather today is sunny.",
54 ],
55 }
56
57 response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
58
59 if response.status_code == 200:
60 result = response.json()
61 print("\n Ranked documents (most relevant first):")
62 for item in result["results"]:
63 doc_idx = item["index"]
64 score = item["relevance_score"]
65 print(f" [{score:.4f}] {data['documents'][doc_idx]}")
66 else:
67 print(f" Request failed: {response.status_code}")
68 print(f" {response.text[:300]}")
69
70
71 def score_text():
72 """Text-only scoring via /score endpoint."""
73 print()
74 print("=" * 60)
75 print("2. Text scoring (/score)")
76 print("=" * 60)
77
78 query = "What is the capital of France?"
79 documents = [
80 "The capital of France is Paris.",
81 "Berlin is the capital of Germany.",
82 "Python is a programming language.",
83 ]
84
85 data = {
86 "model": MODEL,
87 "text_1": query,
88 "text_2": documents,
89 }
90
91 response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
92
93 if response.status_code == 200:
94 result = response.json()
95 print(f"\n Query: {query}\n")
96 for item in result["data"]:
97 idx = item["index"]
98 score = item["score"]
99 print(f" Doc {idx} (score={score:.4f}): {documents[idx]}")
100 else:
101 print(f" Request failed: {response.status_code}")
102 print(f" {response.text[:300]}")
103
104
105 def score_text_top_n():
106 """Text reranking with top_n filtering via /rerank endpoint."""
107 print()
108 print("=" * 60)
109 print("3. Text reranking with top_n=2 (/rerank)")
110 print("=" * 60)
111
112 data = {
113 "model": MODEL,
114 "query": "What is the capital of France?",
115 "documents": [
116 "The capital of France is Paris.",
117 "Berlin is the capital of Germany.",
118 "Python is a programming language.",
119 "The Eiffel Tower is in Paris.",
120 ],
121 "top_n": 2,
122 }
123
124 response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
125
126 if response.status_code == 200:
127 result = response.json()
128 print(f"\n Top {data['top_n']} results:")
129 for item in result["results"]:
130 doc_idx = item["index"]
131 score = item["relevance_score"]
132 print(f" [{score:.4f}] {data['documents'][doc_idx]}")
133 else:
134 print(f" Request failed: {response.status_code}")
135 print(f" {response.text[:300]}")
136
137
138 def main():
139 rerank_text()
140 score_text()
141 score_text_top_n()
142
143
144 if __name__ == "__main__":
145 main()