A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 22:45 UTC
[Frontend][4/n] Make pooling entrypoints request schema consensus | ScoreRequest (#33060)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
wang.yuqi committed
Feb 4, 2026 at 09:48 UTC
1b8fe6f7c4b96ec57172f4fd268341a03c12499b
8 files changed
+432
-205
examples/pooling/score/vision_rerank_api_online.py
+23
@@ -89,6 +89,29 @@ def main(args):
89
response = requests.post(rerank_url, json=prompt)
90
pprint.pprint(response.json())
91
92
+ print("Query: string & Document: text + image url")
93
+ prompt = {
94
+ "model": model,
95
+ "query": query,
96
+ "documents": {"content": [documents[0], documents[1]]},
97
+ }
98
+ response = requests.post(rerank_url, json=prompt)
99
+ pprint.pprint(response.json())
100
+
101
+ print("Query: string & Document: list")
102
+ prompt = {
103
+ "model": model,
104
+ "query": query,
105
+ "documents": [
106
+ document,
107
+ {"content": [documents[0]]},
108
+ {"content": [documents[1]]},
109
+ {"content": [documents[0], documents[1]]},
110
+ ],
111
+ }
112
+ response = requests.post(rerank_url, json=prompt)
113
+ pprint.pprint(response.json())
114
+
115
116
if __name__ == "__main__":
117
args = parse_args()
examples/pooling/score/vision_score_api_online.py
+38
@@ -92,6 +92,44 @@ def main(args):
92
response = requests.post(score_url, json=prompt)
93
pprint.pprint(response.json())
94
95
+ print("Query: string & Document: text + image url")
96
+ prompt = {
97
+ "model": model,
98
+ "queries": query,
99
+ "documents": {"content": [documents[0], documents[1]]},
100
+ }
101
+ response = requests.post(score_url, json=prompt)
102
+ pprint.pprint(response.json())
103
+
104
+ print("Query: string & Document: list")
105
+ prompt = {
106
+ "model": model,
107
+ "queries": query,
108
+ "documents": [
109
+ document,
110
+ {"content": [documents[0]]},
111
+ {"content": [documents[1]]},
112
+ {"content": [documents[0], documents[1]]},
113
+ ],
114
+ }
115
+ response = requests.post(score_url, json=prompt)
116
+ pprint.pprint(response.json())
117
+
118
+ print("Query: list & Document: list")
119
+ data = [
120
+ document,
121
+ {"content": [documents[0]]},
122
+ {"content": [documents[1]]},
123
+ {"content": [documents[0], documents[1]]},
124
+ ]
125
+ prompt = {
126
+ "model": model,
127
+ "queries": data,
128
+ "documents": data,
129
+ }
130
+ response = requests.post(score_url, json=prompt)
131
+ pprint.pprint(response.json())
132
+
133
134
if __name__ == "__main__":
135
args = parse_args()
tests/entrypoints/pooling/score/test_online_score.py
+29
@@ -90,6 +90,35 @@ class TestModel:
90
for i in range(len(vllm_outputs)):
91
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
92
93
+ def test_queries_str_items_str(
94
+ self, server: RemoteOpenAIServer, model: dict[str, Any], runner
95
+ ):
96
+ queries = "What is the capital of France?"
97
+ items = "The capital of France is Paris."
98
+
99
+ score_response = requests.post(
100
+ server.url_for("score"),
101
+ json={
102
+ "model": model["name"],
103
+ "queries": queries,
104
+ "items": items,
105
+ },
106
+ )
107
+ score_response.raise_for_status()
108
+ score = ScoreResponse.model_validate(score_response.json())
109
+
110
+ assert score.id is not None
111
+ assert score.data is not None
112
+ assert len(score.data) == 1
113
+
114
+ vllm_outputs = [d.score for d in score.data]
115
+
116
+ text_pairs = [[queries, items]]
117
+ hf_outputs = run_transformers(runner, model, text_pairs)
118
+
119
+ for i in range(len(vllm_outputs)):
120
+ assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
121
+
122
def test_text_1_str_text_2_str(
123
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
124
):
tests/entrypoints/pooling/score/test_online_score_vision.py
+124
-7
@@ -5,7 +5,7 @@ import pytest
5
import requests
6
7
from tests.utils import VLLM_PATH, RemoteOpenAIServer
8
-from vllm.entrypoints.pooling.score.protocol import ScoreResponse
8
+from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse
9
from vllm.multimodal.utils import encode_image_url, fetch_image
10
11
MODEL_NAME = "Qwen/Qwen3-VL-Reranker-2B"
@@ -16,11 +16,12 @@ HF_OVERRIDES = {
16
}
17
18
query = "A cat standing in the snow."
19
+document = "This product was excellent and exceeded my expectations."
20
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
21
documents = [
22
{
23
"type": "text",
23
- "text": query,
24
+ "text": document,
25
},
26
{
27
"type": "image_url",
@@ -32,6 +33,11 @@ documents = [
33
},
34
]
35
36
+TEXT_VS_TEXT = 0.10040374100208282
37
+TEXT_VS_IMAGE = 0.7423753142356873
38
+TEXT_VS_TEXT_PLUS_IMAGE = 0.5298863053321838
39
+TOL = 0.05
40
+
41
42
@pytest.fixture(scope="module")
43
def server():
@@ -50,15 +56,12 @@ def server():
56
57
58
def test_score_api_queries_str_documents_str(server: RemoteOpenAIServer):
53
- queries = "What is the capital of France?"
54
- documents = "The capital of France is Paris."
55
-
59
score_response = requests.post(
60
server.url_for("score"),
61
json={
62
"model": MODEL_NAME,
60
- "queries": queries,
61
- "documents": documents,
63
+ "queries": query,
64
+ "documents": document,
65
},
66
)
67
score_response.raise_for_status()
@@ -67,6 +70,8 @@ def test_score_api_queries_str_documents_str(server: RemoteOpenAIServer):
70
assert score.id is not None
71
assert score.data is not None
72
assert len(score.data) == 1
73
+ assert score.usage.prompt_tokens == 81
74
+ assert score.data[0].score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
75
76
77
def test_score_api_queries_str_documents_text_content(server: RemoteOpenAIServer):
@@ -84,6 +89,8 @@ def test_score_api_queries_str_documents_text_content(server: RemoteOpenAIServer
89
assert score.id is not None
90
assert score.data is not None
91
assert len(score.data) == 1
92
+ assert score.usage.prompt_tokens == 81
93
+ assert score.data[0].score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
94
95
96
def test_score_api_queries_str_documents_image_url_content(server: RemoteOpenAIServer):
@@ -101,6 +108,8 @@ def test_score_api_queries_str_documents_image_url_content(server: RemoteOpenAIS
108
assert score.id is not None
109
assert score.data is not None
110
assert len(score.data) == 1
111
+ assert score.usage.prompt_tokens == 98
112
+ assert score.data[0].score == pytest.approx(TEXT_VS_IMAGE, rel=TOL)
113
114
115
def test_score_api_queries_str_documents_image_base64_content(
@@ -120,3 +129,111 @@ def test_score_api_queries_str_documents_image_base64_content(
129
assert score.id is not None
130
assert score.data is not None
131
assert len(score.data) == 1
132
+ assert score.usage.prompt_tokens == 98
133
+ assert score.data[0].score == pytest.approx(TEXT_VS_IMAGE, rel=TOL)
134
+
135
+
136
+def test_score_api_queries_str_documents_image_url_plus_text_content(
137
+ server: RemoteOpenAIServer,
138
+):
139
+ score_response = requests.post(
140
+ server.url_for("score"),
141
+ json={
142
+ "model": MODEL_NAME,
143
+ "queries": query,
144
+ "documents": {"content": [documents[0], documents[1]]},
145
+ },
146
+ )
147
+ score_response.raise_for_status()
148
+ score = ScoreResponse.model_validate(score_response.json())
149
+
150
+ assert score.id is not None
151
+ assert score.data is not None
152
+ assert len(score.data) == 1
153
+ assert score.usage.prompt_tokens == 108
154
+ assert score.data[0].score == pytest.approx(TEXT_VS_TEXT_PLUS_IMAGE, rel=TOL)
155
+
156
+
157
+def test_score_api_queries_str_documents_list(server: RemoteOpenAIServer):
158
+ score_response = requests.post(
159
+ server.url_for("score"),
160
+ json={
161
+ "model": MODEL_NAME,
162
+ "queries": query,
163
+ "documents": [
164
+ document,
165
+ {"content": [documents[0]]},
166
+ {"content": [documents[1]]},
167
+ {"content": [documents[0], documents[1]]},
168
+ ],
169
+ },
170
+ )
171
+ score_response.raise_for_status()
172
+ score = ScoreResponse.model_validate(score_response.json())
173
+
174
+ assert score.id is not None
175
+ assert score.data is not None
176
+ assert len(score.data) == 4
177
+ assert score.usage.prompt_tokens == 368
178
+ assert score.data[0].score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
179
+ assert score.data[1].score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
180
+ assert score.data[2].score == pytest.approx(TEXT_VS_IMAGE, rel=TOL)
181
+ assert score.data[3].score == pytest.approx(TEXT_VS_TEXT_PLUS_IMAGE, rel=TOL)
182
+
183
+
184
+def test_rerank_api_queries_str_documents_list(server: RemoteOpenAIServer):
185
+ rerank_response = requests.post(
186
+ server.url_for("rerank"),
187
+ json={
188
+ "model": MODEL_NAME,
189
+ "query": query,
190
+ "documents": [
191
+ document,
192
+ {"content": [documents[0]]},
193
+ {"content": [documents[1]]},
194
+ {"content": [documents[0], documents[1]]},
195
+ ],
196
+ },
197
+ )
198
+ rerank_response.raise_for_status()
199
+ rerank = RerankResponse.model_validate(rerank_response.json())
200
+
201
+ assert rerank.id is not None
202
+ assert rerank.model is not None
203
+ assert rerank.usage is not None
204
+ assert len(rerank.results) == 4
205
+
206
+ rerank.results.sort(key=lambda x: x.index)
207
+ assert rerank.results[0].relevance_score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
208
+ assert rerank.results[1].relevance_score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
209
+ assert rerank.results[2].relevance_score == pytest.approx(TEXT_VS_IMAGE, rel=TOL)
210
+ assert rerank.results[3].relevance_score == pytest.approx(
211
+ TEXT_VS_TEXT_PLUS_IMAGE, rel=TOL
212
+ )
213
+
214
+
215
+def test_score_api_queries_list_documents_list(server: RemoteOpenAIServer):
216
+ score_response = requests.post(
217
+ server.url_for("score"),
218
+ json={
219
+ "model": MODEL_NAME,
220
+ "queries": [query] * 4,
221
+ "documents": [
222
+ document,
223
+ {"content": [documents[0]]},
224
+ {"content": [documents[1]]},
225
+ {"content": [documents[0], documents[1]]},
226
+ ],
227
+ },
228
+ )
229
+ score_response.raise_for_status()
230
+ score = ScoreResponse.model_validate(score_response.json())
231
+
232
+ assert score.id is not None
233
+ assert score.data is not None
234
+ assert len(score.data) == 4
235
+ assert score.usage.prompt_tokens == 368
236
+ assert score.data[0].score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
237
+ assert score.data[1].score == pytest.approx(TEXT_VS_TEXT, rel=TOL)
238
+ assert score.data[2].score == pytest.approx(TEXT_VS_IMAGE, rel=TOL)
239
+ assert score.data[3].score == pytest.approx(TEXT_VS_TEXT_PLUS_IMAGE, rel=TOL)
vllm/entrypoints/llm.py
+37
-73
@@ -40,12 +40,12 @@ from vllm.entrypoints.chat_utils import (
40
ChatTemplateContentFormatOption,
41
)
42
from vllm.entrypoints.pooling.score.utils import (
43
- ScoreContentPartParam,
43
+ ScoreData,
44
ScoreMultiModalParam,
45
_cosine_similarity,
46
- _validate_score_input_lens,
46
compress_token_type_ids,
47
get_score_prompt,
48
+ validate_score_input,
49
)
50
from vllm.entrypoints.utils import log_non_default_args
51
from vllm.inputs import (
@@ -1326,8 +1326,8 @@ class LLM:
1326
1327
def _embedding_score(
1328
self,
1329
- text_1: list[SingletonPrompt],
1330
- text_2: list[SingletonPrompt],
1329
+ data_1: list[ScoreData],
1330
+ data_2: list[ScoreData],
1331
*,
1332
use_tqdm: bool | Callable[..., tqdm],
1333
pooling_params: PoolingParams | None,
@@ -1336,8 +1336,16 @@ class LLM:
1336
) -> list[ScoringRequestOutput]:
1337
tokenizer = self.get_tokenizer()
1338
1339
+ input_texts: list[str] = []
1340
+ for text in data_1 + data_2:
1341
+ if not isinstance(text, str):
1342
+ raise NotImplementedError(
1343
+ "Embedding scores currently do not support multimodal input."
1344
+ )
1345
+ input_texts.append(text)
1346
+
1347
encoded_output = self.encode(
1340
- text_1 + text_2,
1348
+ input_texts,
1349
use_tqdm=use_tqdm,
1350
lora_request=lora_request,
1351
pooling_params=pooling_params,
@@ -1345,8 +1353,8 @@ class LLM:
1353
tokenization_kwargs=tokenization_kwargs,
1354
)
1355
1348
- encoded_output_1 = encoded_output[0 : len(text_1)]
1349
- encoded_output_2 = encoded_output[len(text_1) :]
1356
+ encoded_output_1 = encoded_output[0 : len(data_1)]
1357
+ encoded_output_2 = encoded_output[len(data_1) :]
1358
1359
if len(encoded_output_1) == 1:
1360
encoded_output_1 = encoded_output_1 * len(encoded_output_2)
@@ -1362,8 +1370,8 @@ class LLM:
1370
1371
def _cross_encoding_score(
1372
self,
1365
- data_1: list[str] | list[ScoreContentPartParam],
1366
- data_2: list[str] | list[ScoreContentPartParam],
1373
+ data_1: list[ScoreData],
1374
+ data_2: list[ScoreData],
1375
*,
1376
use_tqdm: bool | Callable[..., tqdm],
1377
pooling_params: PoolingParams | None,
@@ -1424,8 +1432,14 @@ class LLM:
1432
1433
def score(
1434
self,
1427
- data_1: SingletonPrompt | Sequence[SingletonPrompt] | ScoreMultiModalParam,
1428
- data_2: SingletonPrompt | Sequence[SingletonPrompt] | ScoreMultiModalParam,
1435
+ data_1: SingletonPrompt
1436
+ | Sequence[SingletonPrompt]
1437
+ | ScoreMultiModalParam
1438
+ | list[ScoreMultiModalParam],
1439
+ data_2: SingletonPrompt
1440
+ | Sequence[SingletonPrompt]
1441
+ | ScoreMultiModalParam
1442
+ | list[ScoreMultiModalParam],
1443
/,
1444
*,
1445
use_tqdm: bool | Callable[..., tqdm] = True,
@@ -1501,73 +1515,23 @@ class LLM:
1515
"chat_template is only supported for cross-encoder models."
1516
)
1517
1504
- # the tokenizer for models such as
1505
- # "cross-encoder/ms-marco-MiniLM-L-6-v2" doesn't support passing
1506
- # lists of tokens to the `text` and `text_pair` kwargs
1507
- tokenizer = self.get_tokenizer()
1508
-
1509
- if not model_config.is_multimodal_model:
1510
-
1511
- def check_data_type(
1512
- data: SingletonPrompt
1513
- | Sequence[SingletonPrompt]
1514
- | ScoreMultiModalParam,
1515
- ):
1516
- if isinstance(data, dict) and "content" in data:
1517
- raise ValueError(
1518
- "ScoreMultiModalParam is not supported "
1519
- f"for {model_config.architecture}"
1520
- )
1521
-
1522
- check_data_type(data_1)
1523
- check_data_type(data_2)
1524
-
1525
- def ensure_str(prompt: SingletonPrompt):
1526
- if isinstance(prompt, dict):
1527
- if "multi_modal_data" in prompt:
1528
- raise ValueError(
1529
- "Multi-modal prompt is not supported for scoring"
1530
- )
1531
- elif "prompt_token_ids" in prompt:
1532
- prompt = tokenizer.decode(
1533
- cast(TokensPrompt, prompt)["prompt_token_ids"]
1534
- )
1535
- elif "prompt" in prompt:
1536
- prompt = cast(TextPrompt, prompt)["prompt"]
1537
- assert type(prompt) is str
1538
- return prompt
1518
+ is_multimodal_model = model_config.is_multimodal_model
1519
+ architecture = model_config.architecture
1520
1540
- if isinstance(data_1, (str, dict)):
1541
- # Convert a single prompt to a list.
1542
- data_1 = [data_1] # type: ignore[list-item]
1543
-
1544
- data_1 = [ensure_str(t) for t in data_1]
1545
-
1546
- if isinstance(data_2, (str, dict)):
1547
- # Convert a single prompt to a list.
1548
- data_2 = [data_2] # type: ignore[list-item]
1549
-
1550
- data_2 = [ensure_str(t) for t in data_2]
1551
-
1552
- if isinstance(data_1, dict) and "content" in data_1:
1553
- data_1 = data_1.get("content") # type: ignore[assignment]
1554
- elif isinstance(data_1, str):
1555
- data_1 = [data_1]
1556
-
1557
- if isinstance(data_2, dict) and "content" in data_2:
1558
- data_2 = data_2.get("content") # type: ignore[assignment]
1559
- elif isinstance(data_2, str):
1560
- data_2 = [data_2]
1561
-
1562
- _validate_score_input_lens(data_1, data_2) # type: ignore[arg-type]
1521
+ score_data_1, score_data_2 = validate_score_input(
1522
+ data_1, # type: ignore[arg-type]
1523
+ data_2, # type: ignore[arg-type]
1524
+ is_multimodal_model=is_multimodal_model,
1525
+ architecture=architecture,
1526
+ )
1527
1528
tok_params = self._get_cmpl_tok_params(tokenization_kwargs)
1529
encode_kwargs = tok_params.get_encode_kwargs()
1530
1531
if model_config.is_cross_encoder:
1532
return self._cross_encoding_score(
1569
- data_1, # type: ignore[arg-type]
1570
- data_2, # type: ignore[arg-type]
1533
+ score_data_1,
1534
+ score_data_2,
1535
use_tqdm=use_tqdm,
1536
pooling_params=pooling_params,
1537
lora_request=lora_request,
@@ -1576,8 +1540,8 @@ class LLM:
1540
)
1541
else:
1542
return self._embedding_score(
1579
- data_1, # type: ignore[arg-type]
1580
- data_2, # type: ignore[arg-type]
1543
+ score_data_1,
1544
+ score_data_2,
1545
use_tqdm=use_tqdm,
1546
pooling_params=pooling_params,
1547
lora_request=lora_request,
vllm/entrypoints/pooling/score/protocol.py
+28
-11
@@ -14,7 +14,8 @@ from vllm.entrypoints.pooling.base.protocol import (
14
)
15
from vllm.entrypoints.pooling.score.utils import (
16
ScoreContentPartParam,
17
- ScoreMultiModalParam,
17
+ ScoreInput,
18
+ ScoreInputs,
19
)
20
from vllm.renderers import TokenizeParams
21
from vllm.utils import random_uuid
@@ -47,13 +48,13 @@ class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin):
48
49
50
class ScoreDataRequest(ScoreRequestMixin):
50
- data_1: list[str] | str | ScoreMultiModalParam
51
- data_2: list[str] | str | ScoreMultiModalParam
51
+ data_1: ScoreInputs
52
+ data_2: ScoreInputs
53
54
55
class ScoreQueriesDocumentsRequest(ScoreRequestMixin):
55
- queries: list[str] | str | ScoreMultiModalParam
56
- documents: list[str] | str | ScoreMultiModalParam
56
+ queries: ScoreInputs
57
+ documents: ScoreInputs
58
59
@property
60
def data_1(self):
@@ -64,9 +65,22 @@ class ScoreQueriesDocumentsRequest(ScoreRequestMixin):
65
return self.documents
66
67
68
+class ScoreQueriesItemsRequest(ScoreRequestMixin):
69
+ queries: ScoreInputs
70
+ items: ScoreInputs
71
+
72
+ @property
73
+ def data_1(self):
74
+ return self.queries
75
+
76
+ @property
77
+ def data_2(self):
78
+ return self.items
79
+
80
+
81
class ScoreTextRequest(ScoreRequestMixin):
68
- text_1: list[str] | str | ScoreMultiModalParam
69
- text_2: list[str] | str | ScoreMultiModalParam
82
+ text_1: ScoreInputs
83
+ text_2: ScoreInputs
84
85
@property
86
def data_1(self):
@@ -78,13 +92,16 @@ class ScoreTextRequest(ScoreRequestMixin):
92
93
94
ScoreRequest: TypeAlias = (
81
- ScoreQueriesDocumentsRequest | ScoreDataRequest | ScoreTextRequest
95
+ ScoreQueriesDocumentsRequest
96
+ | ScoreQueriesItemsRequest
97
+ | ScoreDataRequest
98
+ | ScoreTextRequest
99
)
100
101
102
class RerankRequest(PoolingBasicRequestMixin, ClassifyRequestMixin):
86
- query: str | ScoreMultiModalParam
87
- documents: list[str] | ScoreMultiModalParam
103
+ query: ScoreInput
104
+ documents: ScoreInputs
105
top_n: int = Field(default_factory=lambda: 0)
106
107
# --8<-- [start:rerank-extra-params]
@@ -108,7 +125,7 @@ class RerankRequest(PoolingBasicRequestMixin, ClassifyRequestMixin):
125
126
class RerankDocument(BaseModel):
127
text: str | None = None
111
- multi_modal: ScoreContentPartParam | None = None
128
+ multi_modal: list[ScoreContentPartParam] | None = None
129
130
131
class RerankResult(BaseModel):
vllm/entrypoints/pooling/score/serving.py
+84
-94
@@ -27,12 +27,12 @@ from vllm.entrypoints.pooling.score.protocol import (
27
ScoreResponseData,
28
)
29
from vllm.entrypoints.pooling.score.utils import (
30
- ScoreContentPartParam,
31
- ScoreMultiModalParam,
30
+ ScoreData,
31
+ ScoreInputs,
32
_cosine_similarity,
33
- _validate_score_input_lens,
33
compress_token_type_ids,
34
get_score_prompt,
35
+ validate_score_input,
36
)
37
from vllm.inputs.data import TokensPrompt
38
from vllm.logger import init_logger
@@ -65,15 +65,32 @@ class ServingScores(OpenAIServing):
65
66
self._tokenizer_executor = ThreadPoolExecutor(max_workers=1)
67
68
+ self.is_cross_encoder = self.model_config.is_cross_encoder
69
+ self.is_multimodal_model = self.model_config.is_multimodal_model
70
+ self.architecture = self.model_config.architecture
71
+
72
+ if self.is_cross_encoder:
73
+ self._score_func = self._cross_encoding_score
74
+ else:
75
+ self._score_func = self._embedding_score
76
+
77
async def _embedding_score(
78
self,
70
- data_1: list[str],
71
- data_2: list[str],
79
+ data_1: list[ScoreData],
80
+ data_2: list[ScoreData],
81
request: RerankRequest | ScoreRequest,
82
request_id: str,
83
lora_request: LoRARequest | None | None = None,
84
trace_headers: Mapping[str, str] | None = None,
85
) -> list[PoolingRequestOutput] | ErrorResponse:
86
+ input_texts: list[str] = []
87
+ for text in data_1 + data_2:
88
+ if not isinstance(text, str):
89
+ raise NotImplementedError(
90
+ "Embedding scores currently do not support multimodal input."
91
+ )
92
+ input_texts.append(text)
93
+
94
model_config = self.model_config
95
tokenizer = self.renderer.get_tokenizer()
96
@@ -82,8 +99,6 @@ class ServingScores(OpenAIServing):
99
executor=self._tokenizer_executor,
100
)
101
85
- input_texts = data_1 + data_2
86
-
102
tokenization_kwargs = request.build_tok_params(model_config).get_encode_kwargs()
103
tokenized_prompts = await asyncio.gather(
104
*(encode_async(t, **tokenization_kwargs) for t in input_texts)
@@ -157,60 +172,30 @@ class ServingScores(OpenAIServing):
172
173
return final_res_batch
174
160
- def _preprocess_score(
161
- self,
162
- request: RerankRequest | ScoreRequest,
163
- tokenizer: TokenizerLike,
164
- tokenization_kwargs: dict[str, Any],
165
- data_1: str | ScoreContentPartParam,
166
- data_2: str | ScoreContentPartParam,
167
- ) -> tuple[str, TokensPrompt]:
168
- model_config = self.model_config
169
-
170
- full_prompt, engine_prompt = get_score_prompt(
171
- model_config=model_config,
172
- data_1=data_1,
173
- data_2=data_2,
174
- tokenizer=tokenizer,
175
- tokenization_kwargs=tokenization_kwargs,
176
- score_template=self.score_template,
177
- )
178
- self._validate_input(request, engine_prompt["prompt_token_ids"], full_prompt)
179
- if request.mm_processor_kwargs is not None:
180
- engine_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs
181
-
182
- return full_prompt, engine_prompt
183
-
175
async def _cross_encoding_score(
176
self,
186
- data_1: list[str] | list[ScoreContentPartParam],
187
- data_2: list[str] | list[ScoreContentPartParam],
177
+ data_1: list[ScoreData],
178
+ data_2: list[ScoreData],
179
request: RerankRequest | ScoreRequest,
180
request_id: str,
181
lora_request: LoRARequest | None | None = None,
182
trace_headers: Mapping[str, str] | None = None,
183
) -> list[PoolingRequestOutput] | ErrorResponse:
193
- model_config = self.model_config
184
tokenizer = self.renderer.get_tokenizer()
185
+ if isinstance(tokenizer, MistralTokenizer):
186
+ raise ValueError("MistralTokenizer not supported for cross-encoding")
187
196
- request_prompts: list[str] = []
197
- engine_prompts: list[TokensPrompt] = []
188
+ model_config = self.model_config
189
190
if len(data_1) == 1:
191
data_1 = data_1 * len(data_2)
192
202
- if isinstance(tokenizer, MistralTokenizer):
203
- raise ValueError("MistralTokenizer not supported for cross-encoding")
204
-
193
tok_kwargs = request.build_tok_params(model_config).get_encode_kwargs()
206
-
194
input_pairs = [(t1, t2) for t1, t2 in zip(data_1, data_2)]
208
-
195
preprocess_async = make_async(
196
self._preprocess_score,
197
executor=self._tokenizer_executor,
198
)
213
-
199
preprocessed_prompts = await asyncio.gather(
200
*(
201
preprocess_async(
@@ -224,6 +209,8 @@ class ServingScores(OpenAIServing):
209
)
210
)
211
212
+ request_prompts: list[str] = []
213
+ engine_prompts: list[TokensPrompt] = []
214
for full_prompt, engine_prompt in preprocessed_prompts:
215
request_prompts.append(full_prompt)
216
engine_prompts.append(engine_prompt)
@@ -278,10 +265,33 @@ class ServingScores(OpenAIServing):
265
266
return [out for out in final_res_batch if out is not None]
267
268
+ def _preprocess_score(
269
+ self,
270
+ request: RerankRequest | ScoreRequest,
271
+ tokenizer: TokenizerLike,
272
+ tokenization_kwargs: dict[str, Any],
273
+ data_1: ScoreData,
274
+ data_2: ScoreData,
275
+ ) -> tuple[str, TokensPrompt]:
276
+ model_config = self.model_config
277
+ full_prompt, engine_prompt = get_score_prompt(
278
+ model_config=model_config,
279
+ data_1=data_1,
280
+ data_2=data_2,
281
+ tokenizer=tokenizer,
282
+ tokenization_kwargs=tokenization_kwargs,
283
+ score_template=self.score_template,
284
+ )
285
+ self._validate_input(request, engine_prompt["prompt_token_ids"], full_prompt)
286
+ if request.mm_processor_kwargs is not None:
287
+ engine_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs
288
+
289
+ return full_prompt, engine_prompt
290
+
291
async def _run_scoring(
292
self,
283
- data_1: list[str] | str | ScoreMultiModalParam,
284
- data_2: list[str] | str | ScoreMultiModalParam,
293
+ data_1: ScoreInputs,
294
+ data_2: ScoreInputs,
295
request: ScoreRequest | RerankRequest,
296
request_id: str,
297
raw_request: Request | None = None,
@@ -294,44 +304,21 @@ class ServingScores(OpenAIServing):
304
else await self._get_trace_headers(raw_request.headers)
305
)
306
297
- if not self.model_config.is_multimodal_model and (
298
- isinstance(data_1, dict) or isinstance(data_2, dict)
299
- ):
300
- raise ValueError(
301
- f"MultiModalParam is not supported for {self.model_config.architecture}" # noqa: E501
302
- )
303
-
304
- if isinstance(data_1, str):
305
- data_1 = [data_1]
306
- elif isinstance(data_1, dict):
307
- data_1 = data_1.get("content") # type: ignore[assignment]
308
-
309
- if isinstance(data_2, str):
310
- data_2 = [data_2]
311
- elif isinstance(data_2, dict):
312
- data_2 = data_2.get("content") # type: ignore[assignment]
313
-
314
- _validate_score_input_lens(data_1, data_2) # type: ignore[arg-type]
315
-
316
- if self.model_config.is_cross_encoder:
317
- return await self._cross_encoding_score(
318
- data_1=data_1, # type: ignore[arg-type]
319
- data_2=data_2, # type: ignore[arg-type]
320
- request=request,
321
- request_id=request_id,
322
- lora_request=lora_request,
323
- trace_headers=trace_headers,
324
- )
307
+ score_data_1, score_data_2 = validate_score_input(
308
+ data_1,
309
+ data_2,
310
+ is_multimodal_model=self.is_multimodal_model,
311
+ architecture=self.architecture,
312
+ )
313
326
- else:
327
- return await self._embedding_score(
328
- data_1=data_1, # type: ignore[arg-type]
329
- data_2=data_2, # type: ignore[arg-type]
330
- request=request,
331
- request_id=request_id,
332
- lora_request=lora_request,
333
- trace_headers=trace_headers,
334
- )
314
+ return await self._score_func(
315
+ data_1=score_data_1,
316
+ data_2=score_data_2,
317
+ request=request,
318
+ request_id=request_id,
319
+ lora_request=lora_request,
320
+ trace_headers=trace_headers,
321
+ )
322
323
async def create_score(
324
self,
@@ -391,15 +378,6 @@ class ServingScores(OpenAIServing):
378
379
request_id = f"rerank-{self._base_request_id(raw_request)}"
380
documents = request.documents
394
- top_n = (
395
- request.top_n
396
- if request.top_n > 0
397
- else (
398
- len(documents)
399
- if isinstance(documents, list)
400
- else len(documents["content"])
401
- )
402
- )
381
382
try:
383
final_res_batch = await self._run_scoring(
@@ -412,6 +390,8 @@ class ServingScores(OpenAIServing):
390
if isinstance(final_res_batch, ErrorResponse):
391
return final_res_batch
392
393
+ top_n = request.top_n if request.top_n > 0 else len(final_res_batch)
394
+
395
return self.request_output_to_rerank_response(
396
final_res_batch,
397
request_id,
@@ -465,22 +445,32 @@ class ServingScores(OpenAIServing):
445
final_res_batch: list[PoolingRequestOutput],
446
request_id: str,
447
model_name: str,
468
- documents: list[str] | ScoreMultiModalParam,
448
+ documents: ScoreInputs,
449
top_n: int,
450
) -> RerankResponse:
451
"""
452
Convert the output of do_rank to a RerankResponse
453
"""
454
+
455
+ if not isinstance(documents, list):
456
+ documents = [documents]
457
+
458
results: list[RerankResult] = []
459
num_prompt_tokens = 0
460
for idx, final_res in enumerate(final_res_batch):
461
classify_res = ScoringRequestOutput.from_base(final_res)
462
463
+ document = documents[idx]
464
+ if isinstance(document, str):
465
+ rerank_document = RerankDocument(text=document)
466
+ else:
467
+ rerank_document = RerankDocument(
468
+ multi_modal=document.get("content", [])
469
+ )
470
+
471
result = RerankResult(
472
index=idx,
481
- document=RerankDocument(text=documents[idx])
482
- if isinstance(documents, list)
483
- else RerankDocument(multi_modal=documents["content"][idx]),
473
+ document=rerank_document,
474
relevance_score=classify_res.outputs.score,
475
)
476
results.append(result)
vllm/entrypoints/pooling/score/utils.py
+69
-20
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Iterable
4
from typing import Any, TypeAlias, cast
5
6
from torch.nn import CosineSimilarity
@@ -10,12 +11,13 @@ from vllm.entrypoints.chat_utils import (
11
BaseMultiModalItemTracker,
12
ChatCompletionContentPartImageEmbedsParam,
13
ChatCompletionContentPartImageParam,
14
+ ChatCompletionContentPartParam,
15
ChatCompletionContentPartTextParam,
16
ChatCompletionContentPartVideoParam,
17
ChatTemplateResolutionError,
18
+ ConversationMessage,
19
MultiModalItemTracker,
17
- _ContentPart,
18
- _parse_chat_message_content_part,
20
+ _parse_chat_message_content_parts,
21
)
22
from vllm.inputs import TokensPrompt
23
from vllm.model_executor.models.interfaces import supports_score_template
@@ -46,6 +48,13 @@ class ScoreMultiModalParam(TypedDict, total=False):
48
"""The multimodal contents"""
49
50
51
+# Raw input data with content key in ScoreMultiModalParam.
52
+ScoreInput = str | ScoreMultiModalParam
53
+ScoreInputs = ScoreInput | list[ScoreInput]
54
+# Score data without content key.
55
+ScoreData = str | list[ScoreContentPartParam]
56
+
57
+
58
def _cosine_similarity(
59
tokenizer: TokenizerLike,
60
embed_1: list[PoolingRequestOutput],
@@ -77,8 +86,8 @@ def _cosine_similarity(
86
87
88
def _validate_score_input_lens(
80
- data_1: list[str] | list[ScoreContentPartParam],
81
- data_2: list[str] | list[ScoreContentPartParam],
89
+ data_1: list[ScoreData],
90
+ data_2: list[ScoreData],
91
):
92
len_1 = len(data_1)
93
len_2 = len(data_2)
@@ -91,19 +100,56 @@ def _validate_score_input_lens(
100
raise ValueError("At least one text_pair element must be given")
101
102
103
+def _validate_mm_score_input(
104
+ data: list[ScoreInput],
105
+ is_multimodal_model: bool,
106
+ architecture: str,
107
+) -> list[ScoreData]:
108
+ out: list[ScoreData] = []
109
+ for d in data:
110
+ if isinstance(d, str):
111
+ out.append(d)
112
+ else:
113
+ if not is_multimodal_model:
114
+ raise ValueError(f"MultiModalParam is not supported for {architecture}")
115
+ content = cast(list[ScoreContentPartParam], d.get("content", []))
116
+ out.append(content)
117
+ return out
118
+
119
+
120
+def validate_score_input(
121
+ data_1: ScoreInputs,
122
+ data_2: ScoreInputs,
123
+ is_multimodal_model: bool,
124
+ architecture: str,
125
+) -> tuple[list[ScoreData], list[ScoreData]]:
126
+ if not isinstance(data_1, list):
127
+ data_1 = [data_1]
128
+
129
+ if not isinstance(data_2, list):
130
+ data_2 = [data_2]
131
+
132
+ score_input_1 = _validate_mm_score_input(data_1, is_multimodal_model, architecture)
133
+ score_input_2 = _validate_mm_score_input(data_2, is_multimodal_model, architecture)
134
+ _validate_score_input_lens(score_input_1, score_input_2)
135
+ return score_input_1, score_input_2
136
+
137
+
138
def parse_score_data(
95
- data_1: str | ScoreContentPartParam,
96
- data_2: str | ScoreContentPartParam,
139
+ data_1: ScoreData,
140
+ data_2: ScoreData,
141
model_config: ModelConfig,
142
) -> tuple[str, str, MultiModalDataDict | None, MultiModalUUIDDict | None]:
143
mm_tracker = MultiModalItemTracker(model_config)
144
101
- content_1 = _parse_score_content(data_1, mm_tracker)
102
- content_2 = _parse_score_content(data_2, mm_tracker)
145
+ content_1 = _parse_score_content("query", data_1, mm_tracker)
146
+ content_2 = _parse_score_content("document", data_2, mm_tracker)
147
104
- def ensure_str(content: _ContentPart | None) -> str:
105
- if content is not None and isinstance(content, str):
106
- return cast(str, content)
148
+ def ensure_str(content: list[ConversationMessage]) -> str:
149
+ assert len(content) == 1
150
+ prompt = content[0]["content"]
151
+ if prompt is not None and isinstance(prompt, str):
152
+ return cast(str, prompt)
153
else:
154
raise ValueError(f"Only string content is supported, but got {content}.")
155
@@ -115,19 +161,22 @@ def parse_score_data(
161
162
163
def _parse_score_content(
118
- data: str | ScoreContentPartParam,
164
+ role: str,
165
+ data: ScoreData,
166
mm_tracker: BaseMultiModalItemTracker,
120
-) -> _ContentPart | None:
167
+) -> list[ConversationMessage]:
168
+ parts: Iterable[ChatCompletionContentPartParam]
169
if isinstance(data, str):
122
- part = ChatCompletionContentPartTextParam(type="text", text=data)
170
+ parts = [ChatCompletionContentPartTextParam(type="text", text=data)]
171
else:
124
- part = data
172
+ parts = cast(Iterable[ChatCompletionContentPartParam], data)
173
174
mm_parser = mm_tracker.create_parser()
175
128
- parse_res = _parse_chat_message_content_part(
129
- part,
130
- mm_parser,
176
+ parse_res = _parse_chat_message_content_parts(
177
+ role=role,
178
+ parts=parts,
179
+ mm_tracker=mm_tracker,
180
wrap_dicts=False,
181
interleave_strings=False,
182
)
@@ -184,8 +233,8 @@ def get_score_prompt(
233
model_config: ModelConfig,
234
tokenizer: TokenizerLike,
235
tokenization_kwargs: dict[str, Any],
187
- data_1: str | ScoreContentPartParam,
188
- data_2: str | ScoreContentPartParam,
236
+ data_1: ScoreData,
237
+ data_2: ScoreData,
238
score_template: str | None = None,
239
) -> tuple[str, TokensPrompt]:
240
prompt_1, prompt_2, mm_data, mm_uuids = parse_score_data(