@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 222 lines 7.16 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import json
5 from typing import TypedDict
6
7 import pytest
8 import requests
9
10 from tests.utils import RemoteOpenAIServer
11 from vllm.entrypoints.pooling.pooling.protocol import IOProcessorResponse
12
13
14 # Test configuration for ColBERT query plugin
15 class ModelConfig(TypedDict):
16 model_name: str
17 plugin: str
18 query_input: str
19 document_input: str
20 hf_overrides: str
21 embedding_dim: int
22 query_maxlen: int
23
24
25 model_config: ModelConfig = {
26 "model_name": "jinaai/jina-colbert-v2",
27 "plugin": "colbert_query_plugin",
28 "query_input": "What is machine learning?",
29 "document_input": "Machine learning is a subset of artificial intelligence.",
30 "hf_overrides": json.dumps({"architectures": ["ColBERTJinaRobertaModel"]}),
31 "embedding_dim": 128,
32 "query_maxlen": 32,
33 }
34
35
36 def _get_attr_or_val(obj: object | dict, key: str):
37 if isinstance(obj, dict) and key in obj:
38 return obj[key]
39 return getattr(obj, key, None)
40
41
42 def _check_token_embeddings(entry, expected_input_type: str):
43 assert _get_attr_or_val(entry, "object") == "embedding"
44 assert _get_attr_or_val(entry, "input_type") == expected_input_type
45
46 embedding = _get_attr_or_val(entry, "embedding")
47 assert isinstance(embedding, list) and len(embedding) > 0
48 for token_embedding in embedding:
49 assert isinstance(token_embedding, list)
50 assert len(token_embedding) == model_config["embedding_dim"]
51 return embedding
52
53
54 @pytest.fixture(scope="module")
55 def server():
56 args = [
57 "--runner",
58 "pooling",
59 "--enforce-eager",
60 "--max-num-seqs",
61 "32",
62 "--trust-remote-code",
63 "--hf_overrides",
64 model_config["hf_overrides"],
65 "--io-processor-plugin",
66 model_config["plugin"],
67 ]
68
69 with RemoteOpenAIServer(model_config["model_name"], args) as remote_server:
70 yield remote_server
71
72
73 def _post_pooling(server: RemoteOpenAIServer, data: dict):
74 request_payload = {
75 "model": model_config["model_name"],
76 "task": "plugin",
77 "data": data,
78 }
79 ret = requests.post(server.url_for("pooling"), json=request_payload)
80 ret.raise_for_status()
81 response = ret.json()
82 parsed_response: object = IOProcessorResponse(**response).data
83 assert parsed_response
84 return parsed_response
85
86
87 def test_colbert_query_plugin_query_online(server: RemoteOpenAIServer):
88 """Queries are expanded to exactly query_maxlen token vectors."""
89 parsed_response = _post_pooling(
90 server, {"input": model_config["query_input"], "input_type": "query"}
91 )
92
93 data = _get_attr_or_val(parsed_response, "data")
94 assert len(data) == 1
95
96 embedding = _check_token_embeddings(data[0], "query")
97 assert len(embedding) == model_config["query_maxlen"]
98
99 usage = _get_attr_or_val(parsed_response, "usage")
100 assert _get_attr_or_val(usage, "prompt_tokens") == model_config["query_maxlen"]
101
102
103 def test_colbert_query_plugin_document_online(server: RemoteOpenAIServer):
104 """Documents return one vector per token, with no mask expansion."""
105 parsed_response = _post_pooling(
106 server, {"input": model_config["document_input"], "input_type": "document"}
107 )
108
109 data = _get_attr_or_val(parsed_response, "data")
110 assert len(data) == 1
111
112 embedding = _check_token_embeddings(data[0], "document")
113 # No query expansion: number of vectors tracks the input length.
114 assert len(embedding) != model_config["query_maxlen"]
115
116 usage = _get_attr_or_val(parsed_response, "usage")
117 assert _get_attr_or_val(usage, "prompt_tokens") == len(embedding)
118
119
120 def test_colbert_query_plugin_missing_input_type_online(server: RemoteOpenAIServer):
121 """input_type is required; omitting it is rejected."""
122 request_payload = {
123 "model": model_config["model_name"],
124 "task": "plugin",
125 "data": {"input": model_config["document_input"]},
126 }
127 ret = requests.post(server.url_for("pooling"), json=request_payload)
128 assert ret.status_code == 400
129
130
131 def test_colbert_query_plugin_batch_online(server: RemoteOpenAIServer):
132 """A list input returns one entry per prompt."""
133 queries = ["What is machine learning?", "What is deep learning?"]
134 parsed_response = _post_pooling(server, {"input": queries, "input_type": "query"})
135
136 data = _get_attr_or_val(parsed_response, "data")
137 assert len(data) == len(queries)
138 for i, entry in enumerate(data):
139 assert _get_attr_or_val(entry, "index") == i
140 embedding = _check_token_embeddings(entry, "query")
141 assert len(embedding) == model_config["query_maxlen"]
142
143
144 @pytest.mark.parametrize("input_type", ["query", "document"])
145 def test_colbert_query_plugin_offline(vllm_runner, input_type: str):
146 """Test the ColBERT query plugin in offline mode."""
147 input_text = (
148 model_config["query_input"]
149 if input_type == "query"
150 else model_config["document_input"]
151 )
152 prompt = {
153 "data": {
154 "input": input_text,
155 "input_type": input_type,
156 }
157 }
158
159 with vllm_runner(
160 model_config["model_name"],
161 runner="pooling",
162 enforce_eager=True,
163 max_num_seqs=32,
164 trust_remote_code=True,
165 io_processor_plugin=model_config["plugin"],
166 hf_overrides=json.loads(model_config["hf_overrides"]),
167 default_torch_num_threads=1,
168 ) as llm_runner:
169 llm = llm_runner.get_llm()
170 pooler_output = llm.encode(prompt, pooling_task="plugin")
171
172 response = pooler_output[0].outputs
173 assert len(response.data) == 1
174
175 embedding = _check_token_embeddings(response.data[0], input_type)
176 if input_type == "query":
177 assert len(embedding) == model_config["query_maxlen"]
178 else:
179 assert len(embedding) != model_config["query_maxlen"]
180
181 assert response.usage.prompt_tokens == len(embedding)
182 assert response.usage.total_tokens == response.usage.prompt_tokens
183
184
185 def test_colbert_query_plugin_offline_multiple_inputs(vllm_runner):
186 """Test the ColBERT query plugin with multiple inputs in offline mode."""
187 queries = [
188 "What is machine learning?",
189 "What is deep learning?",
190 "Why?",
191 ]
192 prompts = {
193 "data": {
194 "input": queries,
195 "input_type": "query",
196 }
197 }
198
199 with vllm_runner(
200 model_config["model_name"],
201 runner="pooling",
202 enforce_eager=True,
203 max_num_seqs=32,
204 trust_remote_code=True,
205 io_processor_plugin=model_config["plugin"],
206 hf_overrides=json.loads(model_config["hf_overrides"]),
207 default_torch_num_threads=1,
208 ) as llm_runner:
209 llm = llm_runner.get_llm()
210 pooler_output = llm.encode(prompts, pooling_task="plugin")
211
212 response = pooler_output[0].outputs
213 assert len(response.data) == len(queries)
214
215 for i, entry in enumerate(response.data):
216 assert entry.index == i
217 embedding = _check_token_embeddings(entry, "query")
218 assert len(embedding) == model_config["query_maxlen"]
219
220 expected_tokens = model_config["query_maxlen"] * len(queries)
221 assert response.usage.prompt_tokens == expected_tokens
222 assert response.usage.total_tokens == response.usage.prompt_tokens