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 5, 18:30 UTC
[Frontend][3/n] Improve pooling entrypoints | scoring. (#28631)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
wang.yuqi committed
Mar 31, 2026 at 15:52 UTC
d9d21eb8e35fa356ebb2a1ffa94da4306905ff0e
37 files changed
+1257
-1780
tests/entrypoints/openai/utils.py
+1
-3
@@ -10,9 +10,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
10
ChatCompletionStreamResponse,
11
ChatMessage,
12
)
13
-from vllm.entrypoints.openai.engine.protocol import (
14
- UsageInfo,
15
-)
13
+from vllm.entrypoints.openai.engine.protocol import UsageInfo
14
15
16
async def accumulate_streaming_response(
tests/entrypoints/pooling/classify/test_offline.py
+1
-1
@@ -105,7 +105,7 @@ def test_pooling_params(llm: LLM):
105
106
@pytest.mark.skip_global_cleanup
107
def test_score_api(llm: LLM):
108
- err_msg = "Score API is only enabled for num_labels == 1."
108
+ err_msg = "Scoring API is only enabled for num_labels == 1."
109
with pytest.raises(ValueError, match=err_msg):
110
llm.score("ping", "pong", use_tqdm=False)
111
tests/entrypoints/pooling/classify/test_online.py
+2
-2
@@ -390,7 +390,7 @@ async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
390
@pytest.mark.asyncio
391
@pytest.mark.parametrize("model_name", [MODEL_NAME])
392
async def test_score(server: RemoteOpenAIServer, model_name: str):
393
- # score api is only enabled for num_labels == 1.
393
+ # Scoring API is only enabled for num_labels == 1.
394
response = requests.post(
395
server.url_for("score"),
396
json={
@@ -405,7 +405,7 @@ async def test_score(server: RemoteOpenAIServer, model_name: str):
405
@pytest.mark.asyncio
406
@pytest.mark.parametrize("model_name", [MODEL_NAME])
407
async def test_rerank(server: RemoteOpenAIServer, model_name: str):
408
- # rerank api is only enabled for num_labels == 1.
408
+ # Scoring API is only enabled for num_labels == 1.
409
response = requests.post(
410
server.url_for("rerank"),
411
json={
tests/entrypoints/pooling/scoring/test_bi_encoder_online.py
+1
-1
@@ -7,7 +7,7 @@ import requests
7
from tests.entrypoints.pooling.scoring.util import EncoderScoringHfRunner
8
from tests.utils import RemoteOpenAIServer
9
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
10
-from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse
10
+from vllm.entrypoints.pooling.scoring.protocol import RerankResponse, ScoreResponse
11
from vllm.platforms import current_platform
12
13
MODEL_NAME = "BAAI/bge-base-en-v1.5"
tests/entrypoints/pooling/scoring/test_cross_encoder_online.py
+1
-1
@@ -8,7 +8,7 @@ import torch.nn.functional as F
8
9
from tests.utils import RemoteOpenAIServer
10
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
11
-from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse
11
+from vllm.entrypoints.pooling.scoring.protocol import RerankResponse, ScoreResponse
12
from vllm.platforms import current_platform
13
14
MODEL_NAME = "BAAI/bge-reranker-base"
tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py
+1
-1
@@ -7,7 +7,7 @@ import pytest
7
import requests
8
9
from tests.utils import VLLM_PATH, RemoteOpenAIServer
10
-from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse
10
+from vllm.entrypoints.pooling.scoring.protocol import RerankResponse, ScoreResponse
11
from vllm.multimodal.utils import encode_image_url, fetch_image
12
from vllm.platforms import current_platform
13
tests/entrypoints/pooling/scoring/test_late_interaction_offline_vision.py
new
+93
@@ -0,0 +1,93 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import weakref
5
+
6
+import pytest
7
+
8
+from vllm import LLM
9
+from vllm.distributed import cleanup_dist_env_and_memory
10
+from vllm.platforms import current_platform
11
+
12
+from .util import make_base64_image, make_image_mm_param
13
+
14
+MODEL_NAME = "vidore/colpali-v1.3-hf"
15
+
16
+
17
+@pytest.fixture(scope="module")
18
+def llm():
19
+ # ROCm: Use FLEX_ATTENTION backend as it's the only attention backend
20
+ # that supports encoder-only models on ROCm.
21
+ attention_config = None
22
+ if current_platform.is_rocm():
23
+ attention_config = {"backend": "FLEX_ATTENTION"}
24
+
25
+ # pytest caches the fixture so we use weakref.proxy to
26
+ # enable garbage collection
27
+ llm = LLM(
28
+ model=MODEL_NAME,
29
+ max_num_batched_tokens=32768,
30
+ tensor_parallel_size=1,
31
+ gpu_memory_utilization=0.75,
32
+ enforce_eager=True,
33
+ seed=0,
34
+ attention_config=attention_config,
35
+ )
36
+
37
+ yield weakref.proxy(llm)
38
+
39
+ del llm
40
+
41
+ cleanup_dist_env_and_memory()
42
+
43
+
44
+@pytest.mark.skip_global_cleanup
45
+def test_query_text_vs_docs_image(llm):
46
+ """Score a text query against image documents via the multimodal path."""
47
+ red_image = make_base64_image(64, 64, color=(255, 0, 0))
48
+ blue_image = make_base64_image(64, 64, color=(0, 0, 255))
49
+
50
+ query = "Describe the red object"
51
+ image_docs = [
52
+ make_image_mm_param(red_image),
53
+ make_image_mm_param(blue_image),
54
+ ]
55
+
56
+ scores = llm.score(query, image_docs)
57
+
58
+ assert len(scores) == 2
59
+ assert scores[0].outputs.score > scores[1].outputs.score
60
+
61
+
62
+@pytest.mark.skip_global_cleanup
63
+def test_query_text_vs_docs_mix(llm) -> None:
64
+ """Score a text query against a mix of text and image documents."""
65
+ red_image = make_base64_image(64, 64, color=(255, 0, 0))
66
+
67
+ query = "What is the capital of France?"
68
+ documents: list = [
69
+ "The capital of France is Paris.",
70
+ make_image_mm_param(red_image),
71
+ ]
72
+
73
+ scores = llm.score(query, documents)
74
+
75
+ assert len(scores) == 2
76
+ assert scores[0].outputs.score > scores[1].outputs.score
77
+
78
+
79
+@pytest.mark.skip_global_cleanup
80
+def test_query_image_vs_docs_text(llm) -> None:
81
+ """Score an image query against text documents."""
82
+ red_image = make_base64_image(64, 64, color=(255, 0, 0))
83
+ image_query = make_image_mm_param(red_image, text="red color")
84
+
85
+ documents = [
86
+ "Describe the red object.",
87
+ "The capital of France is Paris.",
88
+ ]
89
+
90
+ scores = llm.score(image_query, documents)
91
+
92
+ assert len(scores) == 2
93
+ assert scores[0].outputs.score > scores[1].outputs.score
tests/entrypoints/pooling/scoring/test_late_interaction_online.py
+1
-1
@@ -6,7 +6,7 @@ import pytest
6
import requests
7
8
from tests.utils import RemoteOpenAIServer
9
-from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse
9
+from vllm.entrypoints.pooling.scoring.protocol import RerankResponse, ScoreResponse
10
11
from .util import ColBERTScoringHfRunner
12
tests/entrypoints/pooling/scoring/test_late_interaction_online_vision.py
new
+193
@@ -0,0 +1,193 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import requests
6
+
7
+from tests.entrypoints.pooling.scoring.util import (
8
+ make_base64_image,
9
+ make_image_mm_param,
10
+)
11
+from tests.utils import RemoteOpenAIServer
12
+from vllm.entrypoints.pooling.scoring.protocol import RerankResponse, ScoreResponse
13
+
14
+MODEL_NAME = "vidore/colpali-v1.3-hf"
15
+
16
+
17
+@pytest.fixture(scope="module")
18
+def server():
19
+ with RemoteOpenAIServer(MODEL_NAME, []) as remote_server:
20
+ yield remote_server
21
+
22
+
23
+@pytest.mark.asyncio
24
+async def test_score_api_query_text_vs_docs_image(server: RemoteOpenAIServer):
25
+ query = "Describe the red object"
26
+
27
+ red_image = make_base64_image(64, 64, color=(255, 0, 0))
28
+ blue_image = make_base64_image(64, 64, color=(0, 0, 255))
29
+
30
+ documents = [
31
+ make_image_mm_param(red_image),
32
+ make_image_mm_param(blue_image),
33
+ ]
34
+
35
+ score_response = requests.post(
36
+ server.url_for("score"),
37
+ json={
38
+ "model": MODEL_NAME,
39
+ "queries": query,
40
+ "documents": documents,
41
+ },
42
+ )
43
+ score_response.raise_for_status()
44
+ scores = ScoreResponse.model_validate(score_response.json())
45
+
46
+ assert scores.id is not None
47
+ assert scores.data is not None
48
+ assert len(scores.data) == 2
49
+ assert scores.data[0].score > scores.data[1].score
50
+
51
+
52
+@pytest.mark.asyncio
53
+async def test_score_api_query_text_vs_docs_mix(server: RemoteOpenAIServer):
54
+ red_image = make_base64_image(64, 64, color=(255, 0, 0))
55
+ query = "What is the capital of France?"
56
+ documents: list = [
57
+ "The capital of France is Paris.",
58
+ make_image_mm_param(red_image),
59
+ ]
60
+
61
+ score_response = requests.post(
62
+ server.url_for("score"),
63
+ json={
64
+ "model": MODEL_NAME,
65
+ "queries": query,
66
+ "documents": documents,
67
+ },
68
+ )
69
+ score_response.raise_for_status()
70
+ scores = ScoreResponse.model_validate(score_response.json())
71
+
72
+ assert scores.id is not None
73
+ assert scores.data is not None
74
+ assert len(scores.data) == 2
75
+ assert scores.data[0].score > scores.data[1].score
76
+
77
+
78
+@pytest.mark.asyncio
79
+async def test_score_api_query_image_vs_docs_text(server: RemoteOpenAIServer):
80
+ red_image = make_base64_image(64, 64, color=(255, 0, 0))
81
+ image_query = make_image_mm_param(red_image, text="red color")
82
+
83
+ documents = [
84
+ "Describe the red object.",
85
+ "The capital of France is Paris.",
86
+ ]
87
+
88
+ score_response = requests.post(
89
+ server.url_for("score"),
90
+ json={
91
+ "model": MODEL_NAME,
92
+ "queries": image_query,
93
+ "documents": documents,
94
+ },
95
+ )
96
+ score_response.raise_for_status()
97
+ scores = ScoreResponse.model_validate(score_response.json())
98
+
99
+ assert scores.id is not None
100
+ assert scores.data is not None
101
+ assert len(scores.data) == 2
102
+ assert scores.data[0].score > scores.data[1].score
103
+
104
+
105
+@pytest.mark.asyncio
106
+async def test_rerank_api_query_text_vs_docs_image(server: RemoteOpenAIServer):
107
+ query = "Describe the red object"
108
+
109
+ red_image = make_base64_image(64, 64, color=(255, 0, 0))
110
+ blue_image = make_base64_image(64, 64, color=(0, 0, 255))
111
+
112
+ documents = [
113
+ make_image_mm_param(red_image),
114
+ make_image_mm_param(blue_image),
115
+ ]
116
+
117
+ rerank_response = requests.post(
118
+ server.url_for("rerank"),
119
+ json={"model": MODEL_NAME, "query": query, "documents": documents},
120
+ )
121
+
122
+ rerank_response.raise_for_status()
123
+ rerank = RerankResponse.model_validate(rerank_response.json())
124
+
125
+ assert rerank.id is not None
126
+ assert rerank.results is not None
127
+ assert len(rerank.results) == 2
128
+
129
+ red_result = next(r for r in rerank.results if r.index == 0)
130
+ blue_result = next(r for r in rerank.results if r.index == 1)
131
+
132
+ assert red_result.relevance_score > blue_result.relevance_score
133
+
134
+
135
+@pytest.mark.asyncio
136
+async def test_rerank_api_query_text_vs_docs_mix(server: RemoteOpenAIServer):
137
+ red_image = make_base64_image(64, 64, color=(255, 0, 0))
138
+ query = "What is the capital of France?"
139
+ documents: list = [
140
+ "The capital of France is Paris.",
141
+ make_image_mm_param(red_image),
142
+ ]
143
+
144
+ rerank_response = requests.post(
145
+ server.url_for("rerank"),
146
+ json={
147
+ "model": MODEL_NAME,
148
+ "query": query,
149
+ "documents": documents,
150
+ },
151
+ )
152
+ rerank_response.raise_for_status()
153
+ rerank = RerankResponse.model_validate(rerank_response.json())
154
+
155
+ assert rerank.id is not None
156
+ assert rerank.results is not None
157
+ assert len(rerank.results) == 2
158
+
159
+ result0 = next(r for r in rerank.results if r.index == 0)
160
+ result1 = next(r for r in rerank.results if r.index == 1)
161
+
162
+ assert result0.relevance_score > result1.relevance_score
163
+
164
+
165
+@pytest.mark.asyncio
166
+async def test_rerank_api_query_image_vs_docs_text(server: RemoteOpenAIServer):
167
+ red_image = make_base64_image(64, 64, color=(255, 0, 0))
168
+ image_query = make_image_mm_param(red_image, text="red color")
169
+
170
+ documents = [
171
+ "Describe the red object.",
172
+ "The capital of France is Paris.",
173
+ ]
174
+
175
+ rerank_response = requests.post(
176
+ server.url_for("rerank"),
177
+ json={
178
+ "model": MODEL_NAME,
179
+ "query": image_query,
180
+ "documents": documents,
181
+ },
182
+ )
183
+ rerank_response.raise_for_status()
184
+ rerank = RerankResponse.model_validate(rerank_response.json())
185
+
186
+ assert rerank.id is not None
187
+ assert rerank.results is not None
188
+ assert len(rerank.results) == 2
189
+
190
+ result0 = next(r for r in rerank.results if r.index == 0)
191
+ result1 = next(r for r in rerank.results if r.index == 1)
192
+
193
+ assert result0.relevance_score > result1.relevance_score
tests/entrypoints/pooling/scoring/test_utils.py
deleted
-353
@@ -1,353 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from unittest.mock import patch
5
-
6
-import pytest
7
-
8
-from vllm.config import ModelConfig
9
-from vllm.entrypoints.chat_utils import ChatTemplateResolutionError
10
-from vllm.entrypoints.pooling.score.utils import (
11
- get_score_prompt,
12
-)
13
-from vllm.inputs import TokensPrompt
14
-from vllm.tokenizers import get_tokenizer
15
-
16
-# A cross-encoder model for testing
17
-CROSS_ENCODER_MODEL_ID = "cross-encoder/ms-marco-MiniLM-L-6-v2"
18
-
19
-
20
-def assert_prompt_tokenization_consistent(
21
- tokenizer, full_prompt, engine_prompt, add_special_tokens=True
22
-):
23
- """Verify that engine_prompt token_ids match tokenizing full_prompt."""
24
- expected_ids = tokenizer(full_prompt, add_special_tokens=add_special_tokens)[
25
- "input_ids"
26
- ]
27
- actual_ids = engine_prompt["prompt_token_ids"]
28
- assert actual_ids == expected_ids, (
29
- f"Token IDs don't match.\nExpected: {expected_ids}\nActual: {actual_ids}"
30
- )
31
-
32
-
33
-@pytest.fixture(scope="module")
34
-def cross_encoder_model_config():
35
- return ModelConfig(
36
- CROSS_ENCODER_MODEL_ID,
37
- runner="pooling",
38
- )
39
-
40
-
41
-@pytest.fixture(scope="module")
42
-def cross_encoder_tokenizer(cross_encoder_model_config):
43
- return get_tokenizer(
44
- CROSS_ENCODER_MODEL_ID,
45
- trust_remote_code=cross_encoder_model_config.trust_remote_code,
46
- )
47
-
48
-
49
-@pytest.fixture(scope="module")
50
-def llm_reranker_model_config():
51
- """Model config for LLM-as-reranker style (no pad token)."""
52
- config = ModelConfig(
53
- CROSS_ENCODER_MODEL_ID,
54
- runner="pooling",
55
- )
56
- # use_sep_token is a property that reads from hf_config,
57
- # so we set it there to override the default (True)
58
- config.hf_config.use_sep_token = False
59
- return config
60
-
61
-
62
-@pytest.fixture
63
-def tokenization_kwargs():
64
- """Common tokenization kwargs used across tests."""
65
- return {"add_special_tokens": True, "return_tensors": None}
66
-
67
-
68
-@pytest.fixture
69
-def mock_model_with_score_template():
70
- """Mock model class that supports score template and tracks post_process calls."""
71
-
72
- class MockModelWithScoreTemplate:
73
- supports_score_template = True
74
- post_process_called: list[TokensPrompt] = []
75
-
76
- @staticmethod
77
- def get_score_template(p1: str, p2: str) -> str:
78
- return f"[QUERY]{p1}[SEP][DOC]{p2}"
79
-
80
- @staticmethod
81
- def post_process_tokens(prompt: TokensPrompt) -> None:
82
- MockModelWithScoreTemplate.post_process_called.append(prompt)
83
-
84
- return MockModelWithScoreTemplate
85
-
86
-
87
-@pytest.fixture
88
-def mock_model_no_score_template():
89
- """Mock model class that does not support score template."""
90
-
91
- class MockModelNoScoreTemplate:
92
- supports_score_template = False
93
-
94
- return MockModelNoScoreTemplate
95
-
96
-
97
-class TestGetScorePrompt:
98
- """Tests for the get_score_prompt function."""
99
-
100
- def test_tokenization_kwargs_passed_through(
101
- self,
102
- llm_reranker_model_config,
103
- cross_encoder_tokenizer,
104
- ):
105
- """Test that tokenization kwargs are properly passed through."""
106
- data_1 = "Query text"
107
- data_2 = "Document text"
108
-
109
- # Test with truncation - custom kwargs for this test
110
- custom_tokenization_kwargs = {
111
- "add_special_tokens": True,
112
- "return_tensors": None,
113
- "truncation": True,
114
- "max_length": 20,
115
- }
116
-
117
- full_prompt, engine_prompt = get_score_prompt(
118
- llm_reranker_model_config,
119
- cross_encoder_tokenizer,
120
- custom_tokenization_kwargs,
121
- data_1,
122
- data_2,
123
- )
124
-
125
- assert isinstance(full_prompt, str)
126
- assert "prompt_token_ids" in engine_prompt
127
- # With max_length=20 and truncation, should not exceed this
128
- assert len(engine_prompt["prompt_token_ids"]) <= 20
129
- # Since truncation was applied, token_ids should be a prefix of full encoding
130
- full_ids = cross_encoder_tokenizer(full_prompt, add_special_tokens=True)[
131
- "input_ids"
132
- ]
133
- actual_ids = engine_prompt["prompt_token_ids"]
134
- assert full_ids[: len(actual_ids)] == actual_ids, (
135
- f"Token IDs are not a prefix of full encoding.\n"
136
- f"Full IDs: {full_ids}\n"
137
- f"Actual IDs: {actual_ids}"
138
- )
139
-
140
- def test_model_supports_score_template(
141
- self,
142
- cross_encoder_model_config,
143
- cross_encoder_tokenizer,
144
- tokenization_kwargs,
145
- mock_model_with_score_template,
146
- ):
147
- """Test when model supports score template (no score_template arg)."""
148
- with patch(
149
- "vllm.model_executor.model_loader.get_model_cls",
150
- return_value=mock_model_with_score_template,
151
- ):
152
- full_prompt, engine_prompt = get_score_prompt(
153
- cross_encoder_model_config,
154
- cross_encoder_tokenizer,
155
- tokenization_kwargs,
156
- "query text",
157
- "document text",
158
- )
159
-
160
- assert full_prompt == "[QUERY]query text[SEP][DOC]document text"
161
- assert "prompt_token_ids" in engine_prompt
162
- assert len(engine_prompt["prompt_token_ids"]) > 0
163
- assert_prompt_tokenization_consistent(
164
- cross_encoder_tokenizer, full_prompt, engine_prompt
165
- )
166
-
167
- def test_model_supports_score_template_but_custom_template_provided(
168
- self,
169
- cross_encoder_model_config,
170
- cross_encoder_tokenizer,
171
- tokenization_kwargs,
172
- mock_model_with_score_template,
173
- ):
174
- """Test when model supports score template but custom template is provided."""
175
- template = (
176
- 'TEMPLATE_USED {{ messages[0]["content"] }} {{ messages[1]["content"] }}'
177
- )
178
- with (
179
- patch(
180
- "vllm.model_executor.model_loader.get_model_cls",
181
- return_value=mock_model_with_score_template,
182
- ),
183
- ):
184
- full_prompt, engine_prompt = get_score_prompt(
185
- cross_encoder_model_config,
186
- cross_encoder_tokenizer,
187
- tokenization_kwargs,
188
- "query",
189
- "doc",
190
- score_template=template, # Providing a template
191
- )
192
-
193
- assert "prompt_token_ids" in engine_prompt
194
- assert full_prompt == "TEMPLATE_USED query doc"
195
-
196
- assert_prompt_tokenization_consistent(
197
- cross_encoder_tokenizer, full_prompt, engine_prompt
198
- )
199
-
200
- def test_not_using_default_template(
201
- self,
202
- llm_reranker_model_config,
203
- cross_encoder_tokenizer,
204
- tokenization_kwargs,
205
- mock_model_no_score_template,
206
- ):
207
- # FIXME: For now, we only apply a template when one is explicitly provided.
208
- # We cannot rely on the tokenizer's chat template because many models
209
- # inherit junk templates from their base LLM, which breaks both the models
210
- # and the tests that use them.
211
- with (
212
- patch(
213
- "vllm.model_executor.model_loader.get_model_cls",
214
- return_value=mock_model_no_score_template,
215
- ),
216
- patch(
217
- "vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
218
- return_value="test querytest doc",
219
- ),
220
- ):
221
- full_prompt, engine_prompt = get_score_prompt(
222
- llm_reranker_model_config,
223
- cross_encoder_tokenizer,
224
- tokenization_kwargs,
225
- "test query",
226
- "test doc",
227
- )
228
-
229
- assert full_prompt == "test querytest doc"
230
- assert "prompt_token_ids" in engine_prompt
231
- assert_prompt_tokenization_consistent(
232
- cross_encoder_tokenizer, full_prompt, engine_prompt
233
- )
234
-
235
- def test_fallback_with_sep_token(
236
- self,
237
- cross_encoder_model_config,
238
- cross_encoder_tokenizer,
239
- tokenization_kwargs,
240
- mock_model_no_score_template,
241
- ):
242
- """Test fallback path when ChatTemplateResolutionError
243
- and use_sep_token=True."""
244
- with (
245
- patch(
246
- "vllm.model_executor.model_loader.get_model_cls",
247
- return_value=mock_model_no_score_template,
248
- ),
249
- patch(
250
- "vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
251
- side_effect=ChatTemplateResolutionError("No template"),
252
- ),
253
- ):
254
- full_prompt, engine_prompt = get_score_prompt(
255
- cross_encoder_model_config, # use_sep_token=True
256
- cross_encoder_tokenizer,
257
- tokenization_kwargs,
258
- "query",
259
- "document",
260
- )
261
-
262
- assert "prompt_token_ids" in engine_prompt
263
- # Should have token_type_ids from text_pair encoding
264
- assert "token_type_ids" in engine_prompt
265
- assert "query" in full_prompt
266
- assert "document" in full_prompt
267
- assert full_prompt != "querydocument"
268
- assert (
269
- engine_prompt["prompt_token_ids"]
270
- == cross_encoder_tokenizer(
271
- "query", text_pair="document", add_special_tokens=True
272
- )["input_ids"]
273
- )
274
-
275
- # FIXME(?): add_special_tokens=False is needed because in this case
276
- # full_prompt is obtained by decoding the tokenized prompt, which includes
277
- # special tokens and we would get duplicated special tokens otherwise.
278
- # This is inconsistent with other cases.
279
- assert_prompt_tokenization_consistent(
280
- cross_encoder_tokenizer,
281
- full_prompt,
282
- engine_prompt,
283
- add_special_tokens=False,
284
- )
285
-
286
- def test_fallback_without_sep_token(
287
- self,
288
- llm_reranker_model_config,
289
- cross_encoder_tokenizer,
290
- tokenization_kwargs,
291
- mock_model_no_score_template,
292
- ):
293
- """Test fallback path when ChatTemplateResolutionError
294
- and use_sep_token=False."""
295
- with (
296
- patch(
297
- "vllm.model_executor.model_loader.get_model_cls",
298
- return_value=mock_model_no_score_template,
299
- ),
300
- patch(
301
- "vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
302
- side_effect=ChatTemplateResolutionError("No template"),
303
- ),
304
- ):
305
- full_prompt, engine_prompt = get_score_prompt(
306
- llm_reranker_model_config, # use_sep_token=False
307
- cross_encoder_tokenizer,
308
- tokenization_kwargs,
309
- "query",
310
- "document",
311
- )
312
-
313
- assert full_prompt == "querydocument"
314
- assert "prompt_token_ids" in engine_prompt
315
- assert_prompt_tokenization_consistent(
316
- cross_encoder_tokenizer, full_prompt, engine_prompt
317
- )
318
-
319
- def test_post_process_tokens_called(
320
- self,
321
- cross_encoder_model_config,
322
- cross_encoder_tokenizer,
323
- tokenization_kwargs,
324
- mock_model_with_score_template,
325
- ):
326
- """Test that post_process_tokens is called on the engine prompt."""
327
- # Reset the call tracker
328
- mock_model_with_score_template.post_process_called.clear()
329
-
330
- with (
331
- patch(
332
- "vllm.model_executor.model_loader.get_model_cls",
333
- return_value=mock_model_with_score_template,
334
- ),
335
- patch(
336
- "vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
337
- side_effect=ChatTemplateResolutionError("No template"),
338
- ),
339
- ):
340
- full_prompt, engine_prompt = get_score_prompt(
341
- cross_encoder_model_config,
342
- cross_encoder_tokenizer,
343
- tokenization_kwargs,
344
- "query",
345
- "doc",
346
- )
347
-
348
- # post_process_tokens should have been called once
349
- assert len(mock_model_with_score_template.post_process_called) == 1
350
- assert mock_model_with_score_template.post_process_called[0] is engine_prompt
351
- assert_prompt_tokenization_consistent(
352
- cross_encoder_tokenizer, full_prompt, engine_prompt
353
- )
tests/entrypoints/pooling/scoring/util.py
+39
-1
@@ -1,14 +1,23 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from io import BytesIO
5
+
6
+import pybase64 as base64
7
import torch
8
import torch.nn.functional as F
9
from huggingface_hub import hf_hub_download
10
+from PIL import Image
11
from safetensors.torch import load_file
12
from transformers import AutoModel, AutoTokenizer
13
14
from tests.conftest import HfRunner
11
-from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
15
+from vllm.entrypoints.chat_utils import (
16
+ ChatCompletionContentPartImageParam,
17
+ ChatCompletionContentPartTextParam,
18
+)
19
+from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam
20
+from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
21
22
23
class ColBERTScoringHfRunner(torch.nn.Module):
@@ -67,3 +76,32 @@ class EncoderScoringHfRunner(HfRunner):
76
for pair in hf_embeddings
77
]
78
return torch.as_tensor(hf_outputs)
79
+
80
+
81
+def make_base64_image(
82
+ width: int = 64, height: int = 64, color: tuple[int, int, int] = (255, 0, 0)
83
+) -> str:
84
+ """Create a small solid-color PNG image and return its base64 data URI."""
85
+ img = Image.new("RGB", (width, height), color)
86
+ buf = BytesIO()
87
+ img.save(buf, format="PNG")
88
+ b64 = base64.b64encode(buf.getvalue()).decode()
89
+ return f"data:image/png;base64,{b64}"
90
+
91
+
92
+def make_image_mm_param(
93
+ image_uri: str,
94
+ text: str | None = None,
95
+) -> ScoreMultiModalParam:
96
+ """Build a ScoreMultiModalParam containing an image (and optional text)."""
97
+ content: list = [
98
+ ChatCompletionContentPartImageParam(
99
+ type="image_url",
100
+ image_url={"url": image_uri},
101
+ ),
102
+ ]
103
+ if text is not None:
104
+ content.append(
105
+ ChatCompletionContentPartTextParam(type="text", text=text),
106
+ )
107
+ return ScoreMultiModalParam(content=content)
tests/entrypoints/pooling/token_classify/test_offline.py
+1
-1
@@ -60,7 +60,7 @@ def test_token_ids_prompts(llm: LLM):
60
61
@pytest.mark.skip_global_cleanup
62
def test_score_api(llm: LLM):
63
- err_msg = "Score API is only enabled for num_labels == 1."
63
+ err_msg = "Scoring API is only enabled for num_labels == 1."
64
with pytest.raises(ValueError, match=err_msg):
65
llm.score("ping", "pong", use_tqdm=False)
66
tests/models/language/pooling/test_colbert.py
+1
-1
@@ -9,7 +9,7 @@ generic ColBERT support works with different encoder architectures.
9
import pytest
10
import torch
11
12
-from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
12
+from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
13
14
# -----------------------------------------------------------------------
15
# Model definitions: (model_name, colbert_dim, extra vllm_runner kwargs)
tests/models/multimodal/pooling/test_colmodernvbert.py
+1
-1
@@ -10,7 +10,7 @@ embeddings for visual document retrieval.
10
import pytest
11
import torch
12
13
-from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
13
+from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
14
15
MODEL_NAME = "ModernVBERT/colmodernvbert-merged"
16
COLBERT_DIM = 128
tests/models/multimodal/pooling/test_colpali.py
+2
-2
@@ -18,7 +18,7 @@ from vllm.entrypoints.chat_utils import (
18
ChatCompletionContentPartImageParam,
19
ChatCompletionContentPartTextParam,
20
)
21
-from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
21
+from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam
22
23
from ....conftest import VllmRunner
24
@@ -114,7 +114,7 @@ def _run_late_interaction_test(
114
dtype: str,
115
) -> None:
116
"""Verify MaxSim scoring matches manual computation."""
117
- from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
117
+ from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
118
119
with vllm_runner(
120
model,
tests/models/multimodal/pooling/test_colqwen3.py
+2
-2
@@ -18,7 +18,7 @@ from vllm.entrypoints.chat_utils import (
18
ChatCompletionContentPartImageParam,
19
ChatCompletionContentPartTextParam,
20
)
21
-from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
21
+from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam
22
23
from ....conftest import VllmRunner
24
@@ -125,7 +125,7 @@ def _run_late_interaction_test(
125
dtype: str,
126
) -> None:
127
"""Verify MaxSim scoring matches manual computation."""
128
- from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
128
+ from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
129
130
with vllm_runner(
131
model,
tests/models/multimodal/pooling/test_colqwen3_5.py
+1
-1
@@ -73,7 +73,7 @@ def _run_late_interaction_test(
73
dtype: str,
74
) -> None:
75
"""Verify MaxSim scoring matches manual computation."""
76
- from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
76
+ from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
77
78
with vllm_runner(
79
model,
tests/models/multimodal/pooling/test_jinavl_reranker.py
+1
-1
@@ -11,7 +11,7 @@ from vllm.entrypoints.chat_utils import (
11
ChatCompletionContentPartImageParam,
12
ChatCompletionContentPartTextParam,
13
)
14
-from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
14
+from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam
15
16
from ....conftest import HfRunner, VllmRunner
17
tests/models/multimodal/pooling/test_llama_nemotron_vl.py
+1
-1
@@ -21,7 +21,7 @@ from vllm.entrypoints.chat_utils import (
21
ChatCompletionContentPartImageParam,
22
ChatCompletionContentPartTextParam,
23
)
24
-from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
24
+from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam
25
from vllm.platforms import current_platform
26
27
from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
vllm/entrypoints/llm.py
+60
-252
@@ -46,22 +46,16 @@ from vllm.entrypoints.chat_utils import (
46
load_chat_template,
47
)
48
from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors
49
-from vllm.entrypoints.pooling.score.utils import (
50
- ScoreData,
51
- ScoreMultiModalParam,
52
- _cosine_similarity,
53
- compress_token_type_ids,
54
- compute_maxsim_score,
55
- get_score_prompt,
56
- score_data_to_prompts,
57
- validate_score_input,
49
+from vllm.entrypoints.pooling.scoring.io_processor import (
50
+ ScoringIOProcessor,
51
)
52
+from vllm.entrypoints.pooling.scoring.typing import ScoreInput
53
+from vllm.entrypoints.pooling.typing import OfflineInputsContext, OfflineOutputsContext
54
from vllm.entrypoints.utils import log_non_default_args
55
from vllm.inputs import (
56
DataPrompt,
57
EngineInput,
58
PromptType,
64
- SingletonPrompt,
59
TextPrompt,
60
TokensPrompt,
61
)
@@ -1161,7 +1155,9 @@ class LLM:
1155
if pooling_task in self.pooling_io_processors:
1156
io_processor = self.pooling_io_processors[pooling_task]
1157
processor_inputs = io_processor.pre_process_offline(
1164
- prompts_seq, tokenization_kwargs
1158
+ ctx=OfflineInputsContext(
1159
+ prompts=prompts_seq, tokenization_kwargs=tokenization_kwargs
1160
+ )
1161
)
1162
seq_lora_requests = self._lora_request_to_seq(
1163
lora_request, len(prompts_seq)
@@ -1178,7 +1174,9 @@ class LLM:
1174
outputs = self._run_engine(
1175
use_tqdm=use_tqdm, output_type=PoolingRequestOutput
1176
)
1181
- outputs = io_processor.post_process_offline(outputs)
1177
+ outputs = io_processor.post_process_offline(
1178
+ ctx=OfflineOutputsContext(outputs=outputs)
1179
+ )
1180
else:
1181
outputs = self._run_completion(
1182
prompts=prompts_seq,
@@ -1378,188 +1376,10 @@ class LLM:
1376
tokenization_kwargs=tokenization_kwargs,
1377
)
1378
1381
- def _embedding_score(
1382
- self,
1383
- data_1: list[ScoreData],
1384
- data_2: list[ScoreData],
1385
- *,
1386
- use_tqdm: bool | Callable[..., tqdm],
1387
- pooling_params: PoolingParams | None,
1388
- lora_request: list[LoRARequest] | LoRARequest | None,
1389
- tokenization_kwargs: dict[str, Any],
1390
- ) -> list[ScoringRequestOutput]:
1391
- tokenizer = self.get_tokenizer()
1392
-
1393
- input_texts: list[str] = []
1394
- for text in data_1 + data_2:
1395
- if not isinstance(text, str):
1396
- raise NotImplementedError(
1397
- "Embedding scores currently do not support multimodal input."
1398
- )
1399
- input_texts.append(text)
1400
-
1401
- encoded_output = self.encode(
1402
- input_texts,
1403
- use_tqdm=use_tqdm,
1404
- lora_request=lora_request,
1405
- pooling_params=pooling_params,
1406
- pooling_task="embed",
1407
- tokenization_kwargs=tokenization_kwargs,
1408
- )
1409
-
1410
- encoded_output_1 = encoded_output[0 : len(data_1)]
1411
- encoded_output_2 = encoded_output[len(data_1) :]
1412
-
1413
- if len(encoded_output_1) == 1:
1414
- encoded_output_1 = encoded_output_1 * len(encoded_output_2)
1415
-
1416
- scores = _cosine_similarity(
1417
- tokenizer=tokenizer,
1418
- embed_1=encoded_output_1,
1419
- embed_2=encoded_output_2,
1420
- )
1421
-
1422
- return [ScoringRequestOutput.from_base(item) for item in scores]
1423
-
1424
- def _late_interaction_score(
1425
- self,
1426
- data_1: list[ScoreData],
1427
- data_2: list[ScoreData],
1428
- *,
1429
- use_tqdm: bool | Callable[..., tqdm],
1430
- pooling_params: PoolingParams | None,
1431
- lora_request: list[LoRARequest] | LoRARequest | None,
1432
- tokenization_kwargs: dict[str, Any],
1433
- ) -> list[ScoringRequestOutput]:
1434
- """
1435
- Late interaction scoring (ColBERT MaxSim).
1436
-
1437
- Encodes queries and documents into per-token embeddings, then computes
1438
- MaxSim: sum over query tokens of max similarity to any document token.
1439
- """
1440
- from vllm.outputs import PoolingOutput
1441
-
1442
- tokenizer = self.get_tokenizer()
1443
-
1444
- # Convert ScoreData to PromptType (handles both text and multimodal)
1445
- model_config = self.model_config
1446
- prompts_1 = score_data_to_prompts(data_1, "query", model_config)
1447
- prompts_2 = score_data_to_prompts(data_2, "document", model_config)
1448
-
1449
- encoded_output: list[PoolingRequestOutput] = self.encode(
1450
- prompts_1 + prompts_2,
1451
- use_tqdm=use_tqdm,
1452
- lora_request=lora_request,
1453
- pooling_params=pooling_params,
1454
- pooling_task="token_embed",
1455
- tokenization_kwargs=tokenization_kwargs,
1456
- )
1457
-
1458
- encoded_output_1: list[PoolingRequestOutput] = encoded_output[: len(prompts_1)]
1459
- encoded_output_2: list[PoolingRequestOutput] = encoded_output[len(prompts_1) :]
1460
-
1461
- if len(encoded_output_1) == 1:
1462
- encoded_output_1 = encoded_output_1 * len(encoded_output_2)
1463
-
1464
- # Compute MaxSim scores
1465
- scores: list[PoolingRequestOutput] = []
1466
- padding: list[int] = []
1467
- if (pad_token_id := tokenizer.pad_token_id) is not None:
1468
- padding = [pad_token_id]
1469
-
1470
- for emb_1, emb_2 in zip(encoded_output_1, encoded_output_2):
1471
- # emb_1.outputs.data: [query_len, dim]
1472
- # emb_2.outputs.data: [doc_len, dim]
1473
- q_emb = emb_1.outputs.data
1474
- d_emb = emb_2.outputs.data
1475
-
1476
- maxsim_score = compute_maxsim_score(q_emb, d_emb)
1477
-
1478
- tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids
1479
-
1480
- scores.append(
1481
- PoolingRequestOutput(
1482
- request_id=f"{emb_1.request_id}_{emb_2.request_id}",
1483
- outputs=PoolingOutput(data=maxsim_score),
1484
- prompt_token_ids=tokens,
1485
- num_cached_tokens=emb_1.num_cached_tokens + emb_2.num_cached_tokens,
1486
- finished=True,
1487
- )
1488
- )
1489
-
1490
- return [ScoringRequestOutput.from_base(item) for item in scores]
1491
-
1492
- def _cross_encoding_score(
1493
- self,
1494
- data_1: list[ScoreData],
1495
- data_2: list[ScoreData],
1496
- *,
1497
- use_tqdm: bool | Callable[..., tqdm],
1498
- pooling_params: PoolingParams | None,
1499
- lora_request: list[LoRARequest] | LoRARequest | None,
1500
- tokenization_kwargs: dict[str, Any],
1501
- score_template: str | None,
1502
- ) -> list[ScoringRequestOutput]:
1503
- model_config = self.model_config
1504
- tokenizer = self.get_tokenizer()
1505
-
1506
- if is_mistral_tokenizer(tokenizer):
1507
- raise ValueError("Score API is not supported for Mistral tokenizer")
1508
-
1509
- if len(data_1) == 1:
1510
- data_1 = data_1 * len(data_2)
1511
-
1512
- if pooling_params is None:
1513
- pooling_params = PoolingParams(task="classify")
1514
- elif pooling_params.task is None:
1515
- pooling_params.task = "classify"
1516
-
1517
- pooling_params_list = list[PoolingParams]()
1518
-
1519
- prompts = list[PromptType]()
1520
-
1521
- input_pairs = [(t1, t2) for t1, t2 in zip(data_1, data_2)]
1522
-
1523
- for q, d in input_pairs:
1524
- _, engine_prompt = get_score_prompt(
1525
- model_config=model_config,
1526
- data_1=q,
1527
- data_2=d,
1528
- tokenizer=tokenizer,
1529
- tokenization_kwargs=tokenization_kwargs,
1530
- score_template=score_template,
1531
- )
1532
-
1533
- if token_type_ids := engine_prompt.pop("token_type_ids", None):
1534
- params = pooling_params.clone()
1535
- compressed = compress_token_type_ids(token_type_ids)
1536
- params.extra_kwargs = {"compressed_token_type_ids": compressed}
1537
- pooling_params_list.append(params)
1538
- else:
1539
- pooling_params_list.append(pooling_params)
1540
-
1541
- prompts.append(engine_prompt)
1542
-
1543
- outputs = self._run_completion(
1544
- prompts=prompts,
1545
- params=pooling_params_list,
1546
- output_type=PoolingRequestOutput,
1547
- use_tqdm=use_tqdm,
1548
- lora_request=lora_request,
1549
- )
1550
-
1551
- return [ScoringRequestOutput.from_base(item) for item in outputs]
1552
-
1379
def score(
1380
self,
1555
- data_1: SingletonPrompt
1556
- | Sequence[SingletonPrompt]
1557
- | ScoreMultiModalParam
1558
- | list[ScoreMultiModalParam],
1559
- data_2: SingletonPrompt
1560
- | Sequence[SingletonPrompt]
1561
- | ScoreMultiModalParam
1562
- | list[ScoreMultiModalParam],
1381
+ data_1: ScoreInput | list[ScoreInput],
1382
+ data_2: ScoreInput | list[ScoreInput],
1383
/,
1384
*,
1385
use_tqdm: bool | Callable[..., tqdm] = True,
@@ -1606,83 +1426,71 @@ class LLM:
1426
A list of `ScoringRequestOutput` objects containing the
1427
generated scores in the same order as the input prompts.
1428
"""
1609
- model_config = self.model_config
1429
1611
- runner_type = model_config.runner_type
1612
- if runner_type != "pooling":
1430
+ if self.runner_type != "pooling":
1431
raise ValueError(
1432
"LLM.score() is only supported for pooling models. "
1433
"Try passing `--runner pooling` to use the model as a "
1434
"pooling model."
1435
)
1436
1619
- supported_tasks = self.supported_tasks
1437
score_type = self.model_config.score_type
1621
- is_late_interaction = score_type == "late-interaction"
1622
- is_cross_encoder = score_type == "cross-encoder"
1623
-
1624
- # Late interaction models (e.g., ColBERT) use token_embed for scoring
1625
- if not is_late_interaction and all(
1626
- t not in supported_tasks for t in ("embed", "classify")
1438
+ if (
1439
+ score_type == "cross-encoder"
1440
+ and getattr(self.model_config.hf_config, "num_labels", 0) != 1
1441
):
1628
- raise ValueError(
1629
- "Score API is not supported by this model. "
1630
- "Try converting the model using "
1631
- "`--convert embed` or `--convert classify`."
1632
- )
1442
+ raise ValueError("Scoring API is only enabled for num_labels == 1.")
1443
1634
- if is_cross_encoder and getattr(model_config.hf_config, "num_labels", 0) != 1:
1635
- raise ValueError("Score API is only enabled for num_labels == 1.")
1444
+ if score_type is None or score_type not in self.pooling_io_processors:
1445
+ raise ValueError("This model does not support the Scoring API.")
1446
1637
- if not is_cross_encoder and chat_template is not None:
1638
- raise ValueError(
1639
- "chat_template is only supported for cross-encoder models."
1640
- )
1447
+ io_processor = self.pooling_io_processors[score_type]
1448
+ assert isinstance(io_processor, ScoringIOProcessor)
1449
1642
- is_multimodal_model = model_config.is_multimodal_model
1643
- architecture = model_config.architecture
1450
+ pooling_task = io_processor.pooling_task
1451
+ scoring_data = io_processor.valid_inputs(data_1, data_2)
1452
+ offset = len(scoring_data.data_1)
1453
1645
- score_data_1, score_data_2 = validate_score_input(
1646
- data_1, # type: ignore[arg-type]
1647
- data_2, # type: ignore[arg-type]
1648
- is_multimodal_model=is_multimodal_model,
1649
- architecture=architecture,
1454
+ ctx = OfflineInputsContext(
1455
+ prompts=scoring_data,
1456
+ pooling_params=pooling_params,
1457
+ tokenization_kwargs=tokenization_kwargs,
1458
+ chat_template=chat_template,
1459
+ offset=offset,
1460
)
1461
1652
- renderer = self.renderer
1653
- tok_params = renderer.default_cmpl_tok_params.with_kwargs(
1654
- **(tokenization_kwargs or {})
1462
+ processor_inputs = io_processor.pre_process_offline(ctx)
1463
+
1464
+ seq_lora_requests = self._lora_request_to_seq(
1465
+ lora_request, len(processor_inputs)
1466
)
1656
- encode_kwargs = tok_params.get_encode_kwargs()
1467
1658
- if is_cross_encoder:
1659
- return self._cross_encoding_score(
1660
- score_data_1,
1661
- score_data_2,
1662
- use_tqdm=use_tqdm,
1663
- pooling_params=pooling_params,
1664
- lora_request=lora_request,
1665
- tokenization_kwargs=encode_kwargs,
1666
- score_template=chat_template,
1667
- )
1668
- elif is_late_interaction:
1669
- return self._late_interaction_score(
1670
- score_data_1,
1671
- score_data_2,
1672
- use_tqdm=use_tqdm,
1673
- pooling_params=pooling_params,
1674
- lora_request=lora_request,
1675
- tokenization_kwargs=encode_kwargs,
1676
- )
1677
- else:
1678
- return self._embedding_score(
1679
- score_data_1,
1680
- score_data_2,
1681
- use_tqdm=use_tqdm,
1682
- pooling_params=pooling_params,
1683
- lora_request=lora_request,
1684
- tokenization_kwargs=encode_kwargs,
1685
- )
1468
+ if ctx.pooling_params is None:
1469
+ ctx.pooling_params = PoolingParams()
1470
+ params_seq = self._params_to_seq(ctx.pooling_params, len(processor_inputs))
1471
+
1472
+ for param in params_seq:
1473
+ if param.task is None:
1474
+ param.task = pooling_task
1475
+ elif param.task != pooling_task:
1476
+ msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"
1477
+ raise ValueError(msg)
1478
+
1479
+ seq_priority = self._priority_to_seq(None, len(processor_inputs))
1480
+
1481
+ self._render_and_add_requests(
1482
+ prompts=processor_inputs,
1483
+ params=params_seq,
1484
+ lora_requests=seq_lora_requests,
1485
+ priorities=seq_priority,
1486
+ )
1487
+
1488
+ outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput)
1489
+ outputs = io_processor.post_process_offline(
1490
+ ctx=OfflineOutputsContext(outputs=outputs, offset=offset),
1491
+ )
1492
+
1493
+ return [ScoringRequestOutput.from_base(item) for item in outputs]
1494
1495
def start_profile(self, profile_prefix: str | None = None) -> None:
1496
"""Start profiling with optional custom trace prefix.
vllm/entrypoints/openai/engine/serving.py
+4
-104
@@ -11,9 +11,7 @@ from typing import Any, ClassVar, Generic, Protocol, TypeAlias, TypeVar
11
12
import numpy as np
13
from fastapi import Request
14
-from openai.types.responses import (
15
- ToolChoiceFunction,
16
-)
14
+from openai.types.responses import ToolChoiceFunction
15
from pydantic import ConfigDict, TypeAdapter, ValidationError
16
from starlette.datastructures import Headers
17
@@ -21,9 +19,7 @@ import vllm.envs as envs
19
from vllm.beam_search import BeamSearchSequence, create_sort_beams_key_function
20
from vllm.config import ModelConfig
21
from vllm.engine.protocol import EngineClient
24
-from vllm.entrypoints.chat_utils import (
25
- ChatTemplateContentFormatOption,
26
-)
22
+from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
23
from vllm.entrypoints.logger import RequestLogger
24
from vllm.entrypoints.openai.chat_completion.protocol import (
25
BatchChatCompletionRequest,
@@ -42,9 +38,7 @@ from vllm.entrypoints.openai.engine.protocol import (
38
GenerationError,
39
)
40
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
45
-from vllm.entrypoints.openai.responses.protocol import (
46
- ResponsesRequest,
47
-)
41
+from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
42
from vllm.entrypoints.openai.speech_to_text.protocol import (
43
TranscriptionRequest,
44
TranscriptionResponse,
@@ -56,14 +50,6 @@ from vllm.entrypoints.pooling.pooling.protocol import (
50
PoolingCompletionRequest,
51
PoolingResponse,
52
)
59
-from vllm.entrypoints.pooling.score.protocol import (
60
- RerankRequest,
61
- ScoreDataRequest,
62
- ScoreQueriesDocumentsRequest,
63
- ScoreRequest,
64
- ScoreResponse,
65
- ScoreTextRequest,
66
-)
53
from vllm.entrypoints.serve.disagg.protocol import GenerateRequest, GenerateResponse
54
from vllm.entrypoints.serve.tokenize.protocol import (
55
DetokenizeRequest,
@@ -72,8 +58,7 @@ from vllm.entrypoints.serve.tokenize.protocol import (
58
TokenizeResponse,
59
)
60
from vllm.entrypoints.utils import create_error_response
75
-from vllm.exceptions import VLLMValidationError
76
-from vllm.inputs import EngineInput, PromptType, TokensPrompt
61
+from vllm.inputs import EngineInput, PromptType
62
from vllm.logger import init_logger
63
from vllm.logprobs import Logprob, PromptLogprobs
64
from vllm.lora.request import LoRARequest
@@ -119,8 +104,6 @@ CompletionLikeRequest: TypeAlias = (
104
CompletionRequest
105
| TokenizeCompletionRequest
106
| DetokenizeRequest
122
- | RerankRequest
123
- | ScoreRequest
107
| PoolingCompletionRequest
108
)
109
@@ -148,7 +131,6 @@ AnyResponse: TypeAlias = (
131
| TranscriptionResponse
132
| TokenizeResponse
133
| PoolingResponse
151
- | ScoreResponse
134
| GenerateResponse
135
)
136
@@ -692,88 +674,6 @@ class OpenAIServing:
674
message_types.add(content_dict["type"].split("_")[0])
675
return message_types
676
695
- def _validate_input(
696
- self,
697
- request: object,
698
- input_ids: list[int],
699
- input_text: str,
700
- ) -> TokensPrompt:
701
- token_num = len(input_ids)
702
- max_model_len = self.model_config.max_model_len
703
-
704
- # Note: ScoreRequest doesn't have max_tokens
705
- if isinstance(
706
- request,
707
- (
708
- ScoreDataRequest,
709
- ScoreTextRequest,
710
- ScoreQueriesDocumentsRequest,
711
- RerankRequest,
712
- ),
713
- ):
714
- # Note: input length can be up to the entire model context length
715
- # since these requests don't generate tokens.
716
- if token_num > max_model_len:
717
- operations: dict[type[AnyRequest], str] = {
718
- ScoreDataRequest: "score",
719
- ScoreTextRequest: "score",
720
- ScoreQueriesDocumentsRequest: "score",
721
- }
722
- operation = operations.get(type(request), "embedding generation")
723
- raise VLLMValidationError(
724
- f"This model's maximum context length is "
725
- f"{max_model_len} tokens. However, you requested "
726
- f"{token_num} tokens in the input for {operation}. "
727
- f"Please reduce the length of the input prompt.",
728
- parameter="input_tokens",
729
- value=token_num,
730
- )
731
- return TokensPrompt(prompt=input_text, prompt_token_ids=input_ids)
732
-
733
- # Note: TokenizeRequest and DetokenizeRequest doesn't have max_tokens
734
- # and does not require model context length validation
735
- if isinstance(
736
- request,
737
- (TokenizeCompletionRequest, TokenizeChatRequest, DetokenizeRequest),
738
- ):
739
- return TokensPrompt(prompt=input_text, prompt_token_ids=input_ids)
740
-
741
- # chat completion endpoint supports max_completion_tokens
742
- if isinstance(request, ChatCompletionRequest):
743
- # TODO(#9845): remove max_tokens when field dropped from OpenAI API
744
- max_tokens = request.max_completion_tokens or request.max_tokens
745
- else:
746
- max_tokens = getattr(request, "max_tokens", None)
747
-
748
- # Note: input length can be up to model context length - 1 for
749
- # completion-like requests.
750
- if token_num >= max_model_len:
751
- raise VLLMValidationError(
752
- f"This model's maximum context length is "
753
- f"{max_model_len} tokens. However, your request has "
754
- f"{token_num} input tokens. Please reduce the length of "
755
- "the input messages.",
756
- parameter="input_tokens",
757
- value=token_num,
758
- )
759
-
760
- if max_tokens is not None and token_num + max_tokens > max_model_len:
761
- raise VLLMValidationError(
762
- f"This model's maximum context length is "
763
- f"{max_model_len} tokens. However, you requested "
764
- f"{max_tokens} output tokens and your prompt contains "
765
- f"{token_num} input tokens, for a total of "
766
- f"{token_num + max_tokens} tokens "
767
- f"({token_num} + {max_tokens} = "
768
- f"{token_num + max_tokens} > {max_model_len}). "
769
- f"Please reduce the length of the input prompt or the "
770
- f"number of requested output tokens.",
771
- parameter="max_tokens",
772
- value=max_tokens,
773
- )
774
-
775
- return TokensPrompt(prompt=input_text, prompt_token_ids=input_ids)
776
-
677
def _validate_chat_template(
678
self,
679
request_chat_template: str | None,
vllm/entrypoints/openai/run_batch.py
+27
-27
@@ -2,6 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import asyncio
5
+import contextlib
6
+import json
7
import sys
8
import tempfile
9
from argparse import Namespace
@@ -13,12 +15,14 @@ from urllib.parse import urlparse
15
16
import aiohttp
17
import pybase64 as base64
18
+import pydantic
19
import torch
20
from fastapi import UploadFile
21
from prometheus_client import start_http_server
22
from pydantic import Field, TypeAdapter, field_validator, model_validator
23
from pydantic_core.core_schema import ValidationInfo
24
from starlette.datastructures import State
25
+from starlette.responses import JSONResponse
26
from tqdm import tqdm
27
from urllib3.util import parse_url
28
@@ -49,7 +53,7 @@ from vllm.entrypoints.pooling.embed.protocol import (
53
EmbeddingRequest,
54
EmbeddingResponse,
55
)
52
-from vllm.entrypoints.pooling.score.protocol import (
56
+from vllm.entrypoints.pooling.scoring.protocol import (
57
RerankRequest,
58
RerankResponse,
59
ScoreRequest,
@@ -180,6 +184,18 @@ class BatchRequestInput(OpenAIBaseModel):
184
return TypeAdapter(BatchRequestInputBody).validate_python(value)
185
186
187
+AllResponse: TypeAlias = (
188
+ ChatCompletionResponse
189
+ | EmbeddingResponse
190
+ | ScoreResponse
191
+ | RerankResponse
192
+ | TranscriptionResponse
193
+ | TranscriptionResponseVerbose
194
+ | TranslationResponse
195
+ | TranslationResponseVerbose
196
+)
197
+
198
+
199
class BatchResponseData(OpenAIBaseModel):
200
# HTTP status code of the response.
201
status_code: int = 200
@@ -188,17 +204,7 @@ class BatchResponseData(OpenAIBaseModel):
204
request_id: str
205
206
# The body of the response.
191
- body: (
192
- ChatCompletionResponse
193
- | EmbeddingResponse
194
- | ScoreResponse
195
- | RerankResponse
196
- | TranscriptionResponse
197
- | TranscriptionResponseVerbose
198
- | TranslationResponse
199
- | TranslationResponseVerbose
200
- | None
201
- ) = None
207
+ body: AllResponse | None = None
208
209
210
class BatchRequestOutput(OpenAIBaseModel):
@@ -536,19 +542,13 @@ async def run_request(
542
except Exception as e:
543
response = create_error_response(e)
544
539
- if isinstance(
540
- response,
541
- (
542
- ChatCompletionResponse,
543
- EmbeddingResponse,
544
- ScoreResponse,
545
- RerankResponse,
546
- TranscriptionResponse,
547
- TranscriptionResponseVerbose,
548
- TranslationResponse,
549
- TranslationResponseVerbose,
550
- ),
551
- ):
545
+ if isinstance(response, JSONResponse):
546
+ with contextlib.suppress(pydantic.ValidationError):
547
+ response = TypeAdapter(AllResponse | ErrorResponse).validate_python(
548
+ json.loads(response.body)
549
+ )
550
+
551
+ if isinstance(response, AllResponse):
552
batch_output = BatchRequestOutput(
553
id=f"vllm-{random_uuid()}",
554
custom_id=request.custom_id,
@@ -745,14 +745,14 @@ async def build_endpoint_registry(
745
"score": {
746
"url_matcher": lambda url: url.endswith("/score"),
747
"handler_getter": lambda: (
748
- serving_scores.create_score if serving_scores is not None else None
748
+ serving_scores if serving_scores is not None else None
749
),
750
"wrapper_fn": None,
751
},
752
"rerank": {
753
"url_matcher": lambda url: url.endswith("/rerank"),
754
"handler_getter": lambda: (
755
- serving_scores.do_rerank if serving_scores is not None else None
755
+ serving_scores if serving_scores is not None else None
756
),
757
"wrapper_fn": None,
758
},
vllm/entrypoints/pooling/__init__.py
+6
-21
@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING
6
from fastapi import FastAPI
7
8
from vllm.config import ModelConfig
9
+from vllm.entrypoints.pooling.utils import enable_scoring_api
10
from vllm.logger import init_logger
11
12
if TYPE_CHECKING:
@@ -23,23 +24,6 @@ else:
24
logger = init_logger(__name__)
25
26
26
-def enable_scoring_api(
27
- supported_tasks: tuple["SupportedTask", ...],
28
- model_config: ModelConfig | None = None,
29
-) -> bool:
30
- if any(t in supported_tasks for t in ("embed", "token_embed")):
31
- return True
32
-
33
- if model_config is not None and "classify" in supported_tasks:
34
- num_labels = getattr(model_config.hf_config, "num_labels", 0)
35
- if num_labels != 1:
36
- logger.debug_once("Score API is only enabled for num_labels == 1.")
37
- return False
38
- return True
39
-
40
- return False
41
-
42
-
27
def register_pooling_api_routers(
28
app: FastAPI,
29
supported_tasks: tuple["SupportedTask", ...],
@@ -68,7 +52,7 @@ def register_pooling_api_routers(
52
app.include_router(embed_router)
53
54
if enable_scoring_api(supported_tasks, model_config):
71
- from vllm.entrypoints.pooling.score.api_router import router as score_router
55
+ from vllm.entrypoints.pooling.scoring.api_router import router as score_router
56
57
app.include_router(score_router)
58
@@ -84,7 +68,7 @@ def init_pooling_state(
68
from vllm.entrypoints.pooling.classify.serving import ServingClassification
69
from vllm.entrypoints.pooling.embed.serving import ServingEmbedding
70
from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling
87
- from vllm.entrypoints.pooling.score.serving import ServingScores
71
+ from vllm.entrypoints.pooling.scoring.serving import ServingScores
72
from vllm.tasks import POOLING_TASKS
73
74
model_config = engine_client.model_config
@@ -136,8 +120,9 @@ def init_pooling_state(
120
engine_client,
121
state.openai_serving_models,
122
request_logger=request_logger,
139
- score_template=resolved_chat_template,
140
- log_error_stack=args.log_error_stack,
123
+ chat_template=resolved_chat_template,
124
+ chat_template_content_format=args.chat_template_content_format,
125
+ trust_request_chat_template=args.trust_request_chat_template,
126
)
127
if enable_scoring_api(supported_tasks, model_config)
128
else None
vllm/entrypoints/pooling/base/io_processor.py
+21
-25
@@ -13,13 +13,16 @@ from vllm.entrypoints.chat_utils import (
13
ConversationMessage,
14
)
15
from vllm.entrypoints.openai.engine.serving import RendererChatRequest, RendererRequest
16
+from vllm.entrypoints.pooling.scoring.typing import ScoringData
17
from vllm.entrypoints.pooling.typing import (
18
+ OfflineInputsContext,
19
+ OfflineOutputsContext,
20
PoolingChatLikeRequest,
21
PoolingCompletionLikeRequest,
22
PoolingServeContext,
23
)
24
from vllm.inputs import EngineInput, SingletonPrompt
22
-from vllm.renderers import BaseRenderer, merge_kwargs
25
+from vllm.renderers import BaseRenderer, TokenizeParams, merge_kwargs
26
from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq
27
from vllm.tool_parsers import ToolParser
28
from vllm.utils.mistral import is_mistral_tokenizer
@@ -96,29 +99,29 @@ class PoolingIOProcessor:
99
#######################################
100
# offline APIs
101
99
- def pre_process_offline(
100
- self,
101
- prompts: PromptType | Sequence[PromptType],
102
- tokenization_kwargs: dict[str, Any] | None = None,
103
- ) -> Sequence[EngineInput]:
102
+ def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:
103
+ assert not isinstance(ctx.prompts, ScoringData)
104
+ tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(
105
+ **(ctx.tokenization_kwargs or {})
106
+ )
107
return self._preprocess_completion_offline(
105
- prompts=prompts, tokenization_kwargs=tokenization_kwargs
108
+ prompts=ctx.prompts, tok_params=tok_params
109
)
110
108
- async def pre_process_offline_async(self, *args, **kwargs):
109
- return self.pre_process_offline(*args, **kwargs)
111
+ async def pre_process_offline_async(self, ctx: OfflineInputsContext):
112
+ return self.pre_process_offline(ctx)
113
114
def post_process_offline(
115
self,
113
- outputs: list[PoolingRequestOutput],
116
+ ctx: OfflineOutputsContext,
117
) -> list[PoolingRequestOutput]:
115
- return outputs
118
+ return ctx.outputs
119
120
async def post_process_offline_async(
121
self,
119
- outputs: list[PoolingRequestOutput],
122
+ ctx: OfflineOutputsContext,
123
) -> list[PoolingRequestOutput]:
121
- return self.post_process_offline(outputs)
124
+ return self.post_process_offline(ctx)
125
126
#######################################
127
# helpers
@@ -204,28 +207,21 @@ class PoolingIOProcessor:
207
def _preprocess_completion_offline(
208
self,
209
prompts: PromptType | Sequence[PromptType],
207
- tokenization_kwargs: dict[str, Any] | None = None,
210
+ tok_params: TokenizeParams,
211
+ prompt_extras: dict[str, Any] | None = None,
212
) -> Sequence[EngineInput]:
209
- renderer = self.renderer
210
- model_config = self.model_config
211
-
213
prompts = prompt_to_seq(prompts)
213
-
214
parsed_prompts = [
215
(
216
prompt
217
if isinstance(prompt, bytes)
218
- else parse_model_prompt(model_config, prompt)
218
+ else parse_model_prompt(self.model_config, prompt)
219
)
220
for prompt in prompts
221
]
222
- tok_params = renderer.default_cmpl_tok_params.with_kwargs(
223
- **(tokenization_kwargs or {})
224
- )
222
226
- return renderer.render_cmpl(
227
- parsed_prompts,
228
- tok_params,
223
+ return self.renderer.render_cmpl(
224
+ parsed_prompts, tok_params, prompt_extras=prompt_extras
225
)
226
227
def _validate_chat_template(
vllm/entrypoints/pooling/base/serving.py
+18
-4
@@ -117,8 +117,16 @@ class PoolingServing:
117
else await self._get_trace_headers(ctx.raw_request.headers)
118
)
119
120
- pooling_params = self.io_processor.create_pooling_params(ctx.request)
121
- pooling_params.verify(self.model_config)
120
+ if ctx.pooling_params is None:
121
+ pooling_params = self.io_processor.create_pooling_params(ctx.request)
122
+ else:
123
+ pooling_params = ctx.pooling_params
124
+
125
+ if isinstance(pooling_params, list):
126
+ for params in pooling_params:
127
+ params.verify(self.model_config)
128
+ else:
129
+ pooling_params.verify(self.model_config)
130
131
for i, engine_input in enumerate(ctx.engine_inputs):
132
prompt_request_id = (
@@ -127,16 +135,22 @@ class PoolingServing:
135
else ctx.prompt_request_ids[i]
136
)
137
138
+ params = (
139
+ pooling_params[i]
140
+ if isinstance(pooling_params, list)
141
+ else pooling_params
142
+ )
143
+
144
self._log_inputs(
145
prompt_request_id,
146
engine_input,
133
- params=pooling_params,
147
+ params=params,
148
lora_request=ctx.lora_request,
149
)
150
151
generator = self.engine_client.encode(
152
engine_input,
139
- pooling_params,
153
+ params,
154
prompt_request_id,
155
lora_request=ctx.lora_request,
156
trace_headers=trace_headers,
vllm/entrypoints/pooling/io_processor_factories.py
+7
@@ -5,6 +5,8 @@
5
from vllm.config import ModelConfig
6
from vllm.entrypoints.chat_utils import ChatTemplateConfig
7
from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
8
+from vllm.entrypoints.pooling.scoring.io_processor import ScoringIOProcessors
9
+from vllm.entrypoints.pooling.utils import enable_scoring_api
10
from vllm.renderers import BaseRenderer
11
from vllm.tasks import SupportedTask
12
@@ -25,6 +27,11 @@ def init_pooling_io_processors(
27
28
processors.append(("embed", EmbedIOProcessor))
29
30
+ if enable_scoring_api(supported_tasks, model_config):
31
+ score_type = model_config.score_type
32
+ if score_type is not None and score_type in ScoringIOProcessors:
33
+ processors.append((score_type, ScoringIOProcessors[score_type]))
34
+
35
return {
36
task: processor_cls(
37
model_config=model_config,
vllm/entrypoints/pooling/score/serving.py
deleted
-667
@@ -1,667 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import asyncio
4
-import time
5
-from collections.abc import AsyncGenerator, Mapping
6
-from concurrent.futures import ThreadPoolExecutor
7
-from typing import Any
8
-
9
-from fastapi import Request
10
-
11
-from vllm.engine.protocol import EngineClient
12
-from vllm.entrypoints.logger import RequestLogger
13
-from vllm.entrypoints.openai.engine.protocol import (
14
- ErrorResponse,
15
- UsageInfo,
16
-)
17
-from vllm.entrypoints.openai.engine.serving import OpenAIServing
18
-from vllm.entrypoints.openai.models.serving import OpenAIServingModels
19
-from vllm.entrypoints.pooling.score.protocol import (
20
- RerankDocument,
21
- RerankRequest,
22
- RerankResponse,
23
- RerankResult,
24
- RerankUsage,
25
- ScoreRequest,
26
- ScoreResponse,
27
- ScoreResponseData,
28
-)
29
-from vllm.entrypoints.pooling.score.utils import (
30
- ScoreData,
31
- ScoreInputs,
32
- _cosine_similarity,
33
- compress_token_type_ids,
34
- get_score_prompt,
35
- parse_score_data_single,
36
- validate_score_input,
37
-)
38
-from vllm.inputs import EngineInput, TokensPrompt, tokens_input
39
-from vllm.logger import init_logger
40
-from vllm.lora.request import LoRARequest
41
-from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
42
-from vllm.tokenizers import TokenizerLike
43
-from vllm.utils.async_utils import make_async, merge_async_iterators
44
-from vllm.utils.mistral import is_mistral_tokenizer
45
-from vllm.v1.pool.late_interaction import (
46
- build_late_interaction_doc_params,
47
- build_late_interaction_query_params,
48
-)
49
-
50
-logger = init_logger(__name__)
51
-
52
-
53
-class ServingScores(OpenAIServing):
54
- def __init__(
55
- self,
56
- engine_client: EngineClient,
57
- models: OpenAIServingModels,
58
- *,
59
- request_logger: RequestLogger | None,
60
- score_template: str | None = None,
61
- log_error_stack: bool = False,
62
- ) -> None:
63
- super().__init__(
64
- engine_client=engine_client,
65
- models=models,
66
- request_logger=request_logger,
67
- )
68
- self.score_template = score_template
69
-
70
- self._tokenizer_executor = ThreadPoolExecutor(max_workers=1)
71
-
72
- self.score_type = self.model_config.score_type
73
- self.architecture = self.model_config.architecture
74
- self.is_multimodal_model = self.model_config.is_multimodal_model
75
-
76
- if self.score_type == "cross-encoder":
77
- self._score_func = self._cross_encoding_score
78
- elif self.score_type == "late-interaction":
79
- self._score_func = self._late_interaction_score
80
- else: # "bi-encoder"
81
- self._score_func = self._embedding_score
82
-
83
- async def _embedding_score(
84
- self,
85
- data_1: list[ScoreData],
86
- data_2: list[ScoreData],
87
- request: RerankRequest | ScoreRequest,
88
- request_id: str,
89
- lora_request: LoRARequest | None | None = None,
90
- trace_headers: Mapping[str, str] | None = None,
91
- ) -> list[PoolingRequestOutput] | ErrorResponse:
92
- input_texts: list[str] = []
93
- for text in data_1 + data_2:
94
- if not isinstance(text, str):
95
- raise NotImplementedError(
96
- "Embedding scores currently do not support multimodal input."
97
- )
98
- input_texts.append(text)
99
-
100
- model_config = self.model_config
101
- tokenizer = self.renderer.get_tokenizer()
102
-
103
- encode_async = make_async(
104
- tokenizer.encode,
105
- executor=self._tokenizer_executor,
106
- )
107
-
108
- tokenization_kwargs = request.build_tok_params(model_config).get_encode_kwargs()
109
- tokenized_prompts = await asyncio.gather(
110
- *(encode_async(t, **tokenization_kwargs) for t in input_texts)
111
- )
112
-
113
- engine_inputs: list[EngineInput] = []
114
- for tok_result, input_text in zip(tokenized_prompts, input_texts):
115
- text_token_prompt = self._validate_input(request, tok_result, input_text)
116
-
117
- engine_inputs.append(
118
- tokens_input(
119
- text_token_prompt["prompt_token_ids"],
120
- prompt=input_text,
121
- )
122
- )
123
-
124
- # Schedule the request and get the result generator.
125
- generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
126
- pooling_params = request.to_pooling_params("embed")
127
-
128
- for i, engine_input in enumerate(engine_inputs):
129
- request_id_item = f"{request_id}-{i}"
130
-
131
- self._log_inputs(
132
- request_id_item,
133
- engine_input,
134
- params=pooling_params,
135
- lora_request=lora_request,
136
- )
137
-
138
- generators.append(
139
- self.engine_client.encode(
140
- engine_input,
141
- pooling_params,
142
- request_id_item,
143
- lora_request=lora_request,
144
- trace_headers=trace_headers,
145
- priority=request.priority,
146
- )
147
- )
148
-
149
- result_generator = merge_async_iterators(*generators)
150
-
151
- # Non-streaming response
152
- final_res_batch: list[PoolingRequestOutput] = []
153
-
154
- embeddings: list[PoolingRequestOutput | None] = [None] * len(engine_inputs)
155
-
156
- async for i, res in result_generator:
157
- embeddings[i] = res
158
-
159
- emb_data_1: list[PoolingRequestOutput] = []
160
- emb_data_2: list[PoolingRequestOutput] = []
161
-
162
- for i in range(0, len(data_1)):
163
- assert (emb := embeddings[i]) is not None
164
- emb_data_1.append(emb)
165
-
166
- for i in range(len(data_1), len(embeddings)):
167
- assert (emb := embeddings[i]) is not None
168
- emb_data_2.append(emb)
169
-
170
- if len(emb_data_1) == 1:
171
- emb_data_1 = emb_data_1 * len(emb_data_2)
172
-
173
- final_res_batch = _cosine_similarity(
174
- tokenizer=tokenizer, embed_1=emb_data_1, embed_2=emb_data_2
175
- )
176
-
177
- return final_res_batch
178
-
179
- def _preprocess_late_interaction_item(
180
- self,
181
- data: ScoreData,
182
- role: str,
183
- request: RerankRequest | ScoreRequest,
184
- tokenizer: TokenizerLike,
185
- tokenization_kwargs: dict[str, Any],
186
- ) -> TokensPrompt:
187
- """Parse a single ScoreData into a text + optional multimodal
188
- TokensPrompt for late-interaction encoding.
189
-
190
- For plain strings, tokenises directly.
191
- For multimodal content parts, extracts text and multi_modal_data.
192
- """
193
- model_config = self.model_config
194
-
195
- if isinstance(data, str):
196
- text, mm_data, mm_uuids = data, None, None
197
- else:
198
- text, mm_data, mm_uuids = parse_score_data_single(data, role, model_config)
199
-
200
- prompt_ids = tokenizer.encode(text, **tokenization_kwargs)
201
- self._validate_input(request, prompt_ids, text)
202
-
203
- tok_prompt = TokensPrompt(
204
- prompt_token_ids=prompt_ids,
205
- prompt=text,
206
- )
207
-
208
- if mm_data is not None:
209
- tok_prompt["multi_modal_data"] = mm_data
210
- if mm_uuids is not None:
211
- tok_prompt["multi_modal_uuids"] = mm_uuids
212
- if request.mm_processor_kwargs is not None:
213
- tok_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs
214
-
215
- return tok_prompt
216
-
217
- async def _late_interaction_score(
218
- self,
219
- data_1: list[ScoreData],
220
- data_2: list[ScoreData],
221
- request: RerankRequest | ScoreRequest,
222
- request_id: str,
223
- lora_request: LoRARequest | None = None,
224
- trace_headers: Mapping[str, str] | None = None,
225
- ) -> list[PoolingRequestOutput] | ErrorResponse:
226
- """
227
- Late interaction scoring (ColBERT MaxSim).
228
-
229
- Encodes queries and documents into per-token embeddings, then computes
230
- MaxSim: sum over query tokens of max similarity to any document token.
231
- """
232
- model_config = self.model_config
233
- tokenizer = self.renderer.get_tokenizer()
234
- tokenization_kwargs = request.build_tok_params(model_config).get_encode_kwargs()
235
-
236
- all_data = data_1 + data_2
237
- roles = ["query"] * len(data_1) + ["document"] * len(data_2)
238
-
239
- preprocess_async = make_async(
240
- self._preprocess_late_interaction_item,
241
- executor=self._tokenizer_executor,
242
- )
243
-
244
- tok_prompts = await asyncio.gather(
245
- *(
246
- preprocess_async(
247
- data=d,
248
- role=r,
249
- request=request,
250
- tokenizer=tokenizer,
251
- tokenization_kwargs=tokenization_kwargs,
252
- )
253
- for d, r in zip(all_data, roles)
254
- )
255
- )
256
-
257
- query_prompts = tok_prompts[: len(data_1)]
258
- doc_prompts = tok_prompts[len(data_1) :]
259
-
260
- default_pooling_params = request.to_pooling_params("token_embed")
261
-
262
- # stage 1: encode queries and cache token embeddings on workers.
263
- query_keys = [f"{request_id}-query-{i}" for i in range(len(query_prompts))]
264
- query_uses = [len(doc_prompts) if len(query_prompts) == 1 else 1] * len(
265
- query_prompts
266
- )
267
- query_generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
268
- for i, tok_prompt in enumerate(query_prompts):
269
- request_id_item = f"{request_id}-query-{i}"
270
- pooling_params = default_pooling_params.clone()
271
- pooling_params.late_interaction_params = (
272
- build_late_interaction_query_params(
273
- query_key=query_keys[i],
274
- query_uses=query_uses[i],
275
- )
276
- )
277
-
278
- self._log_inputs(
279
- request_id_item,
280
- tok_prompt,
281
- params=pooling_params,
282
- lora_request=lora_request,
283
- )
284
-
285
- query_generators.append(
286
- self.engine_client.encode(
287
- tok_prompt,
288
- pooling_params,
289
- request_id_item,
290
- lora_request=lora_request,
291
- trace_headers=trace_headers,
292
- priority=request.priority,
293
- )
294
- )
295
-
296
- query_outputs: list[PoolingRequestOutput | None] = [None] * len(query_prompts)
297
- if query_generators:
298
- async for i, res in merge_async_iterators(*query_generators):
299
- query_outputs[i] = res
300
-
301
- assert all(res is not None for res in query_outputs)
302
- query_results = [res for res in query_outputs if res is not None]
303
-
304
- # stage 2: encode docs and return scalar scores from workers.
305
- doc_generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
306
- for i, tok_prompt in enumerate(doc_prompts):
307
- request_id_item = f"{request_id}-doc-{i}"
308
- query_idx = 0 if len(query_prompts) == 1 else i
309
- pooling_params = default_pooling_params.clone()
310
- pooling_params.late_interaction_params = build_late_interaction_doc_params(
311
- query_key=query_keys[query_idx]
312
- )
313
-
314
- self._log_inputs(
315
- request_id_item,
316
- tok_prompt,
317
- params=pooling_params,
318
- lora_request=lora_request,
319
- )
320
-
321
- doc_generators.append(
322
- self.engine_client.encode(
323
- tok_prompt,
324
- pooling_params,
325
- request_id_item,
326
- lora_request=lora_request,
327
- trace_headers=trace_headers,
328
- priority=request.priority,
329
- )
330
- )
331
-
332
- doc_outputs: list[PoolingRequestOutput | None] = [None] * len(doc_prompts)
333
- if doc_generators:
334
- async for i, res in merge_async_iterators(*doc_generators):
335
- doc_outputs[i] = res
336
-
337
- assert all(res is not None for res in doc_outputs)
338
- doc_results = [res for res in doc_outputs if res is not None]
339
-
340
- scores: list[PoolingRequestOutput] = []
341
- padding: list[int] = []
342
- if (pad_token_id := tokenizer.pad_token_id) is not None:
343
- padding = [pad_token_id]
344
-
345
- if len(query_results) == 1:
346
- query_results = query_results * len(doc_results)
347
-
348
- for query_result, doc_result in zip(query_results, doc_results):
349
- tokens = (
350
- query_result.prompt_token_ids + padding + doc_result.prompt_token_ids
351
- )
352
-
353
- scores.append(
354
- PoolingRequestOutput(
355
- request_id=f"{query_result.request_id}_{doc_result.request_id}",
356
- outputs=doc_result.outputs,
357
- prompt_token_ids=tokens,
358
- num_cached_tokens=(
359
- query_result.num_cached_tokens + doc_result.num_cached_tokens
360
- ),
361
- finished=True,
362
- )
363
- )
364
-
365
- return scores
366
-
367
- async def _cross_encoding_score(
368
- self,
369
- data_1: list[ScoreData],
370
- data_2: list[ScoreData],
371
- request: RerankRequest | ScoreRequest,
372
- request_id: str,
373
- lora_request: LoRARequest | None | None = None,
374
- trace_headers: Mapping[str, str] | None = None,
375
- ) -> list[PoolingRequestOutput] | ErrorResponse:
376
- tokenizer = self.renderer.get_tokenizer()
377
- if is_mistral_tokenizer(tokenizer):
378
- raise ValueError("MistralTokenizer not supported for cross-encoding")
379
-
380
- model_config = self.model_config
381
-
382
- if len(data_1) == 1:
383
- data_1 = data_1 * len(data_2)
384
-
385
- tok_kwargs = request.build_tok_params(model_config).get_encode_kwargs()
386
- input_pairs = [(t1, t2) for t1, t2 in zip(data_1, data_2)]
387
- preprocess_async = make_async(
388
- self._preprocess_score,
389
- executor=self._tokenizer_executor,
390
- )
391
- preprocessed_prompts = await asyncio.gather(
392
- *(
393
- preprocess_async(
394
- request=request,
395
- tokenizer=tokenizer,
396
- tokenization_kwargs=tok_kwargs,
397
- data_1=t1,
398
- data_2=t2,
399
- )
400
- for t1, t2 in input_pairs
401
- )
402
- )
403
-
404
- # Schedule the request and get the result generator.
405
- generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
406
-
407
- default_pooling_params = request.to_pooling_params("classify")
408
-
409
- for i, (full_prompt, tok_prompt) in enumerate(preprocessed_prompts):
410
- request_id_item = f"{request_id}-{i}"
411
-
412
- self._log_inputs(
413
- request_id_item,
414
- full_prompt,
415
- params=default_pooling_params,
416
- lora_request=lora_request,
417
- )
418
-
419
- if token_type_ids := tok_prompt.pop("token_type_ids", None):
420
- pooling_params = default_pooling_params.clone()
421
- compressed = compress_token_type_ids(token_type_ids)
422
- pooling_params.extra_kwargs = {"compressed_token_type_ids": compressed}
423
- else:
424
- pooling_params = default_pooling_params
425
-
426
- generator = self.engine_client.encode(
427
- tok_prompt,
428
- pooling_params,
429
- request_id_item,
430
- lora_request=lora_request,
431
- trace_headers=trace_headers,
432
- priority=request.priority,
433
- )
434
-
435
- generators.append(generator)
436
-
437
- result_generator = merge_async_iterators(*generators)
438
-
439
- # Non-streaming response
440
- final_res_batch: list[PoolingRequestOutput | None] = [None] * len(
441
- preprocessed_prompts
442
- )
443
-
444
- async for i, res in result_generator:
445
- final_res_batch[i] = res
446
-
447
- return [out for out in final_res_batch if out is not None]
448
-
449
- def _preprocess_score(
450
- self,
451
- request: RerankRequest | ScoreRequest,
452
- tokenizer: TokenizerLike,
453
- tokenization_kwargs: dict[str, Any],
454
- data_1: ScoreData,
455
- data_2: ScoreData,
456
- ) -> tuple[str, TokensPrompt]:
457
- model_config = self.model_config
458
- full_prompt, engine_input = get_score_prompt(
459
- model_config=model_config,
460
- data_1=data_1,
461
- data_2=data_2,
462
- tokenizer=tokenizer,
463
- tokenization_kwargs=tokenization_kwargs,
464
- score_template=self.score_template,
465
- )
466
- self._validate_input(request, engine_input["prompt_token_ids"], full_prompt)
467
- if request.mm_processor_kwargs is not None:
468
- engine_input["mm_processor_kwargs"] = request.mm_processor_kwargs
469
-
470
- return full_prompt, engine_input
471
-
472
- async def _run_scoring(
473
- self,
474
- data_1: ScoreInputs,
475
- data_2: ScoreInputs,
476
- request: ScoreRequest | RerankRequest,
477
- request_id: str,
478
- raw_request: Request | None = None,
479
- ) -> list[PoolingRequestOutput] | ErrorResponse:
480
- lora_request = self._maybe_get_adapters(request)
481
-
482
- trace_headers = (
483
- None
484
- if raw_request is None
485
- else await self._get_trace_headers(raw_request.headers)
486
- )
487
-
488
- score_data_1, score_data_2 = validate_score_input(
489
- data_1,
490
- data_2,
491
- is_multimodal_model=self.is_multimodal_model,
492
- architecture=self.architecture,
493
- )
494
-
495
- return await self._score_func(
496
- data_1=score_data_1,
497
- data_2=score_data_2,
498
- request=request,
499
- request_id=request_id,
500
- lora_request=lora_request,
501
- trace_headers=trace_headers,
502
- )
503
-
504
- async def create_score(
505
- self,
506
- request: ScoreRequest,
507
- raw_request: Request | None = None,
508
- ) -> ScoreResponse | ErrorResponse:
509
- """
510
- Score API similar to Sentence Transformers cross encoder
511
-
512
- See https://sbert.net/docs/package_reference/cross_encoder
513
- """
514
- error_check_ret = await self._check_model(request)
515
- if error_check_ret is not None:
516
- return error_check_ret
517
-
518
- request_id = f"score-{self._base_request_id(raw_request)}"
519
- created_time = int(time.time())
520
-
521
- try:
522
- final_res_batch = await self._run_scoring(
523
- request.data_1,
524
- request.data_2,
525
- request,
526
- request_id,
527
- raw_request,
528
- )
529
- if isinstance(final_res_batch, ErrorResponse):
530
- return final_res_batch
531
-
532
- return self.request_output_to_score_response(
533
- final_res_batch,
534
- request_id,
535
- created_time,
536
- self.models.model_name(),
537
- )
538
- except asyncio.CancelledError:
539
- return self.create_error_response("Client disconnected")
540
-
541
- async def do_rerank(
542
- self, request: RerankRequest, raw_request: Request | None = None
543
- ) -> RerankResponse | ErrorResponse:
544
- """
545
- Rerank API based on JinaAI's rerank API; implements the same
546
- API interface. Designed for compatibility with off-the-shelf
547
- tooling, since this is a common standard for reranking APIs
548
-
549
- See example client implementations at
550
- https://github.com/infiniflow/ragflow/blob/main/rag/llm/rerank_model.py
551
- numerous clients use this standard.
552
- """
553
- error_check_ret = await self._check_model(request)
554
- if error_check_ret is not None:
555
- return error_check_ret
556
-
557
- request_id = f"rerank-{self._base_request_id(raw_request)}"
558
- documents = request.documents
559
-
560
- try:
561
- final_res_batch = await self._run_scoring(
562
- request.query,
563
- documents,
564
- request,
565
- request_id,
566
- raw_request,
567
- )
568
- if isinstance(final_res_batch, ErrorResponse):
569
- return final_res_batch
570
-
571
- top_n = request.top_n if request.top_n > 0 else len(final_res_batch)
572
-
573
- return self.request_output_to_rerank_response(
574
- final_res_batch,
575
- request_id,
576
- self.models.model_name(),
577
- documents,
578
- top_n,
579
- )
580
- except asyncio.CancelledError:
581
- return self.create_error_response("Client disconnected")
582
-
583
- def request_output_to_score_response(
584
- self,
585
- final_res_batch: list[PoolingRequestOutput],
586
- request_id: str,
587
- created_time: int,
588
- model_name: str,
589
- ) -> ScoreResponse:
590
- items: list[ScoreResponseData] = []
591
- num_prompt_tokens = 0
592
-
593
- for idx, final_res in enumerate(final_res_batch):
594
- classify_res = ScoringRequestOutput.from_base(final_res)
595
-
596
- item = ScoreResponseData(
597
- index=idx,
598
- score=classify_res.outputs.score,
599
- )
600
- prompt_token_ids = final_res.prompt_token_ids
601
-
602
- items.append(item)
603
- num_prompt_tokens += len(prompt_token_ids)
604
-
605
- usage = UsageInfo(
606
- prompt_tokens=num_prompt_tokens,
607
- total_tokens=num_prompt_tokens,
608
- )
609
-
610
- return ScoreResponse(
611
- id=request_id,
612
- created=created_time,
613
- model=model_name,
614
- data=items,
615
- usage=usage,
616
- )
617
-
618
- def request_output_to_rerank_response(
619
- self,
620
- final_res_batch: list[PoolingRequestOutput],
621
- request_id: str,
622
- model_name: str,
623
- documents: ScoreInputs,
624
- top_n: int,
625
- ) -> RerankResponse:
626
- """
627
- Convert the output of do_rank to a RerankResponse
628
- """
629
-
630
- if not isinstance(documents, list):
631
- documents = [documents]
632
-
633
- results: list[RerankResult] = []
634
- num_prompt_tokens = 0
635
- for idx, final_res in enumerate(final_res_batch):
636
- classify_res = ScoringRequestOutput.from_base(final_res)
637
-
638
- document = documents[idx]
639
- if isinstance(document, str):
640
- rerank_document = RerankDocument(text=document)
641
- else:
642
- rerank_document = RerankDocument(
643
- multi_modal=document.get("content", [])
644
- )
645
-
646
- result = RerankResult(
647
- index=idx,
648
- document=rerank_document,
649
- relevance_score=classify_res.outputs.score,
650
- )
651
- results.append(result)
652
- prompt_token_ids = final_res.prompt_token_ids
653
- num_prompt_tokens += len(prompt_token_ids)
654
-
655
- # sort by relevance, then return the top n if set
656
- results.sort(key=lambda x: x.relevance_score, reverse=True)
657
- if top_n < len(documents):
658
- results = results[:top_n]
659
-
660
- return RerankResponse(
661
- id=request_id,
662
- model=model_name,
663
- results=results,
664
- usage=RerankUsage(
665
- total_tokens=num_prompt_tokens, prompt_tokens=num_prompt_tokens
666
- ),
667
- )
vllm/entrypoints/pooling/scoring/__init__.py
renamed
vllm/entrypoints/pooling/scoring/api_router.py
renamed
+5
-29
@@ -3,21 +3,15 @@
3
from http import HTTPStatus
4
5
from fastapi import APIRouter, Depends, Request
6
-from fastapi.responses import JSONResponse
7
-from typing_extensions import assert_never
6
7
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
8
from vllm.entrypoints.openai.utils import validate_json_request
11
-from vllm.entrypoints.pooling.score.protocol import (
12
- RerankRequest,
13
- RerankResponse,
14
- ScoreRequest,
15
- ScoreResponse,
16
-)
17
-from vllm.entrypoints.pooling.score.serving import ServingScores
9
from vllm.entrypoints.utils import load_aware_call, with_cancellation
10
from vllm.logger import init_logger
11
12
+from .protocol import RerankRequest, ScoreRequest
13
+from .serving import ServingScores
14
+
15
router = APIRouter()
16
17
logger = init_logger(__name__)
@@ -46,16 +40,7 @@ async def create_score(request: ScoreRequest, raw_request: Request):
40
if handler is None:
41
raise NotImplementedError("The model does not support Score API")
42
49
- generator = await handler.create_score(request, raw_request)
50
-
51
- if isinstance(generator, ErrorResponse):
52
- return JSONResponse(
53
- content=generator.model_dump(), status_code=generator.error.code
54
- )
55
- elif isinstance(generator, ScoreResponse):
56
- return JSONResponse(content=generator.model_dump())
57
-
58
- assert_never(generator)
43
+ return await handler(request, raw_request)
44
45
46
@router.post(
@@ -92,16 +77,7 @@ async def do_rerank(request: RerankRequest, raw_request: Request):
77
if handler is None:
78
raise NotImplementedError("The model does not support Rerank (Score) API")
79
95
- generator = await handler.do_rerank(request, raw_request)
96
-
97
- if isinstance(generator, ErrorResponse):
98
- return JSONResponse(
99
- content=generator.model_dump(), status_code=generator.error.code
100
- )
101
- elif isinstance(generator, RerankResponse):
102
- return JSONResponse(content=generator.model_dump())
103
-
104
- assert_never(generator)
80
+ return await handler(request, raw_request)
81
82
83
@router.post(
vllm/entrypoints/pooling/scoring/io_processor.py
new
+419
@@ -0,0 +1,419 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import time
4
+from collections.abc import Sequence
5
+from typing import Any, TypeAlias, cast
6
+
7
+import torch.nn.functional as F
8
+
9
+from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt
10
+from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
11
+from vllm.entrypoints.pooling.typing import (
12
+ OfflineInputsContext,
13
+ OfflineOutputsContext,
14
+ PoolingServeContext,
15
+)
16
+from vllm.inputs import EngineInput
17
+from vllm.renderers import TokenizeParams
18
+from vllm.renderers.hf import safe_apply_chat_template
19
+from vllm.tasks import PoolingTask, ScoreType
20
+from vllm.utils.mistral import is_mistral_tokenizer
21
+
22
+from ...chat_utils import ChatTemplateResolutionError
23
+from .protocol import RerankRequest, ScoreRequest, ScoringRequest
24
+from .typing import ScoreData, ScoreInput, ScoringData
25
+from .utils import (
26
+ compress_token_type_ids,
27
+ compute_maxsim_score,
28
+ parse_score_data,
29
+ score_data_to_prompts,
30
+ validate_score_input,
31
+)
32
+
33
+ScoringServeContext: TypeAlias = PoolingServeContext[ScoringRequest]
34
+
35
+
36
+class ScoringIOProcessor(PoolingIOProcessor):
37
+ name: ScoreType
38
+ pooling_task: PoolingTask
39
+
40
+ def __init__(self, *args, **kwargs):
41
+ super().__init__(*args, **kwargs)
42
+
43
+ self.tokenizer = self.renderer.get_tokenizer()
44
+ self.architecture = self.model_config.architecture
45
+ self.is_multimodal_model = self.model_config.is_multimodal_model
46
+ self.pad_token_id = self.tokenizer.pad_token_id
47
+
48
+ def create_pooling_params(self, request):
49
+ return request.to_pooling_params(self.pooling_task)
50
+
51
+ def valid_inputs(
52
+ self,
53
+ data_1: ScoreInput | list[ScoreInput],
54
+ data_2: ScoreInput | list[ScoreInput],
55
+ ) -> ScoringData:
56
+ scoring_data = validate_score_input(
57
+ data_1,
58
+ data_2,
59
+ is_multimodal_model=self.is_multimodal_model,
60
+ architecture=self.architecture,
61
+ )
62
+ return scoring_data
63
+
64
+
65
+class BiEncoderIOProcessor(ScoringIOProcessor):
66
+ name: ScoreType = "bi-encoder"
67
+ pooling_task: PoolingTask = "embed"
68
+
69
+ #######################################
70
+ # online APIs
71
+
72
+ def pre_process_online(self, ctx: ScoringServeContext):
73
+ request = ctx.request
74
+
75
+ if isinstance(request, ScoreRequest):
76
+ data_1 = request.data_1
77
+ data_2 = request.data_2
78
+ elif isinstance(request, RerankRequest):
79
+ data_1 = request.query
80
+ data_2 = request.documents
81
+ else:
82
+ raise ValueError(f"Invalid {self.name} request type")
83
+
84
+ scoring_data = self.valid_inputs(data_1, data_2)
85
+ tok_params = request.build_tok_params(self.model_config)
86
+ engine_inputs = self._pre_process(
87
+ scoring_data,
88
+ tok_params,
89
+ prompt_extras={
90
+ k: v
91
+ for k in ("mm_processor_kwargs", "cache_salt")
92
+ if (v := getattr(request, k, None)) is not None
93
+ },
94
+ )
95
+
96
+ ctx.engine_inputs = engine_inputs
97
+ ctx.intermediates = len(scoring_data.data_1)
98
+
99
+ def post_process_online(
100
+ self,
101
+ ctx: ScoringServeContext,
102
+ ):
103
+ if ctx.final_res_batch is None:
104
+ raise ValueError("Final response batch not available")
105
+
106
+ if ctx.intermediates is None:
107
+ raise ValueError("data_1 len not available")
108
+
109
+ ctx.final_res_batch = self._post_process(
110
+ outputs=ctx.final_res_batch, offset=cast(int, ctx.intermediates)
111
+ )
112
+
113
+ #######################################
114
+ # offline APIs
115
+
116
+ def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:
117
+ assert isinstance(ctx.prompts, ScoringData)
118
+ tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(
119
+ **(ctx.tokenization_kwargs or {})
120
+ )
121
+ return self._pre_process(ctx.prompts, tok_params)
122
+
123
+ def post_process_offline(
124
+ self,
125
+ ctx: OfflineOutputsContext,
126
+ ) -> list[PoolingRequestOutput]:
127
+ assert ctx.offset is not None
128
+ return self._post_process(outputs=ctx.outputs, offset=ctx.offset)
129
+
130
+ #######################################
131
+ # helpers
132
+
133
+ def _pre_process(
134
+ self,
135
+ scoring_data: ScoringData,
136
+ tok_params: TokenizeParams,
137
+ prompt_extras: dict[str, Any] | None = None,
138
+ ) -> Sequence[EngineInput]:
139
+ data_1 = score_data_to_prompts(scoring_data.data_1, "query", self.model_config)
140
+ data_2 = score_data_to_prompts(
141
+ scoring_data.data_2, "document", self.model_config
142
+ )
143
+
144
+ return self._preprocess_completion_offline(
145
+ prompts=data_1 + data_2, tok_params=tok_params, prompt_extras=prompt_extras
146
+ )
147
+
148
+ def _post_process(self, outputs: list[PoolingRequestOutput], offset: int):
149
+ emb_data_1 = outputs[:offset]
150
+ emb_data_2 = outputs[offset:]
151
+
152
+ if len(emb_data_1) == 1:
153
+ emb_data_1 = emb_data_1 * len(emb_data_2)
154
+
155
+ final_res_batch: list[PoolingRequestOutput] = []
156
+ for emb_1, emb_2 in zip(emb_data_1, emb_data_2):
157
+ pair_score = F.cosine_similarity(
158
+ emb_1.outputs.data.float(), emb_2.outputs.data.float(), dim=0
159
+ )
160
+
161
+ padding: list[int] = []
162
+ if self.pad_token_id is not None:
163
+ padding = [self.pad_token_id]
164
+
165
+ tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids
166
+
167
+ final_res_batch.append(
168
+ PoolingRequestOutput(
169
+ request_id=f"{emb_1.request_id}_{emb_2.request_id}",
170
+ outputs=pair_score,
171
+ prompt_token_ids=tokens,
172
+ num_cached_tokens=emb_1.num_cached_tokens + emb_2.num_cached_tokens,
173
+ finished=True,
174
+ )
175
+ )
176
+ return final_res_batch
177
+
178
+
179
+class LateInteractionIOProcessor(BiEncoderIOProcessor):
180
+ name: ScoreType = "late-interaction"
181
+ pooling_task: PoolingTask = "token_embed"
182
+
183
+ def _post_process(self, outputs: list[PoolingRequestOutput], offset: int):
184
+ # Split into query and document embeddings
185
+ emb_data_1 = outputs[:offset]
186
+ emb_data_2 = outputs[offset:]
187
+
188
+ # Expand queries if 1:N scoring
189
+ if len(emb_data_1) == 1:
190
+ emb_data_1 = emb_data_1 * len(emb_data_2)
191
+
192
+ final_res_batch: list[PoolingRequestOutput] = []
193
+ padding: list[int] = []
194
+ if (pad_token_id := self.pad_token_id) is not None:
195
+ padding = [pad_token_id]
196
+
197
+ # Compute MaxSim scores
198
+ for emb_1, emb_2 in zip(emb_data_1, emb_data_2):
199
+ # emb_1.outputs.data: [query_len, dim]
200
+ # emb_2.outputs.data: [doc_len, dim]
201
+ q_emb = emb_1.outputs.data
202
+ d_emb = emb_2.outputs.data
203
+
204
+ maxsim_score = compute_maxsim_score(q_emb, d_emb)
205
+
206
+ tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids
207
+
208
+ final_res_batch.append(
209
+ PoolingRequestOutput(
210
+ request_id=f"{emb_1.request_id}_{emb_2.request_id}",
211
+ outputs=maxsim_score,
212
+ prompt_token_ids=tokens,
213
+ num_cached_tokens=emb_1.num_cached_tokens + emb_2.num_cached_tokens,
214
+ finished=True,
215
+ )
216
+ )
217
+ return final_res_batch
218
+
219
+
220
+class CrossEncoderIOProcessor(ScoringIOProcessor):
221
+ name: ScoreType = "cross-encoder"
222
+ pooling_task: PoolingTask = "classify"
223
+
224
+ def __init__(self, *args, **kwargs):
225
+ super().__init__(*args, **kwargs)
226
+
227
+ if is_mistral_tokenizer(self.tokenizer):
228
+ raise ValueError("MistralTokenizer not supported for cross-encoding")
229
+
230
+ from vllm.model_executor.model_loader import get_model_cls
231
+ from vllm.model_executor.models.interfaces import supports_score_template
232
+
233
+ model = get_model_cls(self.model_config)
234
+ self.supports_score_template = supports_score_template(model)
235
+ self.model = model if self.supports_score_template else None
236
+ self.use_sep_token = self.model_config.use_sep_token
237
+
238
+ #######################################
239
+ # online APIs
240
+
241
+ def pre_process_online(self, ctx: ScoringServeContext):
242
+ request = ctx.request
243
+
244
+ if isinstance(request, ScoreRequest):
245
+ data_1 = request.data_1
246
+ data_2 = request.data_2
247
+ elif isinstance(request, RerankRequest):
248
+ data_1 = request.query
249
+ data_2 = request.documents
250
+ else:
251
+ raise ValueError(f"Invalid {self.name} request type")
252
+
253
+ scoring_data = self.valid_inputs(data_1, data_2)
254
+ tok_params = request.build_tok_params(self.model_config)
255
+ pooling_params = self.create_pooling_params(request)
256
+
257
+ engine_inputs, pooling_params_list = self._pre_process(
258
+ scoring_data,
259
+ tok_params,
260
+ pooling_params,
261
+ chat_template=self.chat_template,
262
+ prompt_extras={
263
+ k: v
264
+ for k in ("mm_processor_kwargs", "cache_salt")
265
+ if (v := getattr(request, k, None)) is not None
266
+ },
267
+ )
268
+
269
+ ctx.engine_inputs = engine_inputs
270
+ ctx.pooling_params = pooling_params_list
271
+
272
+ #######################################
273
+ # offline APIs
274
+
275
+ def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:
276
+ assert isinstance(ctx.prompts, ScoringData)
277
+ assert not isinstance(ctx.pooling_params, list)
278
+
279
+ tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(
280
+ **(ctx.tokenization_kwargs or {})
281
+ )
282
+ engine_inputs, pooling_params_list = self._pre_process(
283
+ ctx.prompts, tok_params, ctx.pooling_params, ctx.chat_template
284
+ )
285
+ ctx.pooling_params = pooling_params_list
286
+ return engine_inputs
287
+
288
+ #######################################
289
+ # helpers
290
+
291
+ def _pre_process(
292
+ self,
293
+ scoring_data: ScoringData,
294
+ tok_params: TokenizeParams,
295
+ pooling_params: PoolingParams | None,
296
+ chat_template: str | None = None,
297
+ prompt_extras: dict[str, Any] | None = None,
298
+ ) -> tuple[Sequence[EngineInput], list[PoolingParams]]:
299
+ # todo: support prompt_extras
300
+ arrival_time = time.time()
301
+
302
+ data_1 = scoring_data.data_1
303
+ data_2 = scoring_data.data_2
304
+
305
+ if len(data_1) == 1:
306
+ data_1 = data_1 * len(data_2)
307
+
308
+ if pooling_params is None:
309
+ pooling_params = PoolingParams(task="classify")
310
+
311
+ pooling_params_list = list[PoolingParams]()
312
+ engine_inputs = list[EngineInput]()
313
+ for q, d in zip(data_1, data_2):
314
+ _, engine_prompt = self.get_score_prompt(
315
+ data_1=q,
316
+ data_2=d,
317
+ encode_kwargs=tok_params.get_encode_kwargs(),
318
+ chat_template=chat_template,
319
+ )
320
+
321
+ if token_type_ids := engine_prompt.pop("token_type_ids", None):
322
+ params = pooling_params.clone()
323
+ compressed = compress_token_type_ids(token_type_ids)
324
+ params.extra_kwargs = {"compressed_token_type_ids": compressed}
325
+ pooling_params_list.append(params)
326
+ else:
327
+ pooling_params_list.append(pooling_params)
328
+
329
+ tok_params.apply_post_tokenization(self.tokenizer, engine_prompt)
330
+ engine_inputs.append(
331
+ self.renderer.process_for_engine(engine_prompt, arrival_time)
332
+ )
333
+ return engine_inputs, pooling_params_list
334
+
335
+ def get_score_prompt(
336
+ self,
337
+ data_1: ScoreData,
338
+ data_2: ScoreData,
339
+ encode_kwargs: dict[str, Any],
340
+ chat_template: str | None = None,
341
+ ):
342
+ model_config = self.model_config
343
+ tokenizer = self.tokenizer
344
+
345
+ prompt_1, prompt_2, mm_data, mm_uuids = parse_score_data(
346
+ data_1,
347
+ data_2,
348
+ model_config,
349
+ )
350
+
351
+ def default_tokenizer_encode():
352
+ if self.supports_score_template:
353
+ assert self.model is not None
354
+ full_prompt = self.model.get_score_template(prompt_1, prompt_2)
355
+ if full_prompt is None:
356
+ raise ValueError("Get empty score template from model")
357
+
358
+ prompt_inputs = tokenizer(full_prompt, **encode_kwargs)
359
+ else:
360
+ if self.use_sep_token:
361
+ # cross_encoder models defaults to using separating token.
362
+ prompt_inputs = tokenizer(
363
+ text=prompt_1, text_pair=prompt_2, **encode_kwargs
364
+ )
365
+ full_prompt = tokenizer.decode(prompt_inputs["input_ids"])
366
+ else:
367
+ # `llm as reranker` defaults to not using separating token.
368
+ full_prompt = prompt_1 + prompt_2
369
+ prompt_inputs = tokenizer(text=full_prompt, **encode_kwargs)
370
+ return full_prompt, prompt_inputs
371
+
372
+ # FIXME: For now, we only apply a template when one is explicitly provided.
373
+ # We cannot rely on the tokenizer's chat template because many models
374
+ # inherit junk templates from their base LLM, which breaks both the models
375
+ # and the tests that use them.
376
+ if chat_template is None:
377
+ full_prompt, prompt_inputs = default_tokenizer_encode()
378
+ else:
379
+ # FIXME:
380
+ # Try applying a score template from the CLI arg or tokenizer_config.json
381
+ # If that fails because there is no such template,
382
+ # fall back to the default implementation.
383
+ try:
384
+ full_prompt = safe_apply_chat_template(
385
+ model_config,
386
+ tokenizer,
387
+ [
388
+ {"role": "query", "content": prompt_1},
389
+ {"role": "document", "content": prompt_2},
390
+ ],
391
+ chat_template=chat_template,
392
+ tools=None,
393
+ tokenize=False,
394
+ )
395
+ prompt_inputs = tokenizer(full_prompt, **encode_kwargs)
396
+ except ChatTemplateResolutionError:
397
+ full_prompt, prompt_inputs = default_tokenizer_encode()
398
+
399
+ engine_prompt = TokensPrompt(prompt_token_ids=prompt_inputs["input_ids"])
400
+
401
+ if (token_type_ids := prompt_inputs.get("token_type_ids")) is not None:
402
+ engine_prompt["token_type_ids"] = token_type_ids
403
+
404
+ if self.model is not None:
405
+ self.model.post_process_tokens(engine_prompt)
406
+
407
+ if mm_data is not None:
408
+ engine_prompt["multi_modal_data"] = mm_data
409
+ if mm_uuids is not None:
410
+ engine_prompt["multi_modal_uuids"] = mm_uuids
411
+
412
+ return full_prompt, engine_prompt
413
+
414
+
415
+ScoringIOProcessors: dict[ScoreType, type[ScoringIOProcessor]] = {
416
+ "bi-encoder": BiEncoderIOProcessor,
417
+ "late-interaction": LateInteractionIOProcessor,
418
+ "cross-encoder": CrossEncoderIOProcessor,
419
+}
vllm/entrypoints/pooling/scoring/protocol.py
renamed
+17
-14
@@ -12,15 +12,12 @@ from vllm.entrypoints.pooling.base.protocol import (
12
ClassifyRequestMixin,
13
PoolingBasicRequestMixin,
14
)
15
-from vllm.entrypoints.pooling.score.utils import (
16
- ScoreContentPartParam,
17
- ScoreInput,
18
- ScoreInputs,
19
-)
15
from vllm.renderers import TokenizeParams
16
from vllm.tasks import PoolingTask
17
from vllm.utils import random_uuid
18
19
+from .typing import ScoreContentPartParam, ScoreInput
20
+
21
22
class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin):
23
def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
@@ -43,13 +40,13 @@ class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin):
40
41
42
class ScoreDataRequest(ScoreRequestMixin):
46
- data_1: ScoreInputs
47
- data_2: ScoreInputs
43
+ data_1: ScoreInput | list[ScoreInput]
44
+ data_2: ScoreInput | list[ScoreInput]
45
46
47
class ScoreQueriesDocumentsRequest(ScoreRequestMixin):
51
- queries: ScoreInputs
52
- documents: ScoreInputs
48
+ queries: ScoreInput | list[ScoreInput]
49
+ documents: ScoreInput | list[ScoreInput]
50
51
@property
52
def data_1(self):
@@ -61,8 +58,8 @@ class ScoreQueriesDocumentsRequest(ScoreRequestMixin):
58
59
60
class ScoreQueriesItemsRequest(ScoreRequestMixin):
64
- queries: ScoreInputs
65
- items: ScoreInputs
61
+ queries: ScoreInput | list[ScoreInput]
62
+ items: ScoreInput | list[ScoreInput]
63
64
@property
65
def data_1(self):
@@ -74,8 +71,8 @@ class ScoreQueriesItemsRequest(ScoreRequestMixin):
71
72
73
class ScoreTextRequest(ScoreRequestMixin):
77
- text_1: ScoreInputs
78
- text_2: ScoreInputs
74
+ text_1: ScoreInput | list[ScoreInput]
75
+ text_2: ScoreInput | list[ScoreInput]
76
77
@property
78
def data_1(self):
@@ -96,7 +93,7 @@ ScoreRequest: TypeAlias = (
93
94
class RerankRequest(PoolingBasicRequestMixin, ClassifyRequestMixin):
95
query: ScoreInput
99
- documents: ScoreInputs
96
+ documents: ScoreInput | list[ScoreInput]
97
top_n: int = Field(default_factory=lambda: 0)
98
99
def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
@@ -118,6 +115,9 @@ class RerankRequest(PoolingBasicRequestMixin, ClassifyRequestMixin):
115
)
116
117
118
+ScoringRequest: TypeAlias = ScoreRequest | RerankRequest
119
+
120
+
121
class RerankDocument(BaseModel):
122
text: str | None = None
123
multi_modal: list[ScoreContentPartParam] | None = None
@@ -154,3 +154,6 @@ class ScoreResponse(OpenAIBaseModel):
154
model: str
155
data: list[ScoreResponseData]
156
usage: UsageInfo
157
+
158
+
159
+ScoringResponse: TypeAlias = RerankResponse | ScoreResponse
vllm/entrypoints/pooling/scoring/serving.py
new
+160
@@ -0,0 +1,160 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from fastapi.responses import JSONResponse
5
+
6
+from vllm.config import ModelConfig
7
+from vllm.entrypoints.chat_utils import ChatTemplateConfig
8
+from vllm.entrypoints.openai.engine.protocol import UsageInfo
9
+from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
10
+from vllm.entrypoints.pooling.base.serving import PoolingServing
11
+from vllm.logger import init_logger
12
+from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
13
+from vllm.renderers import BaseRenderer
14
+
15
+from .io_processor import ScoringIOProcessors, ScoringServeContext
16
+from .protocol import (
17
+ RerankDocument,
18
+ RerankRequest,
19
+ RerankResponse,
20
+ RerankResult,
21
+ RerankUsage,
22
+ ScoreRequest,
23
+ ScoreResponse,
24
+ ScoreResponseData,
25
+)
26
+from .typing import ScoreInput
27
+
28
+logger = init_logger(__name__)
29
+
30
+
31
+class ServingScores(PoolingServing):
32
+ request_id_prefix = "score"
33
+
34
+ def init_io_processor(
35
+ self,
36
+ model_config: ModelConfig,
37
+ renderer: BaseRenderer,
38
+ chat_template_config: ChatTemplateConfig,
39
+ ) -> PoolingIOProcessor:
40
+ score_type = model_config.score_type
41
+ assert score_type in ScoringIOProcessors
42
+ processor_cls = ScoringIOProcessors[score_type]
43
+ return processor_cls(
44
+ model_config=model_config,
45
+ renderer=renderer,
46
+ chat_template_config=chat_template_config,
47
+ )
48
+
49
+ async def _build_response(
50
+ self,
51
+ ctx: ScoringServeContext,
52
+ ) -> JSONResponse:
53
+ final_res_batch = ctx.final_res_batch
54
+ request_id = ctx.request_id
55
+ created_time = ctx.created_time
56
+ model_name = self.models.model_name()
57
+
58
+ if isinstance(ctx.request, ScoreRequest):
59
+ return self._request_output_to_score_response(
60
+ final_res_batch,
61
+ request_id,
62
+ created_time,
63
+ model_name,
64
+ )
65
+ elif isinstance(ctx.request, RerankRequest):
66
+ return self._request_output_to_rerank_response(
67
+ final_res_batch,
68
+ request_id,
69
+ model_name,
70
+ ctx.request.documents,
71
+ ctx.request.top_n if ctx.request.top_n > 0 else len(final_res_batch),
72
+ )
73
+ else:
74
+ raise NotImplementedError("")
75
+
76
+ def _request_output_to_score_response(
77
+ self,
78
+ final_res_batch: list[PoolingRequestOutput],
79
+ request_id: str,
80
+ created_time: int,
81
+ model_name: str,
82
+ ) -> JSONResponse:
83
+ items: list[ScoreResponseData] = []
84
+ num_prompt_tokens = 0
85
+
86
+ for idx, final_res in enumerate(final_res_batch):
87
+ classify_res = ScoringRequestOutput.from_base(final_res)
88
+
89
+ item = ScoreResponseData(
90
+ index=idx,
91
+ score=classify_res.outputs.score,
92
+ )
93
+ prompt_token_ids = final_res.prompt_token_ids
94
+
95
+ items.append(item)
96
+ num_prompt_tokens += len(prompt_token_ids)
97
+
98
+ usage = UsageInfo(
99
+ prompt_tokens=num_prompt_tokens,
100
+ total_tokens=num_prompt_tokens,
101
+ )
102
+
103
+ response = ScoreResponse(
104
+ id=request_id,
105
+ created=created_time,
106
+ model=model_name,
107
+ data=items,
108
+ usage=usage,
109
+ )
110
+
111
+ return JSONResponse(content=response.model_dump())
112
+
113
+ def _request_output_to_rerank_response(
114
+ self,
115
+ final_res_batch: list[PoolingRequestOutput],
116
+ request_id: str,
117
+ model_name: str,
118
+ documents: ScoreInput | list[ScoreInput],
119
+ top_n: int,
120
+ ) -> JSONResponse:
121
+ if not isinstance(documents, list):
122
+ documents = [documents]
123
+
124
+ results: list[RerankResult] = []
125
+ num_prompt_tokens = 0
126
+ for idx, final_res in enumerate(final_res_batch):
127
+ classify_res = ScoringRequestOutput.from_base(final_res)
128
+
129
+ document = documents[idx]
130
+ if isinstance(document, str):
131
+ rerank_document = RerankDocument(text=document)
132
+ else:
133
+ rerank_document = RerankDocument(
134
+ multi_modal=document.get("content", [])
135
+ )
136
+
137
+ result = RerankResult(
138
+ index=idx,
139
+ document=rerank_document,
140
+ relevance_score=classify_res.outputs.score,
141
+ )
142
+ results.append(result)
143
+ prompt_token_ids = final_res.prompt_token_ids
144
+ num_prompt_tokens += len(prompt_token_ids)
145
+
146
+ # sort by relevance, then return the top n if set
147
+ results.sort(key=lambda x: x.relevance_score, reverse=True)
148
+ if top_n < len(documents):
149
+ results = results[:top_n]
150
+
151
+ response = RerankResponse(
152
+ id=request_id,
153
+ model=model_name,
154
+ results=results,
155
+ usage=RerankUsage(
156
+ total_tokens=num_prompt_tokens, prompt_tokens=num_prompt_tokens
157
+ ),
158
+ )
159
+
160
+ return JSONResponse(content=response.model_dump())
vllm/entrypoints/pooling/scoring/typing.py
new
+46
@@ -0,0 +1,46 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+from typing import TypeAlias
5
+
6
+from typing_extensions import Required, TypedDict
7
+
8
+from vllm.entrypoints.chat_utils import (
9
+ ChatCompletionContentPartImageEmbedsParam,
10
+ ChatCompletionContentPartImageParam,
11
+ ChatCompletionContentPartTextParam,
12
+ ChatCompletionContentPartVideoParam,
13
+)
14
+
15
+ScoreContentPartParam: TypeAlias = (
16
+ ChatCompletionContentPartImageParam
17
+ | ChatCompletionContentPartImageEmbedsParam
18
+ | ChatCompletionContentPartTextParam
19
+ | ChatCompletionContentPartVideoParam
20
+)
21
+
22
+
23
+class ScoreMultiModalParam(TypedDict, total=False):
24
+ """
25
+ A specialized parameter type for scoring multimodal content
26
+
27
+ The reasons why don't reuse `CustomChatCompletionMessageParam` directly:
28
+ 1. Score tasks don't need the 'role' field (user/assistant/system) that's required in chat completions
29
+ 2. Including chat-specific fields would confuse users about their purpose in scoring
30
+ 3. This is a more focused interface that only exposes what's needed for scoring
31
+ """ # noqa: E501
32
+
33
+ content: Required[list[ScoreContentPartParam]]
34
+ """The multimodal contents"""
35
+
36
+
37
+# Raw input data with content key in ScoreMultiModalParam.
38
+ScoreInput = str | ScoreMultiModalParam
39
+# Score data without content key.
40
+ScoreData = str | list[ScoreContentPartParam]
41
+
42
+
43
+@dataclass
44
+class ScoringData:
45
+ data_1: list[ScoreData]
46
+ data_2: list[ScoreData]
vllm/entrypoints/pooling/scoring/utils.py
renamed
+71
-246
@@ -1,42 +1,27 @@
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
4
+from typing import cast
5
6
import torch
7
-from torch.nn import CosineSimilarity
8
-from typing_extensions import Required, TypedDict
7
8
+from vllm import PromptType, TextPrompt
9
from vllm.config import ModelConfig
10
from vllm.entrypoints.chat_utils import (
11
BaseMultiModalItemTracker,
13
- ChatCompletionContentPartImageEmbedsParam,
14
- ChatCompletionContentPartImageParam,
12
ChatCompletionContentPartParam,
13
ChatCompletionContentPartTextParam,
17
- ChatCompletionContentPartVideoParam,
18
- ChatTemplateResolutionError,
14
ConversationMessage,
15
MultiModalItemTracker,
16
_parse_chat_message_content_parts,
17
)
23
-from vllm.inputs import (
24
- MultiModalDataDict,
25
- MultiModalUUIDDict,
26
- PromptType,
27
- TextPrompt,
28
- TokensPrompt,
29
-)
30
-from vllm.model_executor.models.interfaces import supports_score_template
31
-from vllm.outputs import PoolingRequestOutput
32
-from vllm.renderers.hf import safe_apply_chat_template
33
-from vllm.tokenizers import TokenizerLike
34
-
35
-ScoreContentPartParam: TypeAlias = (
36
- ChatCompletionContentPartImageParam
37
- | ChatCompletionContentPartImageEmbedsParam
38
- | ChatCompletionContentPartTextParam
39
- | ChatCompletionContentPartVideoParam
18
+from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict
19
+
20
+from .typing import (
21
+ ScoreContentPartParam,
22
+ ScoreData,
23
+ ScoreInput,
24
+ ScoringData,
25
)
26
27
@@ -57,72 +42,6 @@ def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tens
42
return token_scores.amax(dim=-1).sum()
43
44
60
-class ScoreMultiModalParam(TypedDict, total=False):
61
- """
62
- A specialized parameter type for scoring multimodal content
63
-
64
- The reasons why don't reuse `CustomChatCompletionMessageParam` directly:
65
- 1. Score tasks don't need the 'role' field (user/assistant/system) that's required in chat completions
66
- 2. Including chat-specific fields would confuse users about their purpose in scoring
67
- 3. This is a more focused interface that only exposes what's needed for scoring
68
- """ # noqa: E501
69
-
70
- content: Required[list[ScoreContentPartParam]]
71
- """The multimodal contents"""
72
-
73
-
74
-# Raw input data with content key in ScoreMultiModalParam.
75
-ScoreInput = str | ScoreMultiModalParam
76
-ScoreInputs = ScoreInput | list[ScoreInput]
77
-# Score data without content key.
78
-ScoreData = str | list[ScoreContentPartParam]
79
-
80
-
81
-def _cosine_similarity(
82
- tokenizer: TokenizerLike,
83
- embed_1: list[PoolingRequestOutput],
84
- embed_2: list[PoolingRequestOutput],
85
-) -> list[PoolingRequestOutput]:
86
- scorer = CosineSimilarity(0)
87
- scores: list[PoolingRequestOutput] = []
88
-
89
- for emb_1, emb_2 in zip(embed_1, embed_2):
90
- pair_score = scorer(emb_1.outputs.data, emb_2.outputs.data)
91
-
92
- padding: list[int] = []
93
- if (pad_token_id := tokenizer.pad_token_id) is not None:
94
- padding = [pad_token_id]
95
-
96
- tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids
97
-
98
- scores.append(
99
- PoolingRequestOutput(
100
- request_id=f"{emb_1.request_id}_{emb_2.request_id}",
101
- outputs=pair_score,
102
- prompt_token_ids=tokens,
103
- num_cached_tokens=emb_1.num_cached_tokens + emb_2.num_cached_tokens,
104
- finished=True,
105
- )
106
- )
107
-
108
- return scores
109
-
110
-
111
-def _validate_score_input_lens(
112
- data_1: list[ScoreData],
113
- data_2: list[ScoreData],
114
-):
115
- len_1 = len(data_1)
116
- len_2 = len(data_2)
117
-
118
- if len_1 > 1 and len_1 != len_2:
119
- raise ValueError("Input lengths must be either 1:1, 1:N or N:N")
120
- if len_1 == 0:
121
- raise ValueError("At least one text element must be given")
122
- if len_2 == 0:
123
- raise ValueError("At least one text_pair element must be given")
124
-
125
-
45
def _validate_mm_score_input(
46
data: list[ScoreInput],
47
is_multimodal_model: bool,
@@ -140,12 +59,27 @@ def _validate_mm_score_input(
59
return out
60
61
62
+def _validate_score_input_lens(
63
+ data_1: list[ScoreData],
64
+ data_2: list[ScoreData],
65
+):
66
+ len_1 = len(data_1)
67
+ len_2 = len(data_2)
68
+
69
+ if len_1 > 1 and len_1 != len_2:
70
+ raise ValueError("Input lengths must be either 1:1, 1:N or N:N")
71
+ if len_1 == 0:
72
+ raise ValueError("At least one text element must be given")
73
+ if len_2 == 0:
74
+ raise ValueError("At least one text_pair element must be given")
75
+
76
+
77
def validate_score_input(
144
- data_1: ScoreInputs,
145
- data_2: ScoreInputs,
78
+ data_1: ScoreInput | list[ScoreInput],
79
+ data_2: ScoreInput | list[ScoreInput],
80
is_multimodal_model: bool,
81
architecture: str,
148
-) -> tuple[list[ScoreData], list[ScoreData]]:
82
+) -> ScoringData:
83
if not isinstance(data_1, list):
84
data_1 = [data_1]
85
@@ -155,62 +89,7 @@ def validate_score_input(
89
score_input_1 = _validate_mm_score_input(data_1, is_multimodal_model, architecture)
90
score_input_2 = _validate_mm_score_input(data_2, is_multimodal_model, architecture)
91
_validate_score_input_lens(score_input_1, score_input_2)
158
- return score_input_1, score_input_2
159
-
160
-
161
-def _ensure_str(content: list[ConversationMessage]) -> str:
162
- """Extract a single string prompt from parsed conversation content."""
163
- assert len(content) == 1
164
- prompt = content[0]["content"]
165
- if prompt is not None and isinstance(prompt, str):
166
- return cast(str, prompt)
167
- raise ValueError(f"Only string content is supported, but got {content}.")
168
-
169
-
170
-def parse_score_data(
171
- data_1: ScoreData,
172
- data_2: ScoreData,
173
- model_config: ModelConfig,
174
-) -> tuple[str, str, MultiModalDataDict | None, MultiModalUUIDDict | None]:
175
- """Parse a query-document pair into text prompts and shared multi-modal
176
- data.
177
-
178
- Uses a **single** :class:`MultiModalItemTracker` so that multi-modal
179
- items from both inputs are merged into one ``mm_data`` dict. This is
180
- the correct behaviour for cross-encoder scoring, where query and
181
- document are concatenated into a single model prompt.
182
- """
183
- mm_tracker = MultiModalItemTracker(model_config)
184
-
185
- content_1 = _parse_score_content("query", data_1, mm_tracker)
186
- content_2 = _parse_score_content("document", data_2, mm_tracker)
187
-
188
- prompt_1 = _ensure_str(content_1)
189
- prompt_2 = _ensure_str(content_2)
190
- mm_items, mm_uuids = mm_tracker.resolve_items()
191
-
192
- return prompt_1, prompt_2, mm_items, mm_uuids
193
-
194
-
195
-def parse_score_data_single(
196
- data: ScoreData,
197
- role: str,
198
- model_config: ModelConfig,
199
-) -> tuple[str, MultiModalDataDict | None, MultiModalUUIDDict | None]:
200
- """Parse **one** ScoreData into a text prompt and its own multi-modal
201
- data.
202
-
203
- Unlike :func:`parse_score_data`, each call creates an **independent**
204
- :class:`MultiModalItemTracker` so multi-modal items are kept separate.
205
- This is the correct behaviour for late-interaction scoring, where
206
- query and document are encoded independently.
207
- """
208
- mm_tracker = MultiModalItemTracker(model_config)
209
- content = _parse_score_content(role, data, mm_tracker)
210
-
211
- prompt = _ensure_str(content)
212
- mm_items, mm_uuids = mm_tracker.resolve_items()
213
- return prompt, mm_items, mm_uuids
92
+ return ScoringData(data_1=score_input_1, data_2=score_input_2)
93
94
95
def score_data_to_prompts(
@@ -243,6 +122,15 @@ def score_data_to_prompts(
122
return prompts
123
124
125
+def _ensure_str(content: list[ConversationMessage]) -> str:
126
+ """Extract a single string prompt from parsed conversation content."""
127
+ assert len(content) == 1
128
+ prompt = content[0]["content"]
129
+ if prompt is not None and isinstance(prompt, str):
130
+ return cast(str, prompt)
131
+ raise ValueError(f"Only string content is supported, but got {content}.")
132
+
133
+
134
def _parse_score_content(
135
role: str,
136
data: ScoreData,
@@ -278,113 +166,50 @@ def _parse_score_content(
166
return next(iter(mm_placeholder_storage.values()))[0]
167
168
281
-def _apply_model_score_template(
282
- model_config: ModelConfig, prompt_1: str, prompt_2: str
283
-) -> str:
284
- # NOTE(Simon): lazy import to avoid bring in all dependencies (e.g. gguf)
285
- from vllm.model_executor.model_loader import get_model_cls
286
-
287
- model = get_model_cls(model_config)
288
- if supports_score_template(model):
289
- full_prompt = model.get_score_template(prompt_1, prompt_2)
290
- if full_prompt is None:
291
- raise ValueError("Get empty score template from model")
292
- return full_prompt
293
-
294
- raise ValueError(f"Unsupported model architecture: {model_config.architecture}")
295
-
296
-
297
-def post_process_tokens(
169
+def parse_score_data_single(
170
+ data: ScoreData,
171
+ role: str,
172
model_config: ModelConfig,
299
- prompt: TokensPrompt,
300
-) -> None:
301
- """
302
- Perform architecture-specific manipulations on the input tokens.
173
+) -> tuple[str, MultiModalDataDict | None, MultiModalUUIDDict | None]:
174
+ """Parse **one** ScoreData into a text prompt and its own multi-modal
175
+ data.
176
304
- Note:
305
- This is an in-place operation.
177
+ Unlike :func:`parse_score_data`, each call creates an **independent**
178
+ :class:`MultiModalItemTracker` so multi-modal items are kept separate.
179
+ This is the correct behaviour for late-interaction scoring, where
180
+ query and document are encoded independently.
181
"""
307
- # NOTE(Simon): lazy import to avoid bring in all dependencies (e.g. gguf)
308
- from vllm.model_executor.model_loader import get_model_cls
182
+ mm_tracker = MultiModalItemTracker(model_config)
183
+ content = _parse_score_content(role, data, mm_tracker)
184
310
- model = get_model_cls(model_config)
311
- if supports_score_template(model):
312
- model.post_process_tokens(prompt)
185
+ prompt = _ensure_str(content)
186
+ mm_items, mm_uuids = mm_tracker.resolve_items()
187
+ return prompt, mm_items, mm_uuids
188
189
315
-def get_score_prompt(
316
- model_config: ModelConfig,
317
- tokenizer: TokenizerLike,
318
- tokenization_kwargs: dict[str, Any],
190
+def parse_score_data(
191
data_1: ScoreData,
192
data_2: ScoreData,
321
- score_template: str | None = None,
322
-) -> tuple[str, TokensPrompt]:
323
- prompt_1, prompt_2, mm_data, mm_uuids = parse_score_data(
324
- data_1,
325
- data_2,
326
- model_config,
327
- )
328
- from vllm.model_executor.model_loader import get_model_cls
193
+ model_config: ModelConfig,
194
+) -> tuple[str, str, MultiModalDataDict | None, MultiModalUUIDDict | None]:
195
+ """Parse a query-document pair into text prompts and shared multi-modal
196
+ data.
197
330
- model = get_model_cls(model_config)
198
+ Uses a **single** :class:`MultiModalItemTracker` so that multi-modal
199
+ items from both inputs are merged into one ``mm_data`` dict. This is
200
+ the correct behaviour for cross-encoder scoring, where query and
201
+ document are concatenated into a single model prompt.
202
+ """
203
+ mm_tracker = MultiModalItemTracker(model_config)
204
332
- def default_tokenizer_encode():
333
- if supports_score_template(model):
334
- full_prompt = _apply_model_score_template(model_config, prompt_1, prompt_2)
335
- prompt_inputs = tokenizer(full_prompt, **tokenization_kwargs)
336
- else:
337
- if model_config.use_sep_token:
338
- # cross_encoder models defaults to using separating token.
339
- prompt_inputs = tokenizer(
340
- text=prompt_1, text_pair=prompt_2, **tokenization_kwargs
341
- )
342
- full_prompt = tokenizer.decode(prompt_inputs["input_ids"])
343
- else:
344
- # `llm as reranker` defaults to not using separating token.
345
- full_prompt = prompt_1 + prompt_2
346
- prompt_inputs = tokenizer(text=full_prompt, **tokenization_kwargs)
347
- return full_prompt, prompt_inputs
348
-
349
- # FIXME: For now, we only apply a template when one is explicitly provided.
350
- # We cannot rely on the tokenizer's chat template because many models
351
- # inherit junk templates from their base LLM, which breaks both the models
352
- # and the tests that use them.
353
- if score_template is None:
354
- full_prompt, prompt_inputs = default_tokenizer_encode()
355
- else:
356
- # FIXME: Try applying a score template from the CLI arg or tokenizer_config.json
357
- # If that fails because there is no such template,
358
- # fall back to the default implementation.
359
- try:
360
- full_prompt = safe_apply_chat_template(
361
- model_config,
362
- tokenizer,
363
- [
364
- {"role": "query", "content": prompt_1},
365
- {"role": "document", "content": prompt_2},
366
- ],
367
- chat_template=score_template,
368
- tools=None,
369
- tokenize=False,
370
- )
371
- prompt_inputs = tokenizer(full_prompt, **tokenization_kwargs)
372
- except ChatTemplateResolutionError:
373
- full_prompt, prompt_inputs = default_tokenizer_encode()
374
-
375
- engine_prompt = TokensPrompt(prompt_token_ids=prompt_inputs["input_ids"])
376
-
377
- if (token_type_ids := prompt_inputs.get("token_type_ids")) is not None:
378
- engine_prompt["token_type_ids"] = token_type_ids
379
-
380
- post_process_tokens(model_config, engine_prompt)
381
-
382
- if mm_data is not None:
383
- engine_prompt["multi_modal_data"] = mm_data
384
- if mm_uuids is not None:
385
- engine_prompt["multi_modal_uuids"] = mm_uuids
386
-
387
- return full_prompt, engine_prompt
205
+ content_1 = _parse_score_content("query", data_1, mm_tracker)
206
+ content_2 = _parse_score_content("document", data_2, mm_tracker)
207
+
208
+ prompt_1 = _ensure_str(content_1)
209
+ prompt_2 = _ensure_str(content_2)
210
+ mm_items, mm_uuids = mm_tracker.resolve_items()
211
+
212
+ return prompt_1, prompt_2, mm_items, mm_uuids
213
214
215
def compress_token_type_ids(token_type_ids: list[int]) -> int:
vllm/entrypoints/pooling/typing.py
+29
-12
@@ -1,14 +1,14 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import time
4
-from collections.abc import AsyncGenerator
4
+from collections.abc import AsyncGenerator, Sequence
5
from dataclasses import dataclass, field
6
from typing import Any, Generic, TypeAlias, TypeVar
7
8
from fastapi import Request
9
from pydantic import ConfigDict
10
11
-from vllm import PoolingRequestOutput
11
+from vllm import PoolingParams, PoolingRequestOutput, PromptType
12
from vllm.entrypoints.pooling.classify.protocol import (
13
ClassificationChatRequest,
14
ClassificationCompletionRequest,
@@ -23,15 +23,13 @@ from vllm.entrypoints.pooling.embed.protocol import (
23
)
24
from vllm.entrypoints.pooling.pooling.protocol import (
25
IOProcessorRequest,
26
+ PoolingBytesResponse,
27
PoolingChatRequest,
28
PoolingCompletionRequest,
29
PoolingResponse,
30
)
30
-from vllm.entrypoints.pooling.score.protocol import (
31
- RerankRequest,
32
- ScoreRequest,
33
- ScoreResponse,
34
-)
31
+from vllm.entrypoints.pooling.scoring.protocol import ScoringRequest, ScoringResponse
32
+from vllm.entrypoints.pooling.scoring.typing import ScoringData
33
from vllm.inputs import EngineInput
34
from vllm.lora.request import LoRARequest
35
@@ -49,8 +47,7 @@ AnyPoolingRequest: TypeAlias = (
47
PoolingCompletionLikeRequest
48
| PoolingChatLikeRequest
49
| IOProcessorRequest
52
- | RerankRequest
53
- | ScoreRequest
50
+ | ScoringRequest
51
| CohereEmbedRequest
52
)
53
@@ -59,7 +56,8 @@ AnyPoolingResponse: TypeAlias = (
56
| EmbeddingResponse
57
| EmbeddingBytesResponse
58
| PoolingResponse
62
- | ScoreResponse
59
+ | PoolingBytesResponse
60
+ | ScoringResponse
61
)
62
63
PoolingRequestT = TypeVar("PoolingRequestT", bound=AnyPoolingRequest)
@@ -73,8 +71,8 @@ class PoolingServeContext(Generic[PoolingRequestT]):
71
request_id: str
72
created_time: int = field(default_factory=lambda: int(time.time()))
73
lora_request: LoRARequest | None = None
76
-
77
- engine_inputs: list[EngineInput] | None = None
74
+ pooling_params: PoolingParams | list[PoolingParams] | None = None
75
+ engine_inputs: Sequence[EngineInput] | None = None
76
prompt_request_ids: list[str] | None = None
77
intermediates: Any | None = None
78
@@ -84,3 +82,22 @@ class PoolingServeContext(Generic[PoolingRequestT]):
82
final_res_batch: list[PoolingRequestOutput] = field(default_factory=list)
83
84
model_config = ConfigDict(arbitrary_types_allowed=True)
85
+
86
+
87
+@dataclass
88
+class OfflineInputsContext:
89
+ prompts: PromptType | Sequence[PromptType] | ScoringData
90
+ pooling_params: PoolingParams | list[PoolingParams] | None = None
91
+ tokenization_kwargs: dict[str, Any] | None = None
92
+ chat_template: str | None = None
93
+
94
+ ## for bi-encoder & late-interaction
95
+ offset: int | None = None
96
+
97
+
98
+@dataclass
99
+class OfflineOutputsContext:
100
+ outputs: list[PoolingRequestOutput]
101
+
102
+ ## for bi-encoder & late-interaction
103
+ offset: int | None = None
vllm/entrypoints/pooling/utils.py
+19
@@ -11,8 +11,10 @@ import pybase64
11
import torch
12
from fastapi.responses import JSONResponse
13
14
+from vllm.config import ModelConfig
15
from vllm.logger import init_logger
16
from vllm.outputs import PoolingRequestOutput
17
+from vllm.tasks import SupportedTask
18
from vllm.utils.serial_utils import (
19
EMBED_DTYPES,
20
EmbedDType,
@@ -133,3 +135,20 @@ def get_json_response_cls() -> type[JSONResponse]:
135
"To make v1/embeddings API fast, please install orjson by `pip install orjson`"
136
)
137
return JSONResponse
138
+
139
+
140
+def enable_scoring_api(
141
+ supported_tasks: tuple["SupportedTask", ...],
142
+ model_config: ModelConfig | None = None,
143
+) -> bool:
144
+ if any(t in supported_tasks for t in ("embed", "token_embed")):
145
+ return True
146
+
147
+ if model_config is not None and "classify" in supported_tasks:
148
+ num_labels = getattr(model_config.hf_config, "num_labels", 0)
149
+ if num_labels != 1:
150
+ logger.debug_once("Scoring API is only enabled for num_labels == 1.")
151
+ return False
152
+ return True
153
+
154
+ return False
vllm/entrypoints/sagemaker/api_router.py
+5
-5
@@ -14,8 +14,8 @@ from vllm.config import ModelConfig
14
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
15
from vllm.entrypoints.openai.engine.serving import OpenAIServing
16
from vllm.entrypoints.openai.utils import validate_json_request
17
-from vllm.entrypoints.pooling import enable_scoring_api
17
from vllm.entrypoints.pooling.base.serving import PoolingServing
18
+from vllm.entrypoints.pooling.utils import enable_scoring_api
19
from vllm.entrypoints.serve.instrumentator.basic import base
20
from vllm.entrypoints.serve.instrumentator.health import health
21
from vllm.tasks import POOLING_TASKS, SupportedTask
@@ -76,15 +76,15 @@ def get_invocation_types(
76
]
77
78
if enable_scoring_api(supported_tasks, model_config):
79
- from vllm.entrypoints.pooling.score.api_router import do_rerank, rerank
80
- from vllm.entrypoints.pooling.score.protocol import RerankRequest
79
+ from vllm.entrypoints.pooling.scoring.api_router import do_rerank, rerank
80
+ from vllm.entrypoints.pooling.scoring.protocol import RerankRequest
81
82
INVOCATION_TYPES += [
83
(RerankRequest, (rerank, do_rerank)),
84
]
85
86
- from vllm.entrypoints.pooling.score.api_router import create_score, score
87
- from vllm.entrypoints.pooling.score.protocol import ScoreRequest
86
+ from vllm.entrypoints.pooling.scoring.api_router import create_score, score
87
+ from vllm.entrypoints.pooling.scoring.protocol import ScoreRequest
88
89
INVOCATION_TYPES += [
90
(ScoreRequest, (score, create_score)),