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 10, 00:15 UTC
[Frontend] Score entrypoint support data_1 & data_2 and queries & documents as inputs (#32577)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
wang.yuqi committed
Jan 19, 2026 at 22:07 UTC
c88860d759715affaf2285afb953cc791c0a5a38
17 files changed
+253
-113
docs/serving/openai_compatible_server.md
+15
-15
@@ -694,7 +694,7 @@ Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](.
694
695
#### Single inference
696
697
-You can pass a string to both `text_1` and `text_2`, forming a single sentence pair.
697
+You can pass a string to both `queries` and `documents`, forming a single sentence pair.
698
699
```bash
700
curl -X 'POST' \
@@ -704,8 +704,8 @@ curl -X 'POST' \
704
-d '{
705
"model": "BAAI/bge-reranker-v2-m3",
706
"encoding_format": "float",
707
- "text_1": "What is the capital of France?",
708
- "text_2": "The capital of France is Paris."
707
+ "queries": "What is the capital of France?",
708
+ "documents": "The capital of France is Paris."
709
}'
710
```
711
@@ -730,9 +730,9 @@ curl -X 'POST' \
730
731
#### Batch inference
732
733
-You can pass a string to `text_1` and a list to `text_2`, forming multiple sentence pairs
734
-where each pair is built from `text_1` and a string in `text_2`.
735
-The total number of pairs is `len(text_2)`.
733
+You can pass a string to `queries` and a list to `documents`, forming multiple sentence pairs
734
+where each pair is built from `queries` and a string in `documents`.
735
+The total number of pairs is `len(documents)`.
736
737
??? console "Request"
738
@@ -743,8 +743,8 @@ The total number of pairs is `len(text_2)`.
743
-H 'Content-Type: application/json' \
744
-d '{
745
"model": "BAAI/bge-reranker-v2-m3",
746
- "text_1": "What is the capital of France?",
747
- "text_2": [
746
+ "queries": "What is the capital of France?",
747
+ "documents": [
748
"The capital of Brazil is Brasilia.",
749
"The capital of France is Paris."
750
]
@@ -775,9 +775,9 @@ The total number of pairs is `len(text_2)`.
775
}
776
```
777
778
-You can pass a list to both `text_1` and `text_2`, forming multiple sentence pairs
779
-where each pair is built from a string in `text_1` and the corresponding string in `text_2` (similar to `zip()`).
780
-The total number of pairs is `len(text_2)`.
778
+You can pass a list to both `queries` and `documents`, forming multiple sentence pairs
779
+where each pair is built from a string in `queries` and the corresponding string in `documents` (similar to `zip()`).
780
+The total number of pairs is `len(documents)`.
781
782
??? console "Request"
783
@@ -789,11 +789,11 @@ The total number of pairs is `len(text_2)`.
789
-d '{
790
"model": "BAAI/bge-reranker-v2-m3",
791
"encoding_format": "float",
792
- "text_1": [
792
+ "queries": [
793
"What is the capital of Brazil?",
794
"What is the capital of France?"
795
],
796
- "text_2": [
796
+ "documents": [
797
"The capital of Brazil is Brasilia.",
798
"The capital of France is Paris."
799
]
@@ -847,8 +847,8 @@ You can pass multi-modal inputs to scoring models by passing `content` including
847
"http://localhost:8000/v1/score",
848
json={
849
"model": "jinaai/jina-reranker-m0",
850
- "text_1": "slm markdown",
851
- "text_2": {
850
+ "queries": "slm markdown",
851
+ "documents": {
852
"content": [
853
{
854
"type": "image_url",
examples/offline_inference/basic/score.py
+5
-5
@@ -21,8 +21,8 @@ def parse_args():
21
22
def main(args: Namespace):
23
# Sample prompts.
24
- text_1 = "What is the capital of France?"
25
- texts_2 = [
24
+ query = "What is the capital of France?"
25
+ documents = [
26
"The capital of Brazil is Brasilia.",
27
"The capital of France is Paris.",
28
]
@@ -32,13 +32,13 @@ def main(args: Namespace):
32
llm = LLM(**vars(args))
33
34
# Generate scores. The output is a list of ScoringRequestOutputs.
35
- outputs = llm.score(text_1, texts_2)
35
+ outputs = llm.score(query, documents)
36
37
# Print the outputs.
38
print("\nGenerated Outputs:\n" + "-" * 60)
39
- for text_2, output in zip(texts_2, outputs):
39
+ for document, output in zip(documents, outputs):
40
score = output.outputs.score
41
- print(f"Pair: {[text_1, text_2]!r} \nScore: {score}")
41
+ print(f"Pair: {[query, document]!r} \nScore: {score}")
42
print("-" * 60)
43
44
examples/offline_inference/openai_batch/README.md
+2
-2
@@ -255,8 +255,8 @@ cat results.jsonl
255
Add score requests to your batch file. The following is an example:
256
257
```text
258
-{"custom_id": "request-1", "method": "POST", "url": "/v1/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "text_1": "What is the capital of France?", "text_2": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
259
-{"custom_id": "request-2", "method": "POST", "url": "/v1/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "text_1": "What is the capital of France?", "text_2": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
258
+{"custom_id": "request-1", "method": "POST", "url": "/v1/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "queries": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
259
+{"custom_id": "request-2", "method": "POST", "url": "/v1/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "queries": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
260
```
261
262
You can mix chat completion, embedding, and score requests in the batch file, as long as the model you are using supports them all (note that all requests must use the same model).
examples/pooling/score/cohere_rerank_client.py
renamed
examples/pooling/score/qwen3_reranker_online.py
+2
-2
@@ -50,8 +50,8 @@ documents = [
50
# Request payload for the score API
51
data = {
52
"model": "Qwen/Qwen3-Reranker-0.6B",
53
- "text_1": queries,
54
- "text_2": documents,
53
+ "queries": queries,
54
+ "documents": documents,
55
}
56
57
examples/pooling/score/score_api_online.py
+18
-12
@@ -30,29 +30,35 @@ def main(args):
30
api_url = f"http://{args.host}:{args.port}/score"
31
model_name = args.model
32
33
- text_1 = "What is the capital of Brazil?"
34
- text_2 = "The capital of Brazil is Brasilia."
35
- prompt = {"model": model_name, "text_1": text_1, "text_2": text_2}
33
+ queries = "What is the capital of Brazil?"
34
+ documents = "The capital of Brazil is Brasilia."
35
+ prompt = {"model": model_name, "queries": queries, "documents": documents}
36
score_response = post_http_request(prompt=prompt, api_url=api_url)
37
- print("\nPrompt when text_1 and text_2 are both strings:")
37
+ print("\nPrompt when queries and documents are both strings:")
38
pprint.pprint(prompt)
39
print("\nScore Response:")
40
pprint.pprint(score_response.json())
41
42
- text_1 = "What is the capital of France?"
43
- text_2 = ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]
44
- prompt = {"model": model_name, "text_1": text_1, "text_2": text_2}
42
+ queries = "What is the capital of France?"
43
+ documents = [
44
+ "The capital of Brazil is Brasilia.",
45
+ "The capital of France is Paris.",
46
+ ]
47
+ prompt = {"model": model_name, "queries": queries, "documents": documents}
48
score_response = post_http_request(prompt=prompt, api_url=api_url)
46
- print("\nPrompt when text_1 is string and text_2 is a list:")
49
+ print("\nPrompt when queries is string and documents is a list:")
50
pprint.pprint(prompt)
51
print("\nScore Response:")
52
pprint.pprint(score_response.json())
53
51
- text_1 = ["What is the capital of Brazil?", "What is the capital of France?"]
52
- text_2 = ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]
53
- prompt = {"model": model_name, "text_1": text_1, "text_2": text_2}
54
+ queries = ["What is the capital of Brazil?", "What is the capital of France?"]
55
+ documents = [
56
+ "The capital of Brazil is Brasilia.",
57
+ "The capital of France is Paris.",
58
+ ]
59
+ prompt = {"model": model_name, "queries": queries, "documents": documents}
60
score_response = post_http_request(prompt=prompt, api_url=api_url)
55
- print("\nPrompt when text_1 and text_2 are both lists:")
61
+ print("\nPrompt when queries and documents are both lists:")
62
pprint.pprint(prompt)
63
print("\nScore Response:")
64
pprint.pprint(score_response.json())
examples/pooling/score/vision_rerank_api_online.py
+20
-2
@@ -18,10 +18,22 @@ e.g.
18
"""
19
20
import argparse
21
+import base64
22
import json
23
24
import requests
25
26
+
27
+def encode_base64_content_from_url(content_url: str) -> dict[str, str]:
28
+ """Encode a content retrieved from a remote url to base64 format."""
29
+
30
+ with requests.get(content_url, headers=headers) as response:
31
+ response.raise_for_status()
32
+ result = base64.b64encode(response.content).decode("utf-8")
33
+
34
+ return {"url": f"data:image/jpeg;base64,{result}"}
35
+
36
+
37
headers = {"accept": "application/json", "Content-Type": "application/json"}
38
39
query = "A woman playing with her dog on a beach at sunset."
@@ -30,8 +42,8 @@ documents = {
42
{
43
"type": "text",
44
"text": (
33
- "A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset, " # noqa: E501
34
- "as the dog offers its paw in a heartwarming display of companionship and trust." # noqa: E501
45
+ "A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset, "
46
+ "as the dog offers its paw in a heartwarming display of companionship and trust."
47
),
48
},
49
{
@@ -40,6 +52,12 @@ documents = {
52
"url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
53
},
54
},
55
+ {
56
+ "type": "image_url",
57
+ "image_url": encode_base64_content_from_url(
58
+ "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
59
+ ),
60
+ },
61
]
62
}
63
examples/pooling/score/vision_score_api_online.py
+22
-4
@@ -17,15 +17,27 @@ e.g.
17
"""
18
19
import argparse
20
+import base64
21
import json
22
import pprint
23
24
import requests
25
26
+
27
+def encode_base64_content_from_url(content_url: str) -> dict[str, str]:
28
+ """Encode a content retrieved from a remote url to base64 format."""
29
+
30
+ with requests.get(content_url, headers=headers) as response:
31
+ response.raise_for_status()
32
+ result = base64.b64encode(response.content).decode("utf-8")
33
+
34
+ return {"url": f"data:image/jpeg;base64,{result}"}
35
+
36
+
37
headers = {"accept": "application/json", "Content-Type": "application/json"}
38
27
-text_1 = "slm markdown"
28
-text_2 = {
39
+queries = "slm markdown"
40
+documents = {
41
"content": [
42
{
43
"type": "image_url",
@@ -39,6 +51,12 @@ text_2 = {
51
"url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
52
},
53
},
54
+ {
55
+ "type": "image_url",
56
+ "image_url": encode_base64_content_from_url(
57
+ "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
58
+ ),
59
+ },
60
]
61
}
62
@@ -58,9 +76,9 @@ def main(args):
76
response = requests.get(models_url, headers=headers)
77
model = response.json()["data"][0]["id"]
78
61
- prompt = {"model": model, "text_1": text_1, "text_2": text_2}
79
+ prompt = {"model": model, "queries": queries, "documents": documents}
80
response = requests.post(score_url, headers=headers, json=prompt)
63
- print("\nPrompt when text_1 is string and text_2 is a image list:")
81
+ print("\nPrompt when queries is string and documents is a image list:")
82
pprint.pprint(prompt)
83
print("\nScore Response:")
84
print(json.dumps(response.json(), indent=2))
tests/entrypoints/openai/test_run_batch.py
+2
-2
@@ -32,8 +32,8 @@ INPUT_EMBEDDING_BATCH = (
32
'{"custom_id": "request-4", "method": "POST", "url": "/v1/embeddings", "body": {"model": "NonExistModel", "input": "Hello world!"}}'
33
)
34
35
-INPUT_SCORE_BATCH = """{"custom_id": "request-1", "method": "POST", "url": "/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "text_1": "What is the capital of France?", "text_2": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
36
-{"custom_id": "request-2", "method": "POST", "url": "/v1/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "text_1": "What is the capital of France?", "text_2": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}"""
35
+INPUT_SCORE_BATCH = """{"custom_id": "request-1", "method": "POST", "url": "/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "queries": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
36
+{"custom_id": "request-2", "method": "POST", "url": "/v1/score", "body": {"model": "BAAI/bge-reranker-v2-m3", "queries": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}"""
37
38
INPUT_RERANK_BATCH = """{"custom_id": "request-1", "method": "POST", "url": "/rerank", "body": {"model": "BAAI/bge-reranker-v2-m3", "query": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
39
{"custom_id": "request-2", "method": "POST", "url": "/v1/rerank", "body": {"model": "BAAI/bge-reranker-v2-m3", "query": "What is the capital of France?", "documents": ["The capital of Brazil is Brasilia.", "The capital of France is Paris."]}}
tests/entrypoints/pooling/classify/test_online.py
+2
-2
@@ -251,8 +251,8 @@ async def test_score(server: RemoteOpenAIServer, model_name: str):
251
server.url_for("score"),
252
json={
253
"model": model_name,
254
- "text_1": "ping",
255
- "text_2": "pong",
254
+ "queries": "ping",
255
+ "documents": "pong",
256
},
257
)
258
assert response.json()["error"]["type"] == "BadRequestError"
tests/entrypoints/pooling/score/test_offline.py
+4
-4
@@ -43,12 +43,12 @@ def llm():
43
44
def test_pooling_params(llm: LLM):
45
def get_outputs(use_activation):
46
- text_1 = "What is the capital of France?"
47
- text_2 = "The capital of France is Paris."
46
+ queries = "What is the capital of France?"
47
+ documents = "The capital of France is Paris."
48
49
outputs = llm.score(
50
- text_1,
51
- text_2,
50
+ queries,
51
+ documents,
52
pooling_params=PoolingParams(use_activation=use_activation),
53
use_tqdm=False,
54
)
tests/entrypoints/pooling/score/test_online_score.py
+95
-37
@@ -61,14 +61,40 @@ def runner(model: dict[str, Any], hf_runner):
61
62
63
class TestModel:
64
- def test_text_1_str_text_2_list(
64
+ def test_queries_str_documents_str(
65
+ self, server: RemoteOpenAIServer, model: dict[str, Any], runner
66
+ ):
67
+ queries = "What is the capital of France?"
68
+ documents = "The capital of France is Paris."
69
+
70
+ score_response = requests.post(
71
+ server.url_for("score"),
72
+ json={
73
+ "model": model["name"],
74
+ "queries": queries,
75
+ "documents": documents,
76
+ },
77
+ )
78
+ score_response.raise_for_status()
79
+ score = ScoreResponse.model_validate(score_response.json())
80
+
81
+ assert score.id is not None
82
+ assert score.data is not None
83
+ assert len(score.data) == 1
84
+
85
+ vllm_outputs = [d.score for d in score.data]
86
+
87
+ text_pairs = [[queries, documents]]
88
+ hf_outputs = run_transformers(runner, model, text_pairs)
89
+
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_text_1_str_text_2_str(
94
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
95
):
96
text_1 = "What is the capital of France?"
68
- text_2 = [
69
- "The capital of Brazil is Brasilia.",
70
- "The capital of France is Paris.",
71
- ]
97
+ text_2 = "The capital of France is Paris."
98
99
score_response = requests.post(
100
server.url_for("score"),
@@ -83,24 +109,50 @@ class TestModel:
109
110
assert score.id is not None
111
assert score.data is not None
86
- assert len(score.data) == 2
112
+ assert len(score.data) == 1
113
114
vllm_outputs = [d.score for d in score.data]
115
90
- text_pairs = [[text_1, text_2[0]], [text_1, text_2[1]]]
116
+ text_pairs = [[text_1, text_2]]
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
96
- def test_text_1_list_text_2_list(
122
+ def test_data_1_str_data_2_str(
123
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
124
):
99
- text_1 = [
100
- "What is the capital of the United States?",
101
- "What is the capital of France?",
102
- ]
103
- text_2 = [
125
+ data_1 = "What is the capital of France?"
126
+ data_2 = "The capital of France is Paris."
127
+
128
+ score_response = requests.post(
129
+ server.url_for("score"),
130
+ json={
131
+ "model": model["name"],
132
+ "data_1": data_1,
133
+ "data_2": data_2,
134
+ },
135
+ )
136
+ score_response.raise_for_status()
137
+ score = ScoreResponse.model_validate(score_response.json())
138
+
139
+ assert score.id is not None
140
+ assert score.data is not None
141
+ assert len(score.data) == 1
142
+
143
+ vllm_outputs = [d.score for d in score.data]
144
+
145
+ text_pairs = [[data_1, data_2]]
146
+ hf_outputs = run_transformers(runner, model, text_pairs)
147
+
148
+ for i in range(len(vllm_outputs)):
149
+ assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
150
+
151
+ def test_queries_str_documents_list(
152
+ self, server: RemoteOpenAIServer, model: dict[str, Any], runner
153
+ ):
154
+ queries = "What is the capital of France?"
155
+ documents = [
156
"The capital of Brazil is Brasilia.",
157
"The capital of France is Paris.",
158
]
@@ -109,8 +161,8 @@ class TestModel:
161
server.url_for("score"),
162
json={
163
"model": model["name"],
112
- "text_1": text_1,
113
- "text_2": text_2,
164
+ "queries": queries,
165
+ "documents": documents,
166
},
167
)
168
score_response.raise_for_status()
@@ -122,24 +174,30 @@ class TestModel:
174
175
vllm_outputs = [d.score for d in score.data]
176
125
- text_pairs = [[text_1[0], text_2[0]], [text_1[1], text_2[1]]]
177
+ text_pairs = [[queries, documents[0]], [queries, documents[1]]]
178
hf_outputs = run_transformers(runner, model, text_pairs)
179
180
for i in range(len(vllm_outputs)):
181
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
182
131
- def test_text_1_str_text_2_str(
183
+ def test_queries_list_documents_list(
184
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
185
):
134
- text_1 = "What is the capital of France?"
135
- text_2 = "The capital of France is Paris."
186
+ queries = [
187
+ "What is the capital of the United States?",
188
+ "What is the capital of France?",
189
+ ]
190
+ documents = [
191
+ "The capital of Brazil is Brasilia.",
192
+ "The capital of France is Paris.",
193
+ ]
194
195
score_response = requests.post(
196
server.url_for("score"),
197
json={
198
"model": model["name"],
141
- "text_1": text_1,
142
- "text_2": text_2,
199
+ "queries": queries,
200
+ "documents": documents,
201
},
202
)
203
score_response.raise_for_status()
@@ -147,11 +205,11 @@ class TestModel:
205
206
assert score.id is not None
207
assert score.data is not None
150
- assert len(score.data) == 1
208
+ assert len(score.data) == 2
209
210
vllm_outputs = [d.score for d in score.data]
211
154
- text_pairs = [[text_1, text_2]]
212
+ text_pairs = [[queries[0], documents[0]], [queries[1], documents[1]]]
213
hf_outputs = run_transformers(runner, model, text_pairs)
214
215
for i in range(len(vllm_outputs)):
@@ -160,8 +218,8 @@ class TestModel:
218
def test_score_max_model_len(
219
self, server: RemoteOpenAIServer, model: dict[str, Any]
220
):
163
- text_1 = "What is the capital of France?" * 20
164
- text_2 = [
221
+ queries = "What is the capital of France?" * 20
222
+ documents = [
223
"The capital of Brazil is Brasilia.",
224
"The capital of France is Paris.",
225
]
@@ -170,8 +228,8 @@ class TestModel:
228
server.url_for("score"),
229
json={
230
"model": model["name"],
173
- "text_1": text_1,
174
- "text_2": text_2,
231
+ "queries": queries,
232
+ "documents": documents,
233
},
234
)
235
assert score_response.status_code == 400
@@ -183,8 +241,8 @@ class TestModel:
241
server.url_for("score"),
242
json={
243
"model": model["name"],
186
- "text_1": text_1,
187
- "text_2": text_2,
244
+ "queries": queries,
245
+ "documents": documents,
246
"truncate_prompt_tokens": 101,
247
},
248
)
@@ -192,13 +250,13 @@ class TestModel:
250
assert "Please, select a smaller truncation size." in score_response.text
251
252
def test_invocations(self, server: RemoteOpenAIServer, model: dict[str, Any]):
195
- text_1 = "What is the capital of France?"
196
- text_2 = "The capital of France is Paris."
253
+ queries = "What is the capital of France?"
254
+ documents = "The capital of France is Paris."
255
256
request_args = {
257
"model": model["name"],
200
- "text_1": text_1,
201
- "text_2": text_2,
258
+ "queries": queries,
259
+ "documents": documents,
260
}
261
262
score_response = requests.post(server.url_for("score"), json=request_args)
@@ -225,14 +283,14 @@ class TestModel:
283
284
def test_use_activation(self, server: RemoteOpenAIServer, model: dict[str, Any]):
285
def get_outputs(use_activation):
228
- text_1 = "What is the capital of France?"
229
- text_2 = "The capital of France is Paris."
286
+ queries = "What is the capital of France?"
287
+ documents = "The capital of France is Paris."
288
response = requests.post(
289
server.url_for("score"),
290
json={
291
"model": model["name"],
234
- "text_1": text_1,
235
- "text_2": text_2,
292
+ "queries": queries,
293
+ "documents": documents,
294
"use_activation": use_activation,
295
},
296
)
tests/models/language/pooling_mteb_test/mteb_score_utils.py
+2
-2
@@ -117,8 +117,8 @@ class ScoreClientMtebEncoder(MtebCrossEncoderMixin):
117
self.url,
118
json={
119
"model": self.model_name,
120
- "text_1": query,
121
- "text_2": corpus,
120
+ "queries": query,
121
+ "documents": corpus,
122
"truncate_prompt_tokens": -1,
123
},
124
).json()
vllm/entrypoints/openai/engine/serving.py
+9
-2
@@ -84,8 +84,11 @@ from vllm.entrypoints.pooling.pooling.protocol import (
84
)
85
from vllm.entrypoints.pooling.score.protocol import (
86
RerankRequest,
87
+ ScoreDataRequest,
88
+ ScoreQueriesDocumentsRequest,
89
ScoreRequest,
90
ScoreResponse,
91
+ ScoreTextRequest,
92
)
93
from vllm.entrypoints.renderer import BaseRenderer, CompletionRenderer, RenderConfig
94
from vllm.entrypoints.serve.disagg.protocol import GenerateRequest, GenerateResponse
@@ -1032,7 +1035,9 @@ class OpenAIServing:
1035
(
1036
EmbeddingChatRequest,
1037
EmbeddingCompletionRequest,
1035
- ScoreRequest,
1038
+ ScoreDataRequest,
1039
+ ScoreTextRequest,
1040
+ ScoreQueriesDocumentsRequest,
1041
RerankRequest,
1042
ClassificationCompletionRequest,
1043
ClassificationChatRequest,
@@ -1042,7 +1047,9 @@ class OpenAIServing:
1047
# since these requests don't generate tokens.
1048
if token_num > self.max_model_len:
1049
operations: dict[type[AnyRequest], str] = {
1045
- ScoreRequest: "score",
1050
+ ScoreDataRequest: "score",
1051
+ ScoreTextRequest: "score",
1052
+ ScoreQueriesDocumentsRequest: "score",
1053
ClassificationCompletionRequest: "classification",
1054
ClassificationChatRequest: "classification",
1055
}
vllm/entrypoints/openai/run_batch.py
+1
-1
@@ -85,7 +85,7 @@ class BatchRequestInput(OpenAIBaseModel):
85
if url == "/v1/embeddings":
86
return TypeAdapter(EmbeddingRequest).validate_python(value)
87
if url.endswith("/score"):
88
- return ScoreRequest.model_validate(value)
88
+ return TypeAdapter(ScoreRequest).validate_python(value)
89
if url.endswith("/rerank"):
90
return RerankRequest.model_validate(value)
91
return TypeAdapter(BatchRequestInputBody).validate_python(value)
vllm/entrypoints/pooling/score/protocol.py
+38
-5
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import time
4
-from typing import Any
4
+from typing import Any, TypeAlias
5
6
from pydantic import (
7
BaseModel,
@@ -19,10 +19,7 @@ from vllm.entrypoints.pooling.score.utils import (
19
from vllm.utils import random_uuid
20
21
22
-class ScoreRequest(PoolingBasicRequestMixin):
23
- text_1: list[str] | str | ScoreMultiModalParam
24
- text_2: list[str] | str | ScoreMultiModalParam
25
-
22
+class ScoreRequestMixin(PoolingBasicRequestMixin):
23
# --8<-- [start:score-extra-params]
24
mm_processor_kwargs: dict[str, Any] | None = Field(
25
default=None,
@@ -53,6 +50,42 @@ class ScoreRequest(PoolingBasicRequestMixin):
50
)
51
52
53
+class ScoreDataRequest(ScoreRequestMixin):
54
+ data_1: list[str] | str | ScoreMultiModalParam
55
+ data_2: list[str] | str | ScoreMultiModalParam
56
+
57
+
58
+class ScoreQueriesDocumentsRequest(ScoreRequestMixin):
59
+ queries: list[str] | str | ScoreMultiModalParam
60
+ documents: list[str] | str | ScoreMultiModalParam
61
+
62
+ @property
63
+ def data_1(self):
64
+ return self.queries
65
+
66
+ @property
67
+ def data_2(self):
68
+ return self.documents
69
+
70
+
71
+class ScoreTextRequest(ScoreRequestMixin):
72
+ text_1: list[str] | str | ScoreMultiModalParam
73
+ text_2: list[str] | str | ScoreMultiModalParam
74
+
75
+ @property
76
+ def data_1(self):
77
+ return self.text_1
78
+
79
+ @property
80
+ def data_2(self):
81
+ return self.text_2
82
+
83
+
84
+ScoreRequest: TypeAlias = (
85
+ ScoreQueriesDocumentsRequest | ScoreDataRequest | ScoreTextRequest
86
+)
87
+
88
+
89
class RerankRequest(PoolingBasicRequestMixin):
90
query: str | ScoreMultiModalParam
91
documents: list[str] | ScoreMultiModalParam
vllm/entrypoints/pooling/score/serving.py
+16
-16
@@ -66,15 +66,15 @@ class ServingScores(OpenAIServing):
66
async def _embedding_score(
67
self,
68
tokenizer: TokenizerLike,
69
- texts_1: list[str],
70
- texts_2: list[str],
69
+ data_1: list[str],
70
+ data_2: list[str],
71
request: RerankRequest | ScoreRequest,
72
request_id: str,
73
tokenization_kwargs: dict[str, Any] | None = None,
74
lora_request: LoRARequest | None | None = None,
75
trace_headers: Mapping[str, str] | None = None,
76
) -> list[PoolingRequestOutput] | ErrorResponse:
77
- input_texts = texts_1 + texts_2
77
+ input_texts = data_1 + data_2
78
79
engine_prompts: list[TokensPrompt] = []
80
tokenize_async = make_async(
@@ -135,22 +135,22 @@ class ServingScores(OpenAIServing):
135
async for i, res in result_generator:
136
embeddings[i] = res
137
138
- emb_texts_1: list[PoolingRequestOutput] = []
139
- emb_texts_2: list[PoolingRequestOutput] = []
138
+ emb_data_1: list[PoolingRequestOutput] = []
139
+ emb_data_2: list[PoolingRequestOutput] = []
140
141
- for i in range(0, len(texts_1)):
141
+ for i in range(0, len(data_1)):
142
assert (emb := embeddings[i]) is not None
143
- emb_texts_1.append(emb)
143
+ emb_data_1.append(emb)
144
145
- for i in range(len(texts_1), len(embeddings)):
145
+ for i in range(len(data_1), len(embeddings)):
146
assert (emb := embeddings[i]) is not None
147
- emb_texts_2.append(emb)
147
+ emb_data_2.append(emb)
148
149
- if len(emb_texts_1) == 1:
150
- emb_texts_1 = emb_texts_1 * len(emb_texts_2)
149
+ if len(emb_data_1) == 1:
150
+ emb_data_1 = emb_data_1 * len(emb_data_2)
151
152
final_res_batch = _cosine_similarity(
153
- tokenizer=tokenizer, embed_1=emb_texts_1, embed_2=emb_texts_2
153
+ tokenizer=tokenizer, embed_1=emb_data_1, embed_2=emb_data_2
154
)
155
156
return final_res_batch
@@ -333,8 +333,8 @@ class ServingScores(OpenAIServing):
333
else:
334
return await self._embedding_score(
335
tokenizer=tokenizer,
336
- texts_1=data_1, # type: ignore[arg-type]
337
- texts_2=data_2, # type: ignore[arg-type]
336
+ data_1=data_1, # type: ignore[arg-type]
337
+ data_2=data_2, # type: ignore[arg-type]
338
request=request,
339
request_id=request_id,
340
tokenization_kwargs=tokenization_kwargs,
@@ -361,8 +361,8 @@ class ServingScores(OpenAIServing):
361
362
try:
363
final_res_batch = await self._run_scoring(
364
- request.text_1,
365
- request.text_2,
364
+ request.data_1,
365
+ request.data_2,
366
request,
367
request_id,
368
raw_request,