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 14, 00:30 UTC
[Frontend][Doc][5/N] Improve all pooling task | Polish encode (pooling) api & Document. (#25524)
Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
wang.yuqi committed
Oct 30, 2025 at 20:13 UTC
4464723f220a74785cd1971cf62a04e3961c2846
27 files changed
+499
-131
docs/design/io_processor_plugins.md
+1
-1
@@ -79,7 +79,7 @@ The `post_process*` methods take `PoolingRequestOutput` objects as input and gen
79
The `validate_or_generate_params` method is used for validating with the plugin any `SamplingParameters`/`PoolingParameters` received with the user request, or to generate new ones if none are specified. The function always returns the validated/generated parameters.
80
The `output_to_response` method is used only for online serving and converts the plugin output to the `IOProcessorResponse` type that is then returned by the API Server. The implementation of the `/pooling` serving endpoint is available here [vllm/entrypoints/openai/serving_pooling.py](../../vllm/entrypoints/openai/serving_pooling.py).
81
82
-An example implementation of a plugin that enables generating geotiff images with the PrithviGeospatialMAE model is available [here](https://github.com/IBM/terratorch/tree/main/terratorch/vllm/plugins/segmentation). Please, also refer to our online ([examples/online_serving/prithvi_geospatial_mae.py](../../examples/online_serving/prithvi_geospatial_mae.py)) and offline ([examples/offline_inference/prithvi_geospatial_mae_io_processor.py](../../examples/offline_inference/prithvi_geospatial_mae_io_processor.py)) inference examples.
82
+An example implementation of a plugin that enables generating geotiff images with the PrithviGeospatialMAE model is available [here](https://github.com/IBM/terratorch/tree/main/terratorch/vllm/plugins/segmentation). Please, also refer to our online ([examples/online_serving/pooling/prithvi_geospatial_mae.py](../../examples/online_serving/pooling/prithvi_geospatial_mae.py)) and offline ([examples/offline_inference/pooling/prithvi_geospatial_mae_io_processor.py](../../examples/offline_inference/pooling/prithvi_geospatial_mae_io_processor.py)) inference examples.
83
84
## Using an IO Processor plugin
85
docs/models/pooling_models.md
+68
-15
@@ -30,11 +30,11 @@ If `--runner pooling` has been set (manually or automatically) but the model doe
30
vLLM will attempt to automatically convert the model according to the architecture names
31
shown in the table below.
32
33
-| Architecture | `--convert` | Supported pooling tasks |
34
-|-------------------------------------------------|-------------|-------------------------------|
35
-| `*ForTextEncoding`, `*EmbeddingModel`, `*Model` | `embed` | `encode`, `embed` |
36
-| `*For*Classification`, `*ClassificationModel` | `classify` | `encode`, `classify`, `score` |
37
-| `*ForRewardModeling`, `*RewardModel` | `reward` | `encode` |
33
+| Architecture | `--convert` | Supported pooling tasks |
34
+|-------------------------------------------------|-------------|---------------------------------------|
35
+| `*ForTextEncoding`, `*EmbeddingModel`, `*Model` | `embed` | `token_embed`, `embed` |
36
+| `*For*Classification`, `*ClassificationModel` | `classify` | `token_classify`, `classify`, `score` |
37
+| `*ForRewardModeling`, `*RewardModel` | `reward` | `token_classify` |
38
39
!!! tip
40
You can explicitly set `--convert <type>` to specify how to convert the model.
@@ -45,12 +45,14 @@ Each pooling model in vLLM supports one or more of these tasks according to
45
[Pooler.get_supported_tasks][vllm.model_executor.layers.pooler.Pooler.get_supported_tasks],
46
enabling the corresponding APIs:
47
48
-| Task | APIs |
49
-|------------|--------------------------------------|
50
-| `encode` | `LLM.reward(...)` |
51
-| `embed` | `LLM.embed(...)`, `LLM.score(...)`\* |
52
-| `classify` | `LLM.classify(...)` |
53
-| `score` | `LLM.score(...)` |
48
+| Task | APIs |
49
+|------------------|-------------------------------------------------------------------------------|
50
+| `embed` | `LLM.embed(...)`, `LLM.score(...)`\*, `LLM.encode(..., pooling_task="embed")` |
51
+| `classify` | `LLM.classify(...)`, `LLM.encode(..., pooling_task="classify")` |
52
+| `score` | `LLM.score(...)` |
53
+| `token_classify` | `LLM.reward(...)`, `LLM.encode(..., pooling_task="token_classify")` |
54
+| `token_embed` | `LLM.encode(..., pooling_task="token_embed")` |
55
+| `plugin` | `LLM.encode(..., pooling_task="plugin")` |
56
57
\* The `LLM.score(...)` API falls back to `embed` task if the model does not support `score` task.
58
@@ -144,7 +146,6 @@ A code example can be found here: [examples/offline_inference/basic/score.py](..
146
### `LLM.reward`
147
148
The [reward][vllm.LLM.reward] method is available to all reward models in vLLM.
147
-It returns the extracted hidden states directly.
149
150
```python
151
from vllm import LLM
@@ -161,15 +162,17 @@ A code example can be found here: [examples/offline_inference/basic/reward.py](.
162
### `LLM.encode`
163
164
The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
164
-It returns the extracted hidden states directly.
165
166
!!! note
167
Please use one of the more specific methods or set the task directly when using `LLM.encode`:
168
169
- For embeddings, use `LLM.embed(...)` or `pooling_task="embed"`.
170
- For classification logits, use `LLM.classify(...)` or `pooling_task="classify"`.
171
- - For rewards, use `LLM.reward(...)` or `pooling_task="reward"`.
171
- For similarity scores, use `LLM.score(...)`.
172
+ - For rewards, use `LLM.reward(...)` or `pooling_task="token_classify"`.
173
+ - For token classification, use `pooling_task="token_classify"`.
174
+ - For multi-vector retrieval, use `pooling_task="token_embed"`
175
+ - For IO Processor Plugins , use `pooling_task="plugin"`
176
177
```python
178
from vllm import LLM
@@ -185,10 +188,47 @@ print(f"Data: {data!r}")
188
189
Our [OpenAI-Compatible Server](../serving/openai_compatible_server.md) provides endpoints that correspond to the offline APIs:
190
188
-- [Pooling API](../serving/openai_compatible_server.md#pooling-api) is similar to `LLM.encode`, being applicable to all types of pooling models.
191
- [Embeddings API](../serving/openai_compatible_server.md#embeddings-api) is similar to `LLM.embed`, accepting both text and [multi-modal inputs](../features/multimodal_inputs.md) for embedding models.
192
- [Classification API](../serving/openai_compatible_server.md#classification-api) is similar to `LLM.classify` and is applicable to sequence classification models.
193
- [Score API](../serving/openai_compatible_server.md#score-api) is similar to `LLM.score` for cross-encoder models.
194
+- [Pooling API](../serving/openai_compatible_server.md#pooling-api) is similar to `LLM.encode`, being applicable to all types of pooling models.
195
+
196
+!!! note
197
+ Please use one of the more specific methods or set the task directly when using [Pooling API](../serving/openai_compatible_server.md#pooling-api) api.:
198
+
199
+ - For embeddings, use [Embeddings API](../serving/openai_compatible_server.md#embeddings-api) or `"task":"embed"`.
200
+ - For classification logits, use [Classification API](../serving/openai_compatible_server.md#classification-api) or `task":"classify"`.
201
+ - For similarity scores, use [Score API](../serving/openai_compatible_server.md#score-api).
202
+ - For rewards, `task":"token_classify"`.
203
+ - For token classification, use `task":"token_classify"`.
204
+ - For multi-vector retrieval, use `task":"token_embed"`
205
+ - For IO Processor Plugins , use `task":"plugin"`
206
+
207
+```python
208
+# start a supported embeddings model server with `vllm serve`, e.g.
209
+# vllm serve intfloat/e5-small
210
+import requests
211
+
212
+host = "localhost"
213
+port = "8000"
214
+model_name = "intfloat/e5-small"
215
+
216
+api_url = f"http://{host}:{port}/pooling"
217
+
218
+prompts = [
219
+ "Hello, my name is",
220
+ "The president of the United States is",
221
+ "The capital of France is",
222
+ "The future of AI is",
223
+]
224
+prompt = {"model": model_name, "input": prompts, "task": "embed"}
225
+
226
+response = requests.post(api_url, json=prompt)
227
+
228
+for output in response.json()["data"]:
229
+ data = output["data"]
230
+ print(f"Data: {data!r} (size={len(data)})")
231
+```
232
233
## Matryoshka Embeddings
234
@@ -265,3 +305,16 @@ Expected output:
305
```
306
307
An OpenAI client example can be found here: [examples/online_serving/pooling/openai_embedding_matryoshka_fy.py](../../examples/online_serving/pooling/openai_embedding_matryoshka_fy.py)
308
+
309
+## Deprecated Features
310
+
311
+### Encode task
312
+
313
+We have split the `encode` task into two more specific token wise tasks: `token_embed` and `token_classify`:
314
+
315
+- `token_embed` is the same as embed, using normalize as activation.
316
+- `token_classify` is the same as classify, default using softmax as activation.
317
+
318
+### Remove softmax from PoolingParams
319
+
320
+We are going to remove `softmax` and `activation` from `PoolingParams`. Instead, you should set `use_activation`, since we actually allow `classify` and `token_classify` to use any activation function.
docs/serving/openai_compatible_server.md
+2
-2
@@ -638,7 +638,7 @@ Usually, the score for a sentence pair refers to the similarity between two sent
638
639
You can find the documentation for cross encoder models at [sbert.net](https://www.sbert.net/docs/package_reference/cross_encoder/cross_encoder.html).
640
641
-Code example: [examples/online_serving/openai_cross_encoder_score.py](../../examples/online_serving/openai_cross_encoder_score.py)
641
+Code example: [examples/online_serving/pooling/openai_cross_encoder_score.py](../../examples/online_serving/pooling/openai_cross_encoder_score.py)
642
643
#### Single inference
644
@@ -819,7 +819,7 @@ You can pass multi-modal inputs to scoring models by passing `content` including
819
print("Scoring output:", response_json["data"][0]["score"])
820
print("Scoring output:", response_json["data"][1]["score"])
821
```
822
-Full example: [examples/online_serving/openai_cross_encoder_score_for_multimodal.py](../../examples/online_serving/openai_cross_encoder_score_for_multimodal.py)
822
+Full example: [examples/online_serving/pooling/openai_cross_encoder_score_for_multimodal.py](../../examples/online_serving/pooling/openai_cross_encoder_score_for_multimodal.py)
823
824
#### Extra parameters
825
examples/offline_inference/pooling/README.md
+12
@@ -38,6 +38,18 @@ python examples/offline_inference/pooling/multi_vector_retrieval.py
38
python examples/offline_inference/pooling/ner.py
39
```
40
41
+## Prithvi Geospatial MAE usage
42
+
43
+```bash
44
+python examples/offline_inference/pooling/prithvi_geospatial_mae.py
45
+```
46
+
47
+## IO Processor Plugins for Prithvi Geospatial MAE
48
+
49
+```bash
50
+python examples/offline_inference/pooling/prithvi_geospatial_mae_io_processor.py
51
+```
52
+
53
## Qwen3 reranker usage
54
55
```bash
examples/offline_inference/pooling/ner.py
+1
-1
@@ -33,7 +33,7 @@ def main(args: Namespace):
33
label_map = llm.llm_engine.vllm_config.model_config.hf_config.id2label
34
35
# Run inference
36
- outputs = llm.encode(prompts)
36
+ outputs = llm.encode(prompts, pooling_task="token_classify")
37
38
for prompt, output in zip(prompts, outputs):
39
logits = output.outputs.data
examples/offline_inference/pooling/prithvi_geospatial_mae.py
renamed
examples/offline_inference/pooling/prithvi_geospatial_mae_io_processor.py
renamed
examples/online_serving/pooling/README.md
+35
-5
@@ -3,65 +3,95 @@
3
## Cohere rerank usage
4
5
```bash
6
+# vllm serve BAAI/bge-reranker-base
7
python examples/online_serving/pooling/cohere_rerank_client.py
8
```
9
10
## Embedding requests base64 encoding_format usage
11
12
```bash
13
+# vllm serve intfloat/e5-small
14
python examples/online_serving/pooling/embedding_requests_base64_client.py
15
```
16
17
## Embedding requests bytes encoding_format usage
18
19
```bash
20
+# vllm serve intfloat/e5-small
21
python examples/online_serving/pooling/embedding_requests_bytes_client.py
22
```
23
24
## Jinaai rerank usage
25
26
```bash
27
+# vllm serve BAAI/bge-reranker-base
28
python examples/online_serving/pooling/jinaai_rerank_client.py
29
```
30
31
## Multi vector retrieval usage
32
33
```bash
34
+# vllm serve BAAI/bge-m3
35
python examples/online_serving/pooling/multi_vector_retrieval_client.py
36
```
37
38
## Named Entity Recognition (NER) usage
39
40
```bash
41
+# vllm serve boltuix/NeuroBERT-NER
42
python examples/online_serving/pooling/ner_client.py
43
```
44
39
-## Openai chat embedding for multimodal usage
45
+## OpenAI chat embedding for multimodal usage
46
47
```bash
48
python examples/online_serving/pooling/openai_chat_embedding_client_for_multimodal.py
49
```
50
45
-## Openai classification usage
51
+## OpenAI classification usage
52
53
```bash
54
+# vllm serve jason9693/Qwen2.5-1.5B-apeach
55
python examples/online_serving/pooling/openai_classification_client.py
56
```
57
51
-## Openai embedding usage
58
+## OpenAI cross_encoder score usage
59
60
```bash
61
+# vllm serve BAAI/bge-reranker-v2-m3
62
+python examples/online_serving/pooling/openai_cross_encoder_score.py
63
+```
64
+
65
+## OpenAI cross_encoder score for multimodal usage
66
+
67
+```bash
68
+# vllm serve jinaai/jina-reranker-m0
69
+python examples/online_serving/pooling/openai_cross_encoder_score_for_multimodal.py
70
+```
71
+
72
+## OpenAI embedding usage
73
+
74
+```bash
75
+# vllm serve intfloat/e5-small
76
python examples/online_serving/pooling/openai_embedding_client.py
77
```
78
57
-## Openai embedding matryoshka dimensions usage
79
+## OpenAI embedding matryoshka dimensions usage
80
81
```bash
82
+# vllm serve jinaai/jina-embeddings-v3 --trust-remote-code
83
python examples/online_serving/pooling/openai_embedding_matryoshka_fy.py
84
```
85
63
-## Openai pooling usage
86
+## OpenAI pooling usage
87
88
```bash
89
+# vllm serve internlm/internlm2-1_8b-reward --trust-remote-code
90
python examples/online_serving/pooling/openai_pooling_client.py
91
```
92
+
93
+## Online Prithvi Geospatial MAE usage
94
+
95
+```bash
96
+python examples/online_serving/pooling/prithvi_geospatial_mae.py
97
+```
examples/online_serving/pooling/openai_cross_encoder_score.py
renamed
examples/online_serving/pooling/openai_cross_encoder_score_for_multimodal.py
renamed
examples/online_serving/pooling/prithvi_geospatial_mae.py
renamed
tests/entrypoints/pooling/llm/test_classify.py
+7
-5
@@ -37,15 +37,17 @@ def llm():
37
38
@pytest.mark.skip_global_cleanup
39
def test_pooling_params(llm: LLM):
40
- def get_outputs(activation):
40
+ def get_outputs(use_activation):
41
outputs = llm.classify(
42
- prompts, pooling_params=PoolingParams(activation=activation), use_tqdm=False
42
+ prompts,
43
+ pooling_params=PoolingParams(use_activation=use_activation),
44
+ use_tqdm=False,
45
)
46
return torch.tensor([x.outputs.probs for x in outputs])
47
46
- default = get_outputs(activation=None)
47
- w_activation = get_outputs(activation=True)
48
- wo_activation = get_outputs(activation=False)
48
+ default = get_outputs(use_activation=None)
49
+ w_activation = get_outputs(use_activation=True)
50
+ wo_activation = get_outputs(use_activation=False)
51
52
assert torch.allclose(default, w_activation, atol=1e-2), (
53
"Default should use activation."
tests/entrypoints/pooling/llm/test_reward.py
+7
-5
@@ -37,15 +37,17 @@ def llm():
37
38
39
def test_pooling_params(llm: LLM):
40
- def get_outputs(activation):
40
+ def get_outputs(use_activation):
41
outputs = llm.reward(
42
- prompts, pooling_params=PoolingParams(activation=activation), use_tqdm=False
42
+ prompts,
43
+ pooling_params=PoolingParams(use_activation=use_activation),
44
+ use_tqdm=False,
45
)
46
return torch.cat([x.outputs.data for x in outputs])
47
46
- default = get_outputs(activation=None)
47
- w_activation = get_outputs(activation=True)
48
- wo_activation = get_outputs(activation=False)
48
+ default = get_outputs(use_activation=None)
49
+ w_activation = get_outputs(use_activation=True)
50
+ wo_activation = get_outputs(use_activation=False)
51
52
assert torch.allclose(default, w_activation, atol=1e-2), (
53
"Default should use activation."
tests/entrypoints/pooling/llm/test_score.py
+5
-5
@@ -34,21 +34,21 @@ def llm():
34
35
36
def test_pooling_params(llm: LLM):
37
- def get_outputs(activation):
37
+ def get_outputs(use_activation):
38
text_1 = "What is the capital of France?"
39
text_2 = "The capital of France is Paris."
40
41
outputs = llm.score(
42
text_1,
43
text_2,
44
- pooling_params=PoolingParams(activation=activation),
44
+ pooling_params=PoolingParams(use_activation=use_activation),
45
use_tqdm=False,
46
)
47
return torch.tensor([x.outputs.score for x in outputs])
48
49
- default = get_outputs(activation=None)
50
- w_activation = get_outputs(activation=True)
51
- wo_activation = get_outputs(activation=False)
49
+ default = get_outputs(use_activation=None)
50
+ w_activation = get_outputs(use_activation=True)
51
+ wo_activation = get_outputs(use_activation=False)
52
53
assert torch.allclose(default, w_activation, atol=1e-2), (
54
"Default should use activation."
tests/entrypoints/pooling/openai/test_classification.py
+72
-20
@@ -7,7 +7,7 @@ import torch
7
import torch.nn.functional as F
8
9
from tests.utils import RemoteOpenAIServer
10
-from vllm.entrypoints.openai.protocol import ClassificationResponse
10
+from vllm.entrypoints.openai.protocol import ClassificationResponse, PoolingResponse
11
12
MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"
13
DTYPE = "float32" # Use float32 to avoid NaN issue
@@ -163,20 +163,24 @@ async def test_invocations(server: RemoteOpenAIServer):
163
164
@pytest.mark.asyncio
165
@pytest.mark.parametrize("model_name", [MODEL_NAME])
166
-async def test_activation(server: RemoteOpenAIServer, model_name: str):
166
+async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
167
input_text = ["This product was excellent and exceeded my expectations"]
168
169
- async def get_outputs(activation):
169
+ async def get_outputs(use_activation):
170
response = requests.post(
171
server.url_for("classify"),
172
- json={"model": model_name, "input": input_text, "activation": activation},
172
+ json={
173
+ "model": model_name,
174
+ "input": input_text,
175
+ "use_activation": use_activation,
176
+ },
177
)
178
outputs = response.json()
179
return torch.tensor([x["probs"] for x in outputs["data"]])
180
177
- default = await get_outputs(activation=None)
178
- w_activation = await get_outputs(activation=True)
179
- wo_activation = await get_outputs(activation=False)
181
+ default = await get_outputs(use_activation=None)
182
+ w_activation = await get_outputs(use_activation=True)
183
+ wo_activation = await get_outputs(use_activation=False)
184
185
assert torch.allclose(default, w_activation, atol=1e-2), (
186
"Default should use activation."
@@ -191,18 +195,7 @@ async def test_activation(server: RemoteOpenAIServer, model_name: str):
195
196
@pytest.mark.asyncio
197
@pytest.mark.parametrize("model_name", [MODEL_NAME])
194
-def test_pooling(server: RemoteOpenAIServer, model_name: str):
195
- # pooling api uses ALL pooling, which does not support chunked prefill.
196
- response = requests.post(
197
- server.url_for("pooling"),
198
- json={"model": model_name, "input": "test", "encoding_format": "float"},
199
- )
200
- assert response.json()["error"]["type"] == "BadRequestError"
201
-
202
-
203
-@pytest.mark.asyncio
204
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
205
-def test_score(server: RemoteOpenAIServer, model_name: str):
198
+async def test_score(server: RemoteOpenAIServer, model_name: str):
199
# score api is only enabled for num_labels == 1.
200
response = requests.post(
201
server.url_for("score"),
@@ -217,7 +210,7 @@ def test_score(server: RemoteOpenAIServer, model_name: str):
210
211
@pytest.mark.asyncio
212
@pytest.mark.parametrize("model_name", [MODEL_NAME])
220
-def test_rerank(server: RemoteOpenAIServer, model_name: str):
213
+async def test_rerank(server: RemoteOpenAIServer, model_name: str):
214
# rerank api is only enabled for num_labels == 1.
215
response = requests.post(
216
server.url_for("rerank"),
@@ -228,3 +221,62 @@ def test_rerank(server: RemoteOpenAIServer, model_name: str):
221
},
222
)
223
assert response.json()["error"]["type"] == "BadRequestError"
224
+
225
+
226
+@pytest.mark.asyncio
227
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
228
+async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
229
+ input_text = "This product was excellent and exceeded my expectations"
230
+ response = requests.post(
231
+ server.url_for("pooling"),
232
+ json={
233
+ "model": model_name,
234
+ "input": input_text,
235
+ "encoding_format": "float",
236
+ "task": "classify",
237
+ },
238
+ )
239
+ poolings = PoolingResponse.model_validate(response.json())
240
+ assert len(poolings.data) == 1
241
+ assert len(poolings.data[0].data) == 2
242
+
243
+
244
+@pytest.mark.asyncio
245
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
246
+async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
247
+ # token_classify uses ALL pooling, which does not support chunked prefill.
248
+ task = "token_classify"
249
+ response = requests.post(
250
+ server.url_for("pooling"),
251
+ json={
252
+ "model": model_name,
253
+ "input": "test",
254
+ "encoding_format": "float",
255
+ "task": task,
256
+ },
257
+ )
258
+ assert response.json()["error"]["type"] == "BadRequestError"
259
+ assert response.json()["error"]["message"].startswith(
260
+ f"Task {task} is not supported"
261
+ )
262
+
263
+
264
+@pytest.mark.asyncio
265
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
266
+@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
267
+async def test_pooling_not_supported(
268
+ server: RemoteOpenAIServer, model_name: str, task: str
269
+):
270
+ response = requests.post(
271
+ server.url_for("pooling"),
272
+ json={
273
+ "model": model_name,
274
+ "input": "test",
275
+ "encoding_format": "float",
276
+ "task": task,
277
+ },
278
+ )
279
+ assert response.json()["error"]["type"] == "BadRequestError"
280
+ assert response.json()["error"]["message"].startswith(
281
+ f"Task {task} is not supported"
282
+ )
tests/entrypoints/pooling/openai/test_embedding.py
+51
-2
@@ -562,12 +562,40 @@ async def test_normalize(server: RemoteOpenAIServer, model_name: str):
562
563
@pytest.mark.asyncio
564
@pytest.mark.parametrize("model_name", [MODEL_NAME])
565
-async def test_pooling(server: RemoteOpenAIServer, model_name: str):
565
+async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
566
+ task = "embed"
567
input_text = ["The chef prepared a delicious meal."]
568
569
response = requests.post(
570
server.url_for("pooling"),
570
- json={"model": model_name, "input": input_text, "encoding_format": "float"},
571
+ json={
572
+ "model": model_name,
573
+ "input": input_text,
574
+ "encoding_format": "float",
575
+ "task": task,
576
+ },
577
+ )
578
+
579
+ poolings = PoolingResponse.model_validate(response.json())
580
+
581
+ assert len(poolings.data) == 1
582
+ assert len(poolings.data[0].data) == 384
583
+
584
+
585
+@pytest.mark.asyncio
586
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
587
+async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
588
+ task = "token_embed"
589
+ input_text = ["The chef prepared a delicious meal."]
590
+
591
+ response = requests.post(
592
+ server.url_for("pooling"),
593
+ json={
594
+ "model": model_name,
595
+ "input": input_text,
596
+ "encoding_format": "float",
597
+ "task": task,
598
+ },
599
)
600
601
poolings = PoolingResponse.model_validate(response.json())
@@ -575,3 +603,24 @@ async def test_pooling(server: RemoteOpenAIServer, model_name: str):
603
assert len(poolings.data) == 1
604
assert len(poolings.data[0].data) == 11
605
assert len(poolings.data[0].data[0]) == 384
606
+
607
+
608
+@pytest.mark.asyncio
609
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
610
+@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
611
+async def test_pooling_not_supported(
612
+ server: RemoteOpenAIServer, model_name: str, task: str
613
+):
614
+ response = requests.post(
615
+ server.url_for("pooling"),
616
+ json={
617
+ "model": model_name,
618
+ "input": "test",
619
+ "encoding_format": "float",
620
+ "task": task,
621
+ },
622
+ )
623
+ assert response.json()["error"]["type"] == "BadRequestError"
624
+ assert response.json()["error"]["message"].startswith(
625
+ f"Task {task} is not supported"
626
+ )
tests/entrypoints/pooling/openai/test_rerank.py
+46
-7
@@ -125,8 +125,8 @@ def test_invocations(server: RemoteOpenAIServer):
125
126
@pytest.mark.asyncio
127
@pytest.mark.parametrize("model_name", [MODEL_NAME])
128
-async def test_activation(server: RemoteOpenAIServer, model_name: str):
129
- async def get_outputs(activation):
128
+async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
129
+ async def get_outputs(use_activation):
130
query = "What is the capital of France?"
131
documents = [
132
"The capital of Brazil is Brasilia.",
@@ -139,16 +139,16 @@ async def test_activation(server: RemoteOpenAIServer, model_name: str):
139
"model": model_name,
140
"query": query,
141
"documents": documents,
142
- "activation": activation,
142
+ "use_activation": use_activation,
143
},
144
)
145
outputs = response.json()
146
147
return torch.tensor([x["relevance_score"] for x in outputs["results"]])
148
149
- default = await get_outputs(activation=None)
150
- w_activation = await get_outputs(activation=True)
151
- wo_activation = await get_outputs(activation=False)
149
+ default = await get_outputs(use_activation=None)
150
+ w_activation = await get_outputs(use_activation=True)
151
+ wo_activation = await get_outputs(use_activation=False)
152
153
assert torch.allclose(default, w_activation, atol=1e-2), (
154
"Default should use activation."
@@ -163,7 +163,25 @@ async def test_activation(server: RemoteOpenAIServer, model_name: str):
163
164
@pytest.mark.asyncio
165
@pytest.mark.parametrize("model_name", [MODEL_NAME])
166
-async def test_pooling(server: RemoteOpenAIServer, model_name: str):
166
+async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
167
+ input_text = "This product was excellent and exceeded my expectations"
168
+ response = requests.post(
169
+ server.url_for("pooling"),
170
+ json={
171
+ "model": model_name,
172
+ "input": input_text,
173
+ "encoding_format": "float",
174
+ "task": "classify",
175
+ },
176
+ )
177
+ poolings = PoolingResponse.model_validate(response.json())
178
+ assert len(poolings.data) == 1
179
+ assert len(poolings.data[0].data) == 1
180
+
181
+
182
+@pytest.mark.asyncio
183
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
184
+async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
185
input_text = ["The chef prepared a delicious meal."]
186
187
response = requests.post(
@@ -176,3 +194,24 @@ async def test_pooling(server: RemoteOpenAIServer, model_name: str):
194
assert len(poolings.data) == 1
195
assert len(poolings.data[0].data) == 11
196
assert len(poolings.data[0].data[0]) == 1
197
+
198
+
199
+@pytest.mark.asyncio
200
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
201
+@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
202
+async def test_pooling_not_supported(
203
+ server: RemoteOpenAIServer, model_name: str, task: str
204
+):
205
+ response = requests.post(
206
+ server.url_for("pooling"),
207
+ json={
208
+ "model": model_name,
209
+ "input": "test",
210
+ "encoding_format": "float",
211
+ "task": task,
212
+ },
213
+ )
214
+ assert response.json()["error"]["type"] == "BadRequestError"
215
+ assert response.json()["error"]["message"].startswith(
216
+ f"Task {task} is not supported"
217
+ )
tests/entrypoints/pooling/openai/test_score.py
+8
-8
@@ -218,8 +218,8 @@ class TestModel:
218
# TODO: reset this tolerance to 0.01 once we find
219
# an alternative to flash_attn with bfloat16
220
221
- def test_activation(self, server: RemoteOpenAIServer, model: dict[str, Any]):
222
- def get_outputs(activation):
221
+ def test_use_activation(self, server: RemoteOpenAIServer, model: dict[str, Any]):
222
+ def get_outputs(use_activation):
223
text_1 = "What is the capital of France?"
224
text_2 = "The capital of France is Paris."
225
response = requests.post(
@@ -228,7 +228,7 @@ class TestModel:
228
"model": model["name"],
229
"text_1": text_1,
230
"text_2": text_2,
231
- "activation": activation,
231
+ "use_activation": use_activation,
232
},
233
)
234
if response.status_code != 200:
@@ -238,9 +238,9 @@ class TestModel:
238
return torch.tensor([x["score"] for x in outputs["data"]])
239
240
if model["is_cross_encoder"]:
241
- default = get_outputs(activation=None)
242
- w_activation = get_outputs(activation=True)
243
- wo_activation = get_outputs(activation=False)
241
+ default = get_outputs(use_activation=None)
242
+ w_activation = get_outputs(use_activation=True)
243
+ wo_activation = get_outputs(use_activation=False)
244
245
assert torch.allclose(default, w_activation, atol=1e-2), (
246
"Default should use activation."
@@ -252,8 +252,8 @@ class TestModel:
252
"w_activation should be close to activation(wo_activation)."
253
)
254
else:
255
- get_outputs(activation=None)
255
+ get_outputs(use_activation=None)
256
257
# The activation parameter only works for the is_cross_encoder model
258
- response = get_outputs(activation=True)
258
+ response = get_outputs(use_activation=True)
259
assert response.status_code == 400
tests/models/language/pooling/test_pooler_config_init_behaviour.py
+4
-4
@@ -24,7 +24,7 @@ def test_classify_models_using_activation(
24
model,
25
max_model_len=512,
26
dtype=dtype,
27
- pooler_config=PoolerConfig(activation=False),
27
+ pooler_config=PoolerConfig(use_activation=False),
28
) as vllm_model:
29
wo_activation_out = vllm_model.classify(example_prompts)
30
@@ -32,7 +32,7 @@ def test_classify_models_using_activation(
32
model,
33
max_model_len=512,
34
dtype=dtype,
35
- pooler_config=PoolerConfig(activation=True),
35
+ pooler_config=PoolerConfig(use_activation=True),
36
) as vllm_model:
37
w_activation_out = vllm_model.classify(example_prompts)
38
@@ -104,7 +104,7 @@ def test_reward_models_using_activation(
104
model,
105
max_model_len=1024,
106
dtype=dtype,
107
- pooler_config=PoolerConfig(activation=False),
107
+ pooler_config=PoolerConfig(use_activation=False),
108
) as vllm_model:
109
wo_activation = vllm_model.reward(example_prompts)
110
@@ -112,7 +112,7 @@ def test_reward_models_using_activation(
112
model,
113
max_model_len=1024,
114
dtype=dtype,
115
- pooler_config=PoolerConfig(activation=True),
115
+ pooler_config=PoolerConfig(use_activation=True),
116
) as vllm_model:
117
w_activation = vllm_model.reward(example_prompts)
118
tests/test_pooling_params.py
+7
-7
@@ -17,7 +17,7 @@ EMBEDDING_MODELS = [
17
),
18
]
19
20
-classify_parameters = ["activation"]
20
+classify_parameters = ["use_activation"]
21
embed_parameters = ["dimensions", "normalize"]
22
step_pooling_parameters = ["step_tag_id", "returned_token_ids"]
23
@@ -88,13 +88,13 @@ def test_embed_dimensions(model_info: EmbedModelInfo):
88
def test_classify(task):
89
model_config = MockModelConfig(pooler_config=PoolerConfig(pooling_type="CLS"))
90
91
- pooling_params = PoolingParams(activation=None)
91
+ pooling_params = PoolingParams(use_activation=None)
92
pooling_params.verify(task=task, model_config=model_config)
93
94
- pooling_params = PoolingParams(activation=True)
94
+ pooling_params = PoolingParams(use_activation=True)
95
pooling_params.verify(task=task, model_config=model_config)
96
97
- pooling_params = PoolingParams(activation=False)
97
+ pooling_params = PoolingParams(use_activation=False)
98
pooling_params.verify(task=task, model_config=model_config)
99
100
invalid_parameters = embed_parameters + step_pooling_parameters
@@ -137,13 +137,13 @@ def test_token_classify(pooling_type: str):
137
pooler_config=PoolerConfig(pooling_type=pooling_type)
138
)
139
140
- pooling_params = PoolingParams(activation=None)
140
+ pooling_params = PoolingParams(use_activation=None)
141
pooling_params.verify(task=task, model_config=model_config)
142
143
- pooling_params = PoolingParams(activation=True)
143
+ pooling_params = PoolingParams(use_activation=True)
144
pooling_params.verify(task=task, model_config=model_config)
145
146
- pooling_params = PoolingParams(activation=False)
146
+ pooling_params = PoolingParams(use_activation=False)
147
pooling_params.verify(task=task, model_config=model_config)
148
149
invalid_parameters = embed_parameters
vllm/config/pooler.py
+32
-6
@@ -7,6 +7,9 @@ from typing import Any
7
from pydantic.dataclasses import dataclass
8
9
from vllm.config.utils import config
10
+from vllm.logger import init_logger
11
+
12
+logger = init_logger(__name__)
13
14
15
@config
@@ -48,7 +51,15 @@ class PoolerConfig:
51
"""
52
53
## for classification models
51
- activation: bool | None = None
54
+ softmax: float | None = None
55
+ """
56
+ softmax will be deprecated, please use use_activation instead.
57
+ """
58
+ activation: float | None = None
59
+ """
60
+ activation will be deprecated, please use use_activation instead.
61
+ """
62
+ use_activation: bool | None = None
63
"""
64
Whether to apply activation function to the classification outputs.
65
Defaults to True.
@@ -59,11 +70,6 @@ class PoolerConfig:
70
"""
71
72
## for reward models
62
- softmax: bool | None = None
63
- """
64
- Whether to apply softmax to the reward outputs.
65
- Defaults to True.
66
- """
73
step_tag_id: int | None = None
74
"""
75
If set, only the score corresponding to the `step_tag_id` in the
@@ -77,6 +83,10 @@ class PoolerConfig:
83
`math-shepherd-mistral-7b-prm` model.
84
"""
85
86
+ def __post_init__(self):
87
+ # raise deprecated warning for softmax and activation
88
+ self.use_activation = get_use_activation(self)
89
+
90
def compute_hash(self) -> str:
91
"""
92
WARNING: Whenever a new field is added to this config,
@@ -94,3 +104,19 @@ class PoolerConfig:
104
factors: list[Any] = []
105
hash_str = hashlib.md5(str(factors).encode(), usedforsecurity=False).hexdigest()
106
return hash_str
107
+
108
+
109
+def get_use_activation(o: object):
110
+ if softmax := getattr(o, "softmax", None) is not None:
111
+ logger.warning_once(
112
+ "softmax will be deprecated, please use use_activation instead."
113
+ )
114
+ return softmax
115
+
116
+ if activation := getattr(o, "activation", None) is not None:
117
+ logger.warning_once(
118
+ "activation will be deprecated, please use use_activation instead."
119
+ )
120
+ return activation
121
+
122
+ return getattr(o, "use_activation", None)
vllm/entrypoints/openai/api_server.py
+2
-6
@@ -107,6 +107,7 @@ from vllm.entrypoints.utils import (
107
)
108
from vllm.logger import init_logger
109
from vllm.reasoning import ReasoningParserManager
110
+from vllm.tasks import POOLING_TASKS
111
from vllm.usage.usage_lib import UsageContext
112
from vllm.utils.argparse_utils import FlexibleArgumentParser
113
from vllm.utils.network_utils import is_valid_ipv6_address
@@ -1748,12 +1749,7 @@ async def init_app_state(
1749
log_error_stack=args.log_error_stack,
1750
)
1751
)
1751
- if (
1752
- any(
1753
- task in supported_tasks
1754
- for task in ["token_embed", "token_classify", "plugin"]
1755
- )
1756
- )
1752
+ if any(task in POOLING_TASKS for task in supported_tasks)
1753
else None
1754
)
1755
state.openai_serving_embedding = (
vllm/entrypoints/openai/protocol.py
+100
-8
@@ -49,6 +49,8 @@ from openai.types.responses.response_reasoning_item import (
49
)
50
from openai_harmony import Message as OpenAIHarmonyMessage
51
52
+from vllm.config.pooler import get_use_activation
53
+from vllm.tasks import PoolingTask
54
from vllm.utils.serial_utils import (
55
EmbedDType,
56
EncodingFormat,
@@ -1669,8 +1671,58 @@ class EmbeddingChatRequest(OpenAIBaseModel):
1671
1672
EmbeddingRequest: TypeAlias = EmbeddingCompletionRequest | EmbeddingChatRequest
1673
1672
-PoolingCompletionRequest = EmbeddingCompletionRequest
1673
-PoolingChatRequest = EmbeddingChatRequest
1674
+
1675
+class PoolingCompletionRequest(EmbeddingCompletionRequest):
1676
+ task: PoolingTask | None = None
1677
+ softmax: bool | None = Field(
1678
+ default=None,
1679
+ description="softmax will be deprecated, please use use_activation instead.",
1680
+ )
1681
+ activation: bool | None = Field(
1682
+ default=None,
1683
+ description="activation will be deprecated, please use use_activation instead.",
1684
+ )
1685
+ use_activation: bool | None = Field(
1686
+ default=None,
1687
+ description="Whether to use activation for classification outputs. "
1688
+ "If it is a classify or token_classify task, the default is True; "
1689
+ "for other tasks, this value should be None.",
1690
+ )
1691
+
1692
+ def to_pooling_params(self):
1693
+ return PoolingParams(
1694
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
1695
+ dimensions=self.dimensions,
1696
+ normalize=self.normalize,
1697
+ use_activation=get_use_activation(self),
1698
+ )
1699
+
1700
+
1701
+class PoolingChatRequest(EmbeddingChatRequest):
1702
+ task: PoolingTask | None = None
1703
+ softmax: bool | None = Field(
1704
+ default=None,
1705
+ description="softmax will be deprecated, please use use_activation instead.",
1706
+ )
1707
+ activation: bool | None = Field(
1708
+ default=None,
1709
+ description="activation will be deprecated, please use use_activation instead.",
1710
+ )
1711
+ use_activation: bool | None = Field(
1712
+ default=None,
1713
+ description="Whether to use activation for classification outputs. "
1714
+ "If it is a classify or token_classify task, the default is True; "
1715
+ "for other tasks, this value should be None.",
1716
+ )
1717
+
1718
+ def to_pooling_params(self):
1719
+ return PoolingParams(
1720
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
1721
+ dimensions=self.dimensions,
1722
+ normalize=self.normalize,
1723
+ use_activation=get_use_activation(self),
1724
+ )
1725
+
1726
1727
T = TypeVar("T")
1728
@@ -1686,6 +1738,7 @@ class IOProcessorRequest(OpenAIBaseModel, Generic[T]):
1738
"""
1739
data: T
1740
1741
+ task: PoolingTask = "plugin"
1742
encoding_format: EncodingFormat = "float"
1743
embed_dtype: EmbedDType = Field(
1744
default="float32",
@@ -1749,14 +1802,27 @@ class ScoreRequest(OpenAIBaseModel):
1802
),
1803
)
1804
1752
- activation: bool | None = None
1805
+ softmax: bool | None = Field(
1806
+ default=None,
1807
+ description="softmax will be deprecated, please use use_activation instead.",
1808
+ )
1809
1810
+ activation: bool | None = Field(
1811
+ default=None,
1812
+ description="activation will be deprecated, please use use_activation instead.",
1813
+ )
1814
+
1815
+ use_activation: bool | None = Field(
1816
+ default=None,
1817
+ description="Whether to use activation for classification outputs. "
1818
+ "Default is True.",
1819
+ )
1820
# --8<-- [end:score-extra-params]
1821
1822
def to_pooling_params(self):
1823
return PoolingParams(
1824
truncate_prompt_tokens=self.truncate_prompt_tokens,
1759
- activation=self.activation,
1825
+ use_activation=get_use_activation(self),
1826
)
1827
1828
@@ -1783,14 +1849,27 @@ class RerankRequest(OpenAIBaseModel):
1849
),
1850
)
1851
1786
- activation: bool | None = None
1852
+ softmax: bool | None = Field(
1853
+ default=None,
1854
+ description="softmax will be deprecated, please use use_activation instead.",
1855
+ )
1856
1857
+ activation: bool | None = Field(
1858
+ default=None,
1859
+ description="activation will be deprecated, please use use_activation instead.",
1860
+ )
1861
+
1862
+ use_activation: bool | None = Field(
1863
+ default=None,
1864
+ description="Whether to use activation for classification outputs. "
1865
+ "Default is True.",
1866
+ )
1867
# --8<-- [end:rerank-extra-params]
1868
1869
def to_pooling_params(self):
1870
return PoolingParams(
1871
truncate_prompt_tokens=self.truncate_prompt_tokens,
1793
- activation=self.activation,
1872
+ use_activation=get_use_activation(self),
1873
)
1874
1875
@@ -1958,14 +2037,27 @@ class ClassificationRequest(OpenAIBaseModel):
2037
),
2038
)
2039
1961
- activation: bool | None = None
2040
+ softmax: bool | None = Field(
2041
+ default=None,
2042
+ description="softmax will be deprecated, please use use_activation instead.",
2043
+ )
2044
+
2045
+ activation: bool | None = Field(
2046
+ default=None,
2047
+ description="activation will be deprecated, please use use_activation instead.",
2048
+ )
2049
2050
+ use_activation: bool | None = Field(
2051
+ default=None,
2052
+ description="Whether to use activation for classification outputs. "
2053
+ "Default is True.",
2054
+ )
2055
# --8<-- [end:classification-extra-params]
2056
2057
def to_pooling_params(self):
2058
return PoolingParams(
2059
truncate_prompt_tokens=self.truncate_prompt_tokens,
1968
- activation=self.activation,
2060
+ use_activation=get_use_activation(self),
2061
)
2062
2063
vllm/entrypoints/openai/serving_pooling.py
+16
-7
@@ -170,15 +170,24 @@ class OpenAIServingPooling(OpenAIServing):
170
pooling_params = request.to_pooling_params()
171
172
pooling_task: PoolingTask
173
- if "token_embed" in self.supported_tasks:
174
- pooling_task = "token_embed"
175
- elif "token_classify" in self.supported_tasks:
176
- pooling_task = "token_classify"
177
- elif "plugin" in self.supported_tasks:
178
- pooling_task = "plugin"
173
+ if request.task is None:
174
+ if "token_embed" in self.supported_tasks:
175
+ pooling_task = "token_embed"
176
+ elif "token_classify" in self.supported_tasks:
177
+ pooling_task = "token_classify"
178
+ elif "plugin" in self.supported_tasks:
179
+ pooling_task = "plugin"
180
+ else:
181
+ return self.create_error_response(
182
+ f"pooling_task must be one of {self.supported_tasks}."
183
+ )
184
else:
185
+ pooling_task = request.task
186
+
187
+ if pooling_task not in self.supported_tasks:
188
return self.create_error_response(
181
- f"pooling_task must be one of {self.supported_tasks}."
189
+ f"Task {pooling_task} is not supported, it"
190
+ f" must be one of {self.supported_tasks}."
191
)
192
193
try:
vllm/model_executor/layers/pooler.py
+2
-2
@@ -607,7 +607,7 @@ class ClassifierPooler(Pooler):
607
pooled_data -= self.logit_bias
608
609
pooling_params = get_pooling_params(pooling_metadata)
610
- flags = [p.activation for p in pooling_params]
610
+ flags = [p.use_activation for p in pooling_params]
611
612
if len(set(flags)) == 1:
613
scores = self.act_fn(pooled_data) if flags[0] else pooled_data
@@ -681,7 +681,7 @@ class TokenClassifierPoolerHead(nn.Module):
681
if self.logit_bias is not None:
682
scores -= self.logit_bias
683
684
- if pooling_param.activation:
684
+ if pooling_param.use_activation:
685
scores = self.act_fn(scores)
686
687
# scores shape: [n_token, num_labels]
vllm/model_executor/models/config.py
+2
-2
@@ -53,8 +53,8 @@ class JambaForSequenceClassificationConfig(VerifyAndUpdateConfig):
53
@staticmethod
54
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
55
pooler_config = vllm_config.model_config.pooler_config
56
- if pooler_config.activation is None:
57
- pooler_config.activation = False
56
+ if pooler_config.use_activation is None:
57
+ pooler_config.use_activation = False
58
59
60
class JinaRobertaModelConfig(VerifyAndUpdateConfig):
vllm/pooling_params.py
+19
-13
@@ -2,16 +2,15 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from copy import deepcopy
5
-from typing import TYPE_CHECKING, Annotated, Any, Optional
5
+from typing import Annotated, Any, Optional
6
7
import msgspec
8
9
+from vllm.config import ModelConfig, PoolerConfig
10
+from vllm.config.pooler import get_use_activation
11
from vllm.sampling_params import RequestOutputKind
12
from vllm.tasks import PoolingTask
13
12
-if TYPE_CHECKING:
13
- from vllm.config import ModelConfig, PoolerConfig
14
-
14
15
class PoolingParams(
16
msgspec.Struct,
@@ -25,10 +24,12 @@ class PoolingParams(
24
Set to -1 to use the model's default truncation size.
25
Set to k to keep only the last k tokens (left truncation).
26
Set to None to disable truncation.
28
- normalize: Whether to normalize the embeddings outputs.
27
dimensions: Reduce the dimensions of embeddings
28
if model support matryoshka representation.
31
- activation: Whether to apply activation function to
29
+ normalize: Whether to normalize the embeddings outputs.
30
+ softmax: softmax will be deprecated, please use use_activation instead.
31
+ activation: activation will be deprecated, please use use_activation instead.
32
+ use_activation: Whether to apply activation function to
33
the classification outputs.
34
"""
35
@@ -44,7 +45,9 @@ class PoolingParams(
45
46
## for classification, scoring and rerank
47
# --8<-- [start:classification-pooling-params]
48
+ softmax: bool | None = None
49
activation: bool | None = None
50
+ use_activation: bool | None = None
51
# --8<-- [end:classification-pooling-params]
52
53
## for step pooling models
@@ -59,16 +62,16 @@ class PoolingParams(
62
63
@property
64
def all_parameters(self) -> list[str]:
62
- return ["dimensions", "normalize", "activation"]
65
+ return ["dimensions", "normalize", "use_activation"]
66
67
@property
68
def valid_parameters(self):
69
return {
70
"embed": ["dimensions", "normalize"],
68
- "classify": ["activation"],
69
- "score": ["activation"],
71
+ "classify": ["use_activation"],
72
+ "score": ["use_activation"],
73
"token_embed": ["dimensions", "normalize"],
71
- "token_classify": ["activation"],
74
+ "token_classify": ["use_activation"],
75
}
76
77
def clone(self) -> "PoolingParams":
@@ -84,6 +87,9 @@ class PoolingParams(
87
msg = f"You cannot overwrite {self.task=!r} with {task=!r}!"
88
raise ValueError(msg)
89
90
+ # raise deprecated warning for softmax and activation
91
+ self.use_activation = get_use_activation(self)
92
+
93
# plugin task uses io_processor.parse_request to verify inputs,
94
# skipping PoolingParams verify
95
if self.task == "plugin":
@@ -168,8 +174,8 @@ class PoolingParams(
174
raise ValueError("Dimensions must be greater than 0")
175
176
elif self.task in ["classify", "score", "token_classify"]:
171
- if self.activation is None:
172
- self.activation = True
177
+ if self.use_activation is None:
178
+ self.use_activation = True
179
else:
180
raise ValueError(f"Unknown pooling task: {self.task}")
181
@@ -197,7 +203,7 @@ class PoolingParams(
203
f"task={self.task}, "
204
f"normalize={self.normalize}, "
205
f"dimensions={self.dimensions}, "
200
- f"activation={self.activation}, "
206
+ f"use_activation={self.use_activation}, "
207
f"step_tag_id={self.step_tag_id}, "
208
f"returned_token_ids={self.returned_token_ids}, "
209
f"requires_token_ids={self.requires_token_ids}, "