@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 061980c36a
Fork 0

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 6, 14:30 UTC

[Feature][Frontend] add support for Cohere Embed v2 API (#37074)

Signed-off-by: walterbm <walter.beller.morales@gmail.com>

Walter Beller-Morales committed Mar 16, 2026 at 19:55 UTC 061980c36a7b78e5d8ea96893b79fd0b9c11a20e
16 files changed +1608 -39
docs/serving/openai_compatible_server.md
+134
@@ -72,6 +72,9 @@ In addition, we have the following custom APIs:
72 - Only applicable to [classification models](../models/pooling_models.md).
73 - [Score API](#score-api) (`/score`)
74 - Applicable to [embedding models and cross-encoder models](../models/pooling_models.md).
75 +- [Cohere Embed API](#cohere-embed-api) (`/v2/embed`)
76 + - Compatible with [Cohere's Embed API](https://docs.cohere.com/reference/embed)
77 + - Works with any [embedding model](../models/pooling_models.md), including multimodal models.
78 - [Re-rank API](#re-rank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
79 - Implements [Jina AI's v1 re-rank API](https://jina.ai/reranker/)
80 - Also compatible with [Cohere's v1 & v2 re-rank APIs](https://docs.cohere.com/v2/reference/rerank)
@@ -429,6 +432,137 @@ these extra parameters are supported instead:
432 --8<-- "vllm/entrypoints/pooling/base/protocol.py:embed-extra-params"
433 ```
434
435 +### Cohere Embed API
436 +
437 +Our API is also compatible with [Cohere's Embed v2 API](https://docs.cohere.com/reference/embed) which adds support for some modern embedding feature such as truncation, output dimensions, embedding types, and input types. This endpoint works with any embedding model (including multimodal models).
438 +
439 +#### Cohere Embed API request parameters
440 +
441 +| Parameter | Type | Required | Description |
442 +| --------- | ---- | -------- | ----------- |
443 +| `model` | string | Yes | Model name |
444 +| `input_type` | string | No | Prompt prefix key (model-dependent, see below) |
445 +| `texts` | list[string] | No | Text inputs (use one of `texts`, `images`, or `inputs`) |
446 +| `images` | list[string] | No | Base64 data URI images |
447 +| `inputs` | list[object] | No | Mixed text and image content objects |
448 +| `embedding_types` | list[string] | No | Output types (default: `["float"]`) |
449 +| `output_dimension` | int | No | Truncate embeddings to this dimension (Matryoshka) |
450 +| `truncate` | string | No | `END`, `START`, or `NONE` (default: `END`) |
451 +
452 +#### Text embedding
453 +
454 +```bash
455 +curl -X POST "http://localhost:8000/v2/embed" \
456 + -H "Content-Type: application/json" \
457 + -d '{
458 + "model": "Snowflake/snowflake-arctic-embed-m-v1.5",
459 + "input_type": "query",
460 + "texts": ["Hello world", "How are you?"],
461 + "embedding_types": ["float"]
462 + }'
463 +```
464 +
465 +??? console "Response"
466 +
467 + ```json
468 + {
469 + "id": "embd-...",
470 + "embeddings": {
471 + "float": [
472 + [0.012, -0.034, ...],
473 + [0.056, 0.078, ...]
474 + ]
475 + },
476 + "texts": ["Hello world", "How are you?"],
477 + "meta": {
478 + "api_version": {"version": "2"},
479 + "billed_units": {"input_tokens": 12}
480 + }
481 + }
482 + ```
483 +
484 +#### Mixed text and image inputs
485 +
486 +For multimodal models, you can embed images by passing base64 data URIs. The `inputs` field accepts a list of objects with mixed text and image content:
487 +
488 +```bash
489 +curl -X POST "http://localhost:8000/v2/embed" \
490 + -H "Content-Type: application/json" \
491 + -d '{
492 + "model": "google/siglip-so400m-patch14-384",
493 + "inputs": [
494 + {
495 + "content": [
496 + {"type": "text", "text": "A photo of a cat"},
497 + {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBOR..."}}
498 + ]
499 + }
500 + ],
501 + "embedding_types": ["float"]
502 + }'
503 +```
504 +
505 +#### Embedding types
506 +
507 +The `embedding_types` parameter controls the output format. Multiple types can be requested in a single call:
508 +
509 +| Type | Description |
510 +| ---- | ----------- |
511 +| `float` | Raw float32 embeddings (default) |
512 +| `binary` | Bit-packed signed binary |
513 +| `ubinary` | Bit-packed unsigned binary |
514 +| `base64` | Little-endian float32 encoded as base64 |
515 +
516 +```bash
517 +curl -X POST "http://localhost:8000/v2/embed" \
518 + -H "Content-Type: application/json" \
519 + -d '{
520 + "model": "Snowflake/snowflake-arctic-embed-m-v1.5",
521 + "input_type": "query",
522 + "texts": ["What is machine learning?"],
523 + "embedding_types": ["float", "binary"]
524 + }'
525 +```
526 +
527 +??? console "Response"
528 +
529 + ```json
530 + {
531 + "id": "embd-...",
532 + "embeddings": {
533 + "float": [[0.012, -0.034, ...]],
534 + "binary": [[42, -117, ...]]
535 + },
536 + "texts": ["What is machine learning?"],
537 + "meta": {
538 + "api_version": {"version": "2"},
539 + "billed_units": {"input_tokens": 8}
540 + }
541 + }
542 + ```
543 +
544 +#### Truncation
545 +
546 +The `truncate` parameter controls how inputs exceeding the model's maximum sequence length are handled:
547 +
548 +| Value | Behavior |
549 +| ----- | --------- |
550 +| `END` (default) | Keep the first tokens, drop the end |
551 +| `START` | Keep the last tokens, drop the beginning |
552 +| `NONE` | Return an error if the input is too long |
553 +
554 +#### Input type and prompt prefixes
555 +
556 +The `input_type` field selects a prompt prefix to prepend to each text input. The available values
557 +depend on the model:
558 +
559 +- **Models with `task_instructions` in `config.json`**: The keys from the `task_instructions` dict are
560 + the valid `input_type` values and the corresponding value is prepended to each text.
561 +- **Models with `config_sentence_transformers.json` prompts**: The keys from the `prompts` dict are
562 + the valid `input_type` values. For example, `Snowflake/snowflake-arctic-embed-xs` defines `"query"`,
563 + so setting `input_type: "query"` prepends `"Represent this sentence for searching relevant passages: "`.
564 +- **Other models**: `input_type` is not accepted and will raise a validation error if passed.
565 +
566 ### Transcriptions API
567
568 Our Transcriptions API is compatible with [OpenAI's Transcriptions API](https://platform.openai.com/docs/api-reference/audio/createTranscription);
tests/entrypoints/pooling/embed/test_cohere_online.py new
+310
@@ -0,0 +1,310 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the Cohere /v2/embed API with generic (non-Cohere) models.
4 +
5 +Validates that the Cohere v2 embed endpoint works correctly with standard
6 +embedding models, covering text embedding, embedding type conversions,
7 +response structure, batching, normalisation, and semantic similarity.
8 +"""
9 +
10 +import base64
11 +import struct
12 +
13 +import numpy as np
14 +import pytest
15 +import requests
16 +
17 +from tests.utils import RemoteOpenAIServer
18 +
19 +DTYPE = "bfloat16"
20 +
21 +MODELS: list[tuple[str, list[str]]] = [
22 + ("intfloat/multilingual-e5-small", []),
23 + (
24 + "Snowflake/snowflake-arctic-embed-m-v1.5",
25 + [
26 + "--trust_remote_code",
27 + "--hf_overrides",
28 + '{"matryoshka_dimensions":[256]}',
29 + ],
30 + ),
31 +]
32 +
33 +
34 +@pytest.fixture(scope="module", params=MODELS, ids=lambda m: m[0])
35 +def model_config(request):
36 + return request.param
37 +
38 +
39 +@pytest.fixture(scope="module")
40 +def model_name(model_config):
41 + return model_config[0]
42 +
43 +
44 +@pytest.fixture(scope="module")
45 +def server(model_config):
46 + name, extra_args = model_config
47 + args = [
48 + "--runner",
49 + "pooling",
50 + "--dtype",
51 + DTYPE,
52 + "--enforce-eager",
53 + "--max-model-len",
54 + "512",
55 + "--gpu-memory-utilization",
56 + "0.02",
57 + ] + extra_args
58 + with RemoteOpenAIServer(name, args) as remote_server:
59 + yield remote_server
60 +
61 +
62 +def _cohere_embed(
63 + server: RemoteOpenAIServer,
64 + model_name: str,
65 + texts: list[str] | None = None,
66 + images: list[str] | None = None,
67 + input_type: str | None = None,
68 + embedding_types: list[str] | None = None,
69 +) -> dict:
70 + body: dict = {"model": model_name}
71 + if input_type is not None:
72 + body["input_type"] = input_type
73 + if texts is not None:
74 + body["texts"] = texts
75 + if images is not None:
76 + body["images"] = images
77 + if embedding_types is not None:
78 + body["embedding_types"] = embedding_types
79 + resp = requests.post(server.url_for("/v2/embed"), json=body)
80 + resp.raise_for_status()
81 + return resp.json()
82 +
83 +
84 +def _openai_embed(
85 + server: RemoteOpenAIServer, model_name: str, texts: list[str]
86 +) -> dict:
87 + body = {"model": model_name, "input": texts, "encoding_format": "float"}
88 + resp = requests.post(server.url_for("/v1/embeddings"), json=body)
89 + resp.raise_for_status()
90 + return resp.json()
91 +
92 +
93 +def _cosine_sim(a: list[float], b: list[float]) -> float:
94 + va, vb = np.array(a), np.array(b)
95 + return float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb)))
96 +
97 +
98 +# -----------------------------------------------------------
99 +# Text embedding tests
100 +# -----------------------------------------------------------
101 +
102 +
103 +def test_basic_embed(server: RemoteOpenAIServer, model_name: str):
104 + r = _cohere_embed(
105 + server, model_name, texts=["hello world"], embedding_types=["float"]
106 + )
107 + assert "embeddings" in r
108 + assert len(r["embeddings"]["float"]) == 1
109 + assert len(r["embeddings"]["float"][0]) > 0
110 +
111 +
112 +def test_unsupported_input_type_rejected(server: RemoteOpenAIServer, model_name: str):
113 + """An input_type not defined in the model's prompt config should be
114 + rejected with a 400 error."""
115 + body = {
116 + "model": model_name,
117 + "input_type": "nonexistent_type",
118 + "texts": ["hello world"],
119 + "embedding_types": ["float"],
120 + }
121 + resp = requests.post(server.url_for("/v2/embed"), json=body)
122 + assert resp.status_code == 400
123 + assert "Unsupported input_type" in resp.json()["error"]["message"]
124 +
125 +
126 +def test_omitted_input_type_accepted(server: RemoteOpenAIServer, model_name: str):
127 + """Omitting input_type should always work (no prompt prefix applied)."""
128 + body = {
129 + "model": model_name,
130 + "texts": ["hello world"],
131 + "embedding_types": ["float"],
132 + }
133 + resp = requests.post(server.url_for("/v2/embed"), json=body)
134 + assert resp.status_code == 200
135 + data = resp.json()
136 + assert len(data["embeddings"]["float"]) == 1
137 +
138 +
139 +def test_v1_v2_parity(server: RemoteOpenAIServer, model_name: str):
140 + """v1 (OpenAI) and v2 (Cohere) endpoints should produce the same
141 + float embeddings for a generic model."""
142 + texts = ["hello world"]
143 + v2 = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
144 + v1 = _openai_embed(server, model_name, texts)
145 + cos = _cosine_sim(v2["embeddings"]["float"][0], v1["data"][0]["embedding"])
146 + assert cos > 0.9999, f"v1/v2 parity failed, cosine={cos}"
147 +
148 +
149 +def test_embedding_types(server: RemoteOpenAIServer, model_name: str):
150 + r = _cohere_embed(
151 + server,
152 + model_name,
153 + texts=["test"],
154 + embedding_types=["float", "binary", "ubinary"],
155 + )
156 + dim = len(r["embeddings"]["float"][0])
157 + assert len(r["embeddings"]["binary"][0]) == dim // 8
158 + assert len(r["embeddings"]["ubinary"][0]) == dim // 8
159 +
160 +
161 +def test_response_structure(server: RemoteOpenAIServer, model_name: str):
162 + r = _cohere_embed(server, model_name, texts=["test"], embedding_types=["float"])
163 + assert "id" in r
164 + assert "embeddings" in r
165 + assert "texts" in r
166 + assert r["texts"] == ["test"]
167 + assert "meta" in r
168 + assert r["meta"]["api_version"]["version"] == "2"
169 + assert "billed_units" in r["meta"]
170 + assert r["meta"]["billed_units"]["input_tokens"] > 0
171 + assert r["meta"]["billed_units"]["image_tokens"] == 0
172 +
173 +
174 +def test_batch(server: RemoteOpenAIServer, model_name: str):
175 + texts = ["apple", "banana", "cherry"]
176 + r = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
177 + assert len(r["embeddings"]["float"]) == 3
178 + dim = len(r["embeddings"]["float"][0])
179 + for emb in r["embeddings"]["float"]:
180 + assert len(emb) == dim
181 +
182 +
183 +def test_l2_normalized(server: RemoteOpenAIServer, model_name: str):
184 + r = _cohere_embed(
185 + server, model_name, texts=["hello world"], embedding_types=["float"]
186 + )
187 + emb = np.array(r["embeddings"]["float"][0])
188 + assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
189 +
190 +
191 +def test_semantic_similarity(server: RemoteOpenAIServer, model_name: str):
192 + r = _cohere_embed(
193 + server,
194 + model_name,
195 + texts=["machine learning", "deep learning", "chocolate cake recipe"],
196 + embedding_types=["float"],
197 + )
198 + embs = r["embeddings"]["float"]
199 + cos_related = _cosine_sim(embs[0], embs[1])
200 + cos_unrelated = _cosine_sim(embs[0], embs[2])
201 + assert cos_related > cos_unrelated
202 +
203 +
204 +def test_missing_input_returns_error(server: RemoteOpenAIServer, model_name: str):
205 + body = {"model": model_name}
206 + resp = requests.post(server.url_for("/v2/embed"), json=body)
207 + assert resp.status_code == 400
208 +
209 +
210 +def test_base64_embedding_type(server: RemoteOpenAIServer, model_name: str):
211 + r = _cohere_embed(
212 + server,
213 + model_name,
214 + texts=["test encoding"],
215 + embedding_types=["float", "base64"],
216 + )
217 + float_emb = r["embeddings"]["float"][0]
218 + b64_str = r["embeddings"]["base64"][0]
219 + decoded = struct.unpack(f"<{len(float_emb)}f", base64.b64decode(b64_str))
220 + np.testing.assert_allclose(float_emb, decoded, rtol=1e-5)
221 +
222 +
223 +# -----------------------------------------------------------
224 +# Truncation tests
225 +# -----------------------------------------------------------
226 +
227 +
228 +def _cohere_embed_raw(
229 + server: RemoteOpenAIServer,
230 + body: dict,
231 +) -> requests.Response:
232 + return requests.post(server.url_for("/v2/embed"), json=body)
233 +
234 +
235 +def test_truncate_end_succeeds(server: RemoteOpenAIServer, model_name: str):
236 + """truncate=END should silently truncate long input."""
237 + long_text = " ".join(["word"] * 2000)
238 + body = {
239 + "model": model_name,
240 + "texts": [long_text],
241 + "embedding_types": ["float"],
242 + "truncate": "END",
243 + }
244 + resp = _cohere_embed_raw(server, body)
245 + assert resp.status_code == 200
246 + data = resp.json()
247 + assert len(data["embeddings"]["float"]) == 1
248 +
249 +
250 +def test_truncate_start_succeeds(server: RemoteOpenAIServer, model_name: str):
251 + """truncate=START should silently truncate long input from the start."""
252 + long_text = " ".join(["word"] * 2000)
253 + body = {
254 + "model": model_name,
255 + "texts": [long_text],
256 + "embedding_types": ["float"],
257 + "truncate": "START",
258 + }
259 + resp = _cohere_embed_raw(server, body)
260 + assert resp.status_code == 200
261 + data = resp.json()
262 + assert len(data["embeddings"]["float"]) == 1
263 +
264 +
265 +def test_truncate_none_rejects_long_input(server: RemoteOpenAIServer, model_name: str):
266 + """truncate=NONE should error when input exceeds model context."""
267 + long_text = " ".join(["word"] * 2000)
268 + body = {
269 + "model": model_name,
270 + "texts": [long_text],
271 + "embedding_types": ["float"],
272 + "truncate": "NONE",
273 + }
274 + resp = _cohere_embed_raw(server, body)
275 + assert resp.status_code == 400
276 +
277 +
278 +def test_truncate_start_vs_end_differ(server: RemoteOpenAIServer, model_name: str):
279 + """START and END truncation should produce different embeddings
280 + when the input is long enough to actually be truncated.
281 +
282 + We construct input with distinct tokens at the start vs end
283 + so that keeping different halves produces different embeddings.
284 + """
285 + start_words = " ".join([f"alpha{i}" for i in range(300)])
286 + end_words = " ".join([f"omega{i}" for i in range(300)])
287 + long_text = start_words + " " + end_words
288 +
289 + body_end = {
290 + "model": model_name,
291 + "texts": [long_text],
292 + "embedding_types": ["float"],
293 + "truncate": "END",
294 + }
295 + body_start = {
296 + "model": model_name,
297 + "texts": [long_text],
298 + "embedding_types": ["float"],
299 + "truncate": "START",
300 + }
301 + r_end = _cohere_embed_raw(server, body_end).json()
302 + r_start = _cohere_embed_raw(server, body_start).json()
303 +
304 + emb_end = r_end["embeddings"]["float"][0]
305 + emb_start = r_start["embeddings"]["float"][0]
306 + cos = _cosine_sim(emb_end, emb_start)
307 + assert cos < 0.99, (
308 + f"START and END truncation should produce different embeddings "
309 + f"for long input, but cosine similarity was {cos}"
310 + )
tests/entrypoints/pooling/embed/test_cohere_online_vision.py new
+135
@@ -0,0 +1,135 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the Cohere /v2/embed API with a multimodal model (SigLIP).
4 +
5 +Validates image embedding, batching, normalisation, and embedding type
6 +conversions through the /v2/embed endpoint.
7 +"""
8 +
9 +import base64
10 +import struct
11 +import zlib
12 +
13 +import numpy as np
14 +import pytest
15 +import requests
16 +
17 +from tests.utils import RemoteOpenAIServer
18 +
19 +MODEL_NAME = "google/siglip-so400m-patch14-384"
20 +DTYPE = "bfloat16"
21 +
22 +
23 +@pytest.fixture(scope="module")
24 +def server():
25 + args = [
26 + "--runner",
27 + "pooling",
28 + "--dtype",
29 + DTYPE,
30 + "--enforce-eager",
31 + "--max-model-len",
32 + "64",
33 + "--gpu-memory-utilization",
34 + "0.3",
35 + ]
36 + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
37 + yield remote_server
38 +
39 +
40 +def _make_tiny_png(r: int, g: int, b: int, w: int = 2, h: int = 2) -> str:
41 + raw = b""
42 + for _ in range(h):
43 + raw += b"\x00" + bytes([r, g, b]) * w
44 + compressed = zlib.compress(raw)
45 +
46 + def chunk(ctype: bytes, cdata: bytes) -> bytes:
47 + c = ctype + cdata
48 + return (
49 + struct.pack(">I", len(cdata))
50 + + c
51 + + struct.pack(">I", zlib.crc32(c) & 0xFFFFFFFF)
52 + )
53 +
54 + ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)
55 + png = (
56 + b"\x89PNG\r\n\x1a\n"
57 + + chunk(b"IHDR", ihdr)
58 + + chunk(b"IDAT", compressed)
59 + + chunk(b"IEND", b"")
60 + )
61 + return "data:image/png;base64," + base64.b64encode(png).decode()
62 +
63 +
64 +def _cohere_embed(
65 + server: RemoteOpenAIServer,
66 + texts: list[str] | None = None,
67 + images: list[str] | None = None,
68 + embedding_types: list[str] | None = None,
69 +) -> dict:
70 + body: dict = {"model": MODEL_NAME}
71 + if texts is not None:
72 + body["texts"] = texts
73 + if images is not None:
74 + body["images"] = images
75 + if embedding_types is not None:
76 + body["embedding_types"] = embedding_types
77 + resp = requests.post(server.url_for("/v2/embed"), json=body)
78 + resp.raise_for_status()
79 + return resp.json()
80 +
81 +
82 +def test_image_embed(server: RemoteOpenAIServer):
83 + img_uri = _make_tiny_png(255, 0, 0)
84 + r = _cohere_embed(
85 + server,
86 + images=[img_uri],
87 + embedding_types=["float"],
88 + )
89 + assert "embeddings" in r
90 + assert len(r["embeddings"]["float"]) == 1
91 + assert len(r["embeddings"]["float"][0]) > 0
92 + assert r["meta"]["billed_units"]["image_tokens"] > 0
93 + assert r["meta"]["billed_units"]["input_tokens"] == 0
94 +
95 +
96 +def test_image_batch(server: RemoteOpenAIServer):
97 + red = _make_tiny_png(255, 0, 0)
98 + blue = _make_tiny_png(0, 0, 255)
99 + r = _cohere_embed(
100 + server,
101 + images=[red, blue],
102 + embedding_types=["float"],
103 + )
104 + assert len(r["embeddings"]["float"]) == 2
105 +
106 +
107 +def test_image_l2_normalized(server: RemoteOpenAIServer):
108 + img_uri = _make_tiny_png(0, 255, 0)
109 + r = _cohere_embed(
110 + server,
111 + images=[img_uri],
112 + embedding_types=["float"],
113 + )
114 + emb = np.array(r["embeddings"]["float"][0])
115 + assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
116 +
117 +
118 +def test_image_embedding_types(server: RemoteOpenAIServer):
119 + img_uri = _make_tiny_png(128, 128, 128)
120 + r = _cohere_embed(
121 + server,
122 + images=[img_uri],
123 + embedding_types=["float", "binary", "ubinary"],
124 + )
125 + dim = len(r["embeddings"]["float"][0])
126 + assert len(r["embeddings"]["binary"][0]) == dim // 8
127 + assert len(r["embeddings"]["ubinary"][0]) == dim // 8
128 +
129 +
130 +def test_text_embed_on_multimodal(server: RemoteOpenAIServer):
131 + """SigLIP also supports text-only embedding via /v2/embed."""
132 + r = _cohere_embed(server, texts=["hello world"], embedding_types=["float"])
133 + assert "embeddings" in r
134 + assert len(r["embeddings"]["float"]) == 1
135 + assert len(r["embeddings"]["float"][0]) > 0
tests/entrypoints/pooling/embed/test_cohere_openai_parity.py new
+102
@@ -0,0 +1,102 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Parity test between Cohere /v2/embed and OpenAI /v1/embeddings.
4 +
5 +Verifies that both endpoints produce identical float embeddings when
6 +no prompt prefix is applied (input_type omitted for Cohere /v2/embed).
7 +"""
8 +
9 +import numpy as np
10 +import pytest
11 +import requests
12 +
13 +from tests.utils import RemoteOpenAIServer
14 +
15 +MODEL_NAME = "BAAI/bge-base-en-v1.5"
16 +DTYPE = "bfloat16"
17 +
18 +
19 +@pytest.fixture(scope="module")
20 +def server():
21 + args = [
22 + "--runner",
23 + "pooling",
24 + "--dtype",
25 + DTYPE,
26 + "--enforce-eager",
27 + "--max-model-len",
28 + "512",
29 + "--gpu-memory-utilization",
30 + "0.02",
31 + ]
32 + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
33 + yield remote_server
34 +
35 +
36 +def _cohere_embed(
37 + server: RemoteOpenAIServer,
38 + texts: list[str],
39 +) -> list[list[float]]:
40 + body = {
41 + "model": MODEL_NAME,
42 + "texts": texts,
43 + "embedding_types": ["float"],
44 + }
45 + resp = requests.post(server.url_for("/v2/embed"), json=body)
46 + resp.raise_for_status()
47 + return resp.json()["embeddings"]["float"]
48 +
49 +
50 +def _openai_embed(
51 + server: RemoteOpenAIServer,
52 + texts: list[str],
53 +) -> list[list[float]]:
54 + body = {"model": MODEL_NAME, "input": texts, "encoding_format": "float"}
55 + resp = requests.post(server.url_for("/v1/embeddings"), json=body)
56 + resp.raise_for_status()
57 + return [item["embedding"] for item in resp.json()["data"]]
58 +
59 +
60 +def test_single_text_parity(server: RemoteOpenAIServer):
61 + """A single text should produce identical embeddings via both APIs."""
62 + texts = ["the quick brown fox jumps over the lazy dog"]
63 + v2 = _cohere_embed(server, texts)
64 + v1 = _openai_embed(server, texts)
65 + np.testing.assert_allclose(v2[0], v1[0], rtol=1e-5)
66 +
67 +
68 +def test_batch_parity(server: RemoteOpenAIServer):
69 + """A batch of texts should produce identical embeddings via both APIs,
70 + in the same order."""
71 + texts = [
72 + "machine learning",
73 + "deep learning",
74 + "natural language processing",
75 + ]
76 + v2 = _cohere_embed(server, texts)
77 + v1 = _openai_embed(server, texts)
78 + assert len(v2) == len(v1) == 3
79 + for i in range(3):
80 + np.testing.assert_allclose(v2[i], v1[i], rtol=1e-5, err_msg=f"index {i}")
81 +
82 +
83 +def test_token_count_parity(server: RemoteOpenAIServer):
84 + """Both APIs should report the same prompt token count."""
85 + texts = ["hello world"]
86 + v2_resp = requests.post(
87 + server.url_for("/v2/embed"),
88 + json={
89 + "model": MODEL_NAME,
90 + "texts": texts,
91 + "embedding_types": ["float"],
92 + },
93 + )
94 + v1_resp = requests.post(
95 + server.url_for("/v1/embeddings"),
96 + json={"model": MODEL_NAME, "input": texts, "encoding_format": "float"},
97 + )
98 + v2_resp.raise_for_status()
99 + v1_resp.raise_for_status()
100 + v2_tokens = v2_resp.json()["meta"]["billed_units"]["input_tokens"]
101 + v1_tokens = v1_resp.json()["usage"]["prompt_tokens"]
102 + assert v2_tokens == v1_tokens
tests/entrypoints/pooling/embed/test_io_processor.py new
+208
@@ -0,0 +1,208 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for EmbedIOProcessor."""
4 +
5 +import pytest
6 +
7 +from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
8 +from vllm.entrypoints.pooling.embed.protocol import (
9 + CohereEmbedRequest,
10 +)
11 +
12 +
13 +class TestResolveTruncation:
14 + """Unit tests for EmbedIOProcessor._resolve_cohere_truncation."""
15 +
16 + @staticmethod
17 + def _make_request(**kwargs) -> CohereEmbedRequest:
18 + defaults = {
19 + "model": "test",
20 + "input_type": "search_document",
21 + "texts": ["hello"],
22 + }
23 + return CohereEmbedRequest(**(defaults | kwargs))
24 +
25 + def test_truncate_end_default(self):
26 + req = self._make_request()
27 + tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
28 + assert tokens == -1
29 + assert side is None
30 +
31 + def test_truncate_end_explicit(self):
32 + req = self._make_request(truncate="END")
33 + tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
34 + assert tokens == -1
35 + assert side is None
36 +
37 + def test_truncate_end_with_max_tokens(self):
38 + req = self._make_request(truncate="END", max_tokens=128)
39 + tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
40 + assert tokens == 128
41 + assert side is None
42 +
43 + def test_truncate_none(self):
44 + req = self._make_request(truncate="NONE")
45 + tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
46 + assert tokens is None
47 + assert side is None
48 +
49 + def test_truncate_none_with_max_tokens(self):
50 + """truncate=NONE should NOT set truncate_prompt_tokens; the
51 + max_tokens limit is enforced separately via _check_max_tokens."""
52 + req = self._make_request(truncate="NONE", max_tokens=10)
53 + tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
54 + assert tokens is None
55 + assert side is None
56 +
57 + def test_truncate_start(self):
58 + req = self._make_request(truncate="START")
59 + tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
60 + assert tokens == -1
61 + assert side == "left"
62 +
63 + def test_truncate_start_with_max_tokens(self):
64 + req = self._make_request(truncate="START", max_tokens=64)
65 + tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
66 + assert tokens == 64
67 + assert side == "left"
68 +
69 +
70 +class TestApplyStPrompt:
71 + """Unit tests for EmbedIOProcessor._apply_task_instruction."""
72 +
73 + @staticmethod
74 + def _make_handler(task_instructions: dict[str, str] | None):
75 + handler = object.__new__(EmbedIOProcessor)
76 + handler.task_instructions = task_instructions
77 + return handler
78 +
79 + def test_no_prompts_configured(self):
80 + handler = self._make_handler(None)
81 + texts = ["hello", "world"]
82 + assert handler._apply_task_instruction(texts, "query") is texts
83 +
84 + def test_matching_input_type(self):
85 + handler = self._make_handler({"query": "search_query: "})
86 + result = handler._apply_task_instruction(["hello"], "query")
87 + assert result == ["search_query: hello"]
88 +
89 + def test_non_matching_input_type(self):
90 + handler = self._make_handler({"query": "search_query: "})
91 + texts = ["hello"]
92 + assert handler._apply_task_instruction(texts, "document") is texts
93 +
94 + def test_multiple_texts(self):
95 + handler = self._make_handler(
96 + {"query": "Represent this sentence for searching: "}
97 + )
98 + result = handler._apply_task_instruction(["a", "b", "c"], "query")
99 + assert result == [
100 + "Represent this sentence for searching: a",
101 + "Represent this sentence for searching: b",
102 + "Represent this sentence for searching: c",
103 + ]
104 +
105 + def test_empty_prefix_returns_unchanged(self):
106 + handler = self._make_handler({"passage": ""})
107 + texts = ["hello"]
108 + assert handler._apply_task_instruction(texts, "passage") is texts
109 +
110 +
111 +class TestLoadTaskInstructions:
112 + """Unit tests for EmbedIOProcessor._load_task_instructions."""
113 +
114 + def test_no_attribute(self):
115 + class FakeConfig:
116 + pass
117 +
118 + assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
119 +
120 + def test_with_task_instructions(self):
121 + class FakeConfig:
122 + task_instructions = {
123 + "retrieval.query": "Represent the query: ",
124 + "retrieval.passage": "",
125 + }
126 +
127 + result = EmbedIOProcessor._load_task_instructions(FakeConfig())
128 + assert result == {
129 + "retrieval.query": "Represent the query: ",
130 + "retrieval.passage": "",
131 + }
132 +
133 + def test_empty_dict(self):
134 + class FakeConfig:
135 + task_instructions = {}
136 +
137 + assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
138 +
139 + def test_non_dict(self):
140 + class FakeConfig:
141 + task_instructions = "not a dict"
142 +
143 + assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
144 +
145 +
146 +class TestCheckMaxTokens:
147 + """Unit tests for EmbedIOProcessor._check_cohere_max_tokens."""
148 +
149 + @staticmethod
150 + def _fake_output(n_tokens: int):
151 + class _Out:
152 + def __init__(self, n: int):
153 + self.prompt_token_ids = list(range(n))
154 +
155 + return _Out(n_tokens)
156 +
157 + def test_none_check_is_noop(self):
158 + outs = [self._fake_output(100)]
159 + EmbedIOProcessor._check_cohere_max_tokens(outs, None)
160 +
161 + def test_within_limit(self):
162 + outs = [self._fake_output(5), self._fake_output(3)]
163 + EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
164 +
165 + def test_exceeds_limit(self):
166 + outs = [self._fake_output(3), self._fake_output(10)]
167 + with pytest.raises(ValueError, match="exceeds max_tokens=5"):
168 + EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
169 +
170 + def test_exact_limit(self):
171 + outs = [self._fake_output(5)]
172 + EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
173 +
174 +
175 +class TestValidateInputType:
176 + """Unit tests for EmbedIOProcessor._validate_input_type."""
177 +
178 + @staticmethod
179 + def _make_handler(task_instructions: dict[str, str] | None):
180 + handler = object.__new__(EmbedIOProcessor)
181 + handler.task_instructions = task_instructions
182 + return handler
183 +
184 + def test_none_input_type_always_accepted(self):
185 + handler = self._make_handler(None)
186 + handler._validate_input_type(None)
187 + handler_with = self._make_handler({"query": "q: "})
188 + handler_with._validate_input_type(None)
189 +
190 + def test_no_prompts_rejects(self):
191 + handler = self._make_handler(None)
192 + with pytest.raises(ValueError, match="does not define any input_type"):
193 + handler._validate_input_type("anything")
194 +
195 + def test_known_type_accepted(self):
196 + handler = self._make_handler({"query": "q: ", "document": "d: "})
197 + handler._validate_input_type("query")
198 + handler._validate_input_type("document")
199 +
200 + def test_unknown_type_rejected(self):
201 + handler = self._make_handler({"query": "q: ", "document": "d: "})
202 + with pytest.raises(ValueError, match="Unsupported input_type 'other'"):
203 + handler._validate_input_type("other")
204 +
205 + def test_error_lists_supported(self):
206 + handler = self._make_handler({"a": "", "b": ""})
207 + with pytest.raises(ValueError, match="Supported values: a, b"):
208 + handler._validate_input_type("z")
tests/entrypoints/pooling/embed/test_protocol.py new
+129
@@ -0,0 +1,129 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for Cohere embed protocol: build_typed_embeddings and its
4 +underlying packing helpers, plus Cohere-specific serving helpers."""
5 +
6 +import base64
7 +import struct
8 +
9 +import numpy as np
10 +import pytest
11 +
12 +from vllm.entrypoints.pooling.embed.protocol import (
13 + build_typed_embeddings,
14 +)
15 +
16 +
17 +@pytest.fixture
18 +def sample_embeddings() -> list[list[float]]:
19 + return [
20 + [0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, -0.8],
21 + [-0.05, 0.15, -0.25, 0.35, -0.45, 0.55, -0.65, 0.75],
22 + ]
23 +
24 +
25 +class TestBuildTypedEmbeddingsFloat:
26 + def test_float_passthrough(self, sample_embeddings: list[list[float]]):
27 + result = build_typed_embeddings(sample_embeddings, ["float"])
28 + assert result.float == sample_embeddings
29 + assert result.binary is None
30 +
31 + def test_empty_input(self):
32 + result = build_typed_embeddings([], ["float"])
33 + assert result.float == []
34 +
35 +
36 +class TestBuildTypedEmbeddingsBinary:
37 + def test_binary_packing(self):
38 + # 8 values: positive->1, negative->0 => bits: 10101010 = 0xAA = 170
39 + # signed: 170 - 128 = 42
40 + embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
41 + result = build_typed_embeddings(embs, ["binary"])
42 + assert result.binary is not None
43 + assert result.binary[0] == [42]
44 +
45 + def test_ubinary_packing(self):
46 + embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
47 + result = build_typed_embeddings(embs, ["ubinary"])
48 + assert result.ubinary is not None
49 + assert result.ubinary[0] == [170] # 0b10101010
50 +
51 + def test_binary_all_positive(self):
52 + embs = [[0.1] * 8]
53 + result = build_typed_embeddings(embs, ["binary"])
54 + assert result.binary is not None
55 + # all bits = 1 => 0xFF = 255, signed: 255 - 128 = 127
56 + assert result.binary[0] == [127]
57 +
58 + def test_binary_all_negative(self):
59 + embs = [[-0.1] * 8]
60 + result = build_typed_embeddings(embs, ["binary"])
61 + assert result.binary is not None
62 + # all bits = 0, signed: 0 - 128 = -128
63 + assert result.binary[0] == [-128]
64 +
65 + def test_binary_dimension_is_eighth(self, sample_embeddings: list[list[float]]):
66 + result = build_typed_embeddings(sample_embeddings, ["binary"])
67 + assert result.binary is not None
68 + for orig, packed in zip(sample_embeddings, result.binary):
69 + assert len(packed) == len(orig) // 8
70 +
71 + def test_zero_treated_as_positive(self):
72 + embs = [[0.0] * 8]
73 + result = build_typed_embeddings(embs, ["binary"])
74 + assert result.binary is not None
75 + # 0.0 >= 0 is True, so bit=1 for all => 127 (signed)
76 + assert result.binary[0] == [127]
77 +
78 + def test_non_multiple_of_8_raises(self):
79 + embs = [[0.1] * 7]
80 + with pytest.raises(ValueError, match="multiple of 8"):
81 + build_typed_embeddings(embs, ["binary"])
82 +
83 + def test_ubinary_non_multiple_of_8_raises(self):
84 + embs = [[0.1] * 10]
85 + with pytest.raises(ValueError, match="multiple of 8"):
86 + build_typed_embeddings(embs, ["ubinary"])
87 +
88 +
89 +class TestBuildTypedEmbeddingsBase64:
90 + def test_base64_roundtrip(self, sample_embeddings: list[list[float]]):
91 + result = build_typed_embeddings(sample_embeddings, ["base64"])
92 + assert result.base64 is not None
93 + assert len(result.base64) == 2
94 +
95 + for orig, b64_str in zip(sample_embeddings, result.base64):
96 + decoded = base64.b64decode(b64_str)
97 + n = len(orig)
98 + values = struct.unpack(f"<{n}f", decoded)
99 + np.testing.assert_allclose(orig, values, rtol=1e-5)
100 +
101 + def test_base64_byte_length(self):
102 + embs = [[0.1, 0.2, 0.3]]
103 + result = build_typed_embeddings(embs, ["base64"])
104 + assert result.base64 is not None
105 + raw = base64.b64decode(result.base64[0])
106 + assert len(raw) == 3 * 4 # 3 floats * 4 bytes each
107 +
108 +
109 +class TestBuildTypedEmbeddingsMultiple:
110 + def test_all_types_at_once(self, sample_embeddings: list[list[float]]):
111 + result = build_typed_embeddings(
112 + sample_embeddings,
113 + ["float", "binary", "ubinary", "base64"],
114 + )
115 + assert result.float is not None
116 + assert result.binary is not None
117 + assert result.ubinary is not None
118 + assert result.base64 is not None
119 +
120 + def test_subset_types(self, sample_embeddings: list[list[float]]):
121 + result = build_typed_embeddings(sample_embeddings, ["float", "binary"])
122 + assert result.float is not None
123 + assert result.binary is not None
124 + assert result.ubinary is None
125 + assert result.base64 is None
126 +
127 + def test_unknown_type_ignored(self, sample_embeddings: list[list[float]]):
128 + result = build_typed_embeddings(sample_embeddings, ["float", "unknown_type"])
129 + assert result.float is not None
vllm/entrypoints/pooling/base/protocol.py
+9 -1
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 -from typing import Annotated, Any
5 +from typing import Annotated, Any, Literal
6
7 from pydantic import Field, model_validator
8
@@ -24,6 +24,14 @@ class PoolingBasicRequestMixin(OpenAIBaseModel):
24
25 # --8<-- [start:pooling-common-extra-params]
26 truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
27 + truncation_side: Literal["left", "right"] | None = Field(
28 + default=None,
29 + description=(
30 + "Which side to truncate from when truncate_prompt_tokens is active. "
31 + "'right' keeps the first N tokens. "
32 + "'left' keeps the last N tokens."
33 + ),
34 + )
35 request_id: str = Field(
36 default_factory=random_uuid,
37 description=(
vllm/entrypoints/pooling/classify/protocol.py
+2
@@ -32,6 +32,7 @@ class ClassificationCompletionRequest(
32 max_total_tokens=model_config.max_model_len,
33 max_output_tokens=0,
34 truncate_prompt_tokens=self.truncate_prompt_tokens,
35 + truncation_side=self.truncation_side,
36 do_lower_case=encoder_config.get("do_lower_case", False),
37 add_special_tokens=self.add_special_tokens,
38 max_total_tokens_param="max_model_len",
@@ -54,6 +55,7 @@ class ClassificationChatRequest(
55 max_total_tokens=model_config.max_model_len,
56 max_output_tokens=0,
57 truncate_prompt_tokens=self.truncate_prompt_tokens,
58 + truncation_side=self.truncation_side,
59 do_lower_case=encoder_config.get("do_lower_case", False),
60 add_special_tokens=self.add_special_tokens,
61 max_total_tokens_param="max_model_len",
vllm/entrypoints/pooling/embed/api_router.py
+26 -5
@@ -7,12 +7,12 @@ from fastapi import APIRouter, Depends, Request
7
8 from vllm.entrypoints.openai.engine.protocol import ErrorResponse
9 from vllm.entrypoints.openai.utils import validate_json_request
10 -from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest
11 -from vllm.entrypoints.pooling.embed.serving import ServingEmbedding
12 -from vllm.entrypoints.utils import (
13 - load_aware_call,
14 - with_cancellation,
10 +from vllm.entrypoints.pooling.embed.protocol import (
11 + CohereEmbedRequest,
12 + EmbeddingRequest,
13 )
14 +from vllm.entrypoints.pooling.embed.serving import ServingEmbedding
15 +from vllm.entrypoints.utils import load_aware_call, with_cancellation
16
17 router = APIRouter()
18
@@ -40,3 +40,24 @@ async def create_embedding(
40 raise NotImplementedError("The model does not support Embeddings API")
41
42 return await handler(request, raw_request)
43 +
44 +
45 +@router.post(
46 + "/v2/embed",
47 + dependencies=[Depends(validate_json_request)],
48 + responses={
49 + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
50 + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
51 + },
52 +)
53 +@with_cancellation
54 +@load_aware_call
55 +async def create_cohere_embedding(
56 + request: CohereEmbedRequest,
57 + raw_request: Request,
58 +):
59 + handler = embedding(raw_request)
60 + if handler is None:
61 + raise NotImplementedError("The model does not support Embeddings API")
62 +
63 + return await handler(request, raw_request)
vllm/entrypoints/pooling/embed/io_processor.py
+302 -17
@@ -1,14 +1,37 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Any, cast
3 +from collections.abc import Sequence
4 +from typing import Any, Literal, cast
5
6 import torch
6 -
7 +from openai.types.chat import (
8 + ChatCompletionContentPartImageParam,
9 + ChatCompletionContentPartTextParam,
10 +)
11 +from openai.types.chat.chat_completion_content_part_image_param import ImageURL
12 +
13 +from vllm import PoolingParams
14 +from vllm.entrypoints.chat_utils import (
15 + ChatCompletionContentPartParam,
16 + ChatCompletionMessageParam,
17 + CustomChatCompletionMessageParam,
18 +)
19 from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
20 +from vllm.entrypoints.pooling.embed.protocol import (
21 + CohereEmbedInput,
22 + CohereEmbedRequest,
23 + EmbeddingChatRequest,
24 + EmbeddingCompletionRequest,
25 +)
26 from vllm.entrypoints.pooling.typing import PoolingServeContext
27 from vllm.inputs.data import ProcessorInputs, token_inputs
28 +from vllm.logger import init_logger
29 from vllm.outputs import PoolingOutput, PoolingRequestOutput
30 +from vllm.renderers import merge_kwargs
31 from vllm.utils.collection_utils import chunk_list
32 +from vllm.utils.mistral import is_mistral_tokenizer
33 +
34 +logger = init_logger(__name__)
35
36
37 class EmbedIOProcessor(PoolingIOProcessor):
@@ -21,16 +44,45 @@ class EmbedIOProcessor(PoolingIOProcessor):
44 self.pooler_config = self.model_config.pooler_config
45 self.enable_chunked_processing = self.pooler_config.enable_chunked_processing
46
24 - #################################################################
25 - # Long Text Embedding with Chunked Processing
26 - # PTAL: examples/pooling/embed/openai_embedding_long_text
47 + # Load task instructions from HF config or sentence-transformers config
48 + self.task_instructions: dict[str, str] | None = self._load_task_instructions(
49 + self.model_config.hf_config
50 + ) or self._load_st_prompts(self.model_config.model, self.model_config.revision)
51 + if self.task_instructions:
52 + logger.info(
53 + "Loaded prompt prefixes for input_type: %s",
54 + list(self.task_instructions.keys()),
55 + )
56
57 def pre_process_online(self, ctx: PoolingServeContext):
29 - super().pre_process_online(ctx)
58 + if isinstance(ctx.request, CohereEmbedRequest):
59 + self._pre_process_cohere_online(ctx)
60 + else:
61 + super().pre_process_online(ctx)
62 +
63 + if self.enable_chunked_processing:
64 + self._pre_process_chunked(ctx)
65 +
66 + def post_process_online(
67 + self,
68 + ctx: PoolingServeContext,
69 + ):
70 + if ctx.final_res_batch is None:
71 + raise ValueError("Final response batch not available")
72
73 if not self.enable_chunked_processing:
32 - return None
74 + self._enforce_cohere_max_tokens(ctx)
75 + return super().post_process_online(ctx)
76
77 + self._post_process_chunked(ctx)
78 + self._enforce_cohere_max_tokens(ctx)
79 +
80 + #################################################################
81 + # Long Text Embedding with Chunked Processing
82 + # PTAL: examples/pooling/embed/openai_embedding_long_text
83 + #################################################################
84 +
85 + def _pre_process_chunked(self, ctx: PoolingServeContext) -> None:
86 if ctx.engine_prompts is None:
87 raise ValueError("Engine prompts not available")
88
@@ -61,18 +113,10 @@ class EmbedIOProcessor(PoolingIOProcessor):
113
114 ctx.engine_prompts = chunked_engine_prompts
115 ctx.prompt_request_ids = prompt_request_ids
64 - return None
116
66 - def post_process_online(
67 - self,
68 - ctx: PoolingServeContext,
69 - ):
70 - if ctx.final_res_batch is None:
71 - raise ValueError("Final response batch not available")
72 -
73 - if not self.enable_chunked_processing:
74 - return super().post_process_online(ctx)
117 + return None
118
119 + def _post_process_chunked(self, ctx: PoolingServeContext) -> None:
120 # Online aggregation for chunked requests to
121 # minimize memory usage
122 # Track aggregation state for each prompt
@@ -195,4 +239,245 @@ class EmbedIOProcessor(PoolingIOProcessor):
239 raise ValueError(f"Result not found for prompt {prompt_idx}")
240
241 ctx.final_res_batch = final_res_batch
242 +
243 return None
244 +
245 + #################################################################
246 + # Cohere Request Preprocessing & Postprocessing
247 + #################################################################
248 +
249 + @staticmethod
250 + def _load_task_instructions(hf_config: Any) -> dict[str, str] | None:
251 + """Extract ``task_instructions`` from the HF model config."""
252 + ti = getattr(hf_config, "task_instructions", None)
253 + if not isinstance(ti, dict) or not ti:
254 + return None
255 + return {k: v for k, v in ti.items() if isinstance(v, str)}
256 +
257 + @staticmethod
258 + def _load_st_prompts(
259 + model: str | Any,
260 + revision: str | None,
261 + ) -> dict[str, str] | None:
262 + """Load ``task_instructions`` from ``config_sentence_transformers.json``."""
263 + from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
264 +
265 + try:
266 + cfg = get_hf_file_to_dict(
267 + "config_sentence_transformers.json", str(model), revision
268 + )
269 + except (ValueError, OSError):
270 + return None
271 +
272 + if cfg is None:
273 + return None
274 + prompts = cfg.get("prompts")
275 + if not isinstance(prompts, dict) or not prompts:
276 + return None
277 + return {k: v for k, v in prompts.items() if isinstance(v, str)}
278 +
279 + @staticmethod
280 + def _mixed_input_to_messages(
281 + inp: CohereEmbedInput,
282 + *,
283 + task_prefix: str | None = None,
284 + ) -> list[ChatCompletionMessageParam]:
285 + """Build chat messages from a mixed text+image input.
286 +
287 + When *task_prefix* is given, it is prepended to each text part.
288 + """
289 + parts: list[ChatCompletionContentPartParam] = []
290 + for item in inp.content:
291 + if item.type == "text" and item.text is not None:
292 + text = task_prefix + item.text if task_prefix else item.text
293 + parts.append(ChatCompletionContentPartTextParam(type="text", text=text))
294 + elif item.type == "image_url" and item.image_url is not None:
295 + parts.append(
296 + ChatCompletionContentPartImageParam(
297 + type="image_url",
298 + image_url=ImageURL(url=item.image_url["url"]),
299 + )
300 + )
301 + return [CustomChatCompletionMessageParam(role="user", content=parts)]
302 +
303 + @staticmethod
304 + def _check_cohere_max_tokens(
305 + outputs: list[PoolingRequestOutput],
306 + max_tokens_check: int | None,
307 + ) -> None:
308 + """Raise if any output exceeds *max_tokens_check* tokens.
309 +
310 + Used to enforce ``truncate=NONE`` with an explicit ``max_tokens``:
311 + the pipeline runs without truncation and we reject afterwards.
312 + """
313 + if max_tokens_check is None:
314 + return
315 + for out in outputs:
316 + n = len(out.prompt_token_ids)
317 + if n > max_tokens_check:
318 + raise ValueError(
319 + f"Input of {n} tokens exceeds max_tokens={max_tokens_check} "
320 + "with truncate=NONE. Set truncate to END or START to "
321 + "allow truncation."
322 + )
323 +
324 + @staticmethod
325 + def _resolve_cohere_truncation(
326 + request: CohereEmbedRequest,
327 + ) -> tuple[int | None, Literal["left", "right"] | None]:
328 + """Return ``(truncate_prompt_tokens, truncation_side)``."""
329 + if request.truncate == "NONE":
330 + return None, None
331 + if request.truncate == "START":
332 + tokens = request.max_tokens if request.max_tokens is not None else -1
333 + return tokens, "left"
334 + if request.max_tokens is not None:
335 + return request.max_tokens, None
336 + return -1, None
337 +
338 + def create_pooling_params(self, request):
339 + if isinstance(request, CohereEmbedRequest):
340 + return PoolingParams(
341 + task="embed",
342 + dimensions=request.output_dimension,
343 + )
344 + return super().create_pooling_params(request)
345 +
346 + def _pre_process_cohere_online(self, ctx: PoolingServeContext) -> None:
347 + """Convert a ``CohereEmbedRequest`` into engine prompts.
348 +
349 + For texts, a single batched completion request path is used.
350 + For images and mixed inputs, conversations are batch-rendered
351 + through the chat template in one ``render_chat`` call.
352 + """
353 + request = ctx.request
354 + assert isinstance(request, CohereEmbedRequest)
355 +
356 + if request.texts is None and request.images is None and request.inputs is None:
357 + raise ValueError("One of texts, images, or inputs must be provided")
358 +
359 + truncate_prompt_tokens, truncation_side = self._resolve_cohere_truncation(
360 + request
361 + )
362 + input_type = request.input_type
363 + self._validate_input_type(input_type)
364 +
365 + if request.images is not None:
366 + all_messages: list[list[ChatCompletionMessageParam]] = [
367 + [
368 + CustomChatCompletionMessageParam(
369 + role="user",
370 + content=[{"type": "image_url", "image_url": {"url": uri}}],
371 + )
372 + ]
373 + for uri in request.images
374 + ]
375 + ctx.engine_prompts = self._batch_render_chat(
376 + request, all_messages, truncate_prompt_tokens, truncation_side
377 + )
378 +
379 + elif request.inputs is not None:
380 + task_prefix = self._get_task_instruction_prefix(input_type)
381 + all_messages = [
382 + self._mixed_input_to_messages(inp, task_prefix=task_prefix)
383 + for inp in request.inputs
384 + ]
385 + ctx.engine_prompts = self._batch_render_chat(
386 + request, all_messages, truncate_prompt_tokens, truncation_side
387 + )
388 +
389 + else:
390 + prefixed = self._apply_task_instruction(request.texts or [], input_type)
391 + proxy = EmbeddingCompletionRequest(
392 + model=request.model,
393 + input=prefixed,
394 + dimensions=request.output_dimension,
395 + encoding_format="float",
396 + truncate_prompt_tokens=truncate_prompt_tokens,
397 + truncation_side=truncation_side,
398 + )
399 + ctx.engine_prompts = self._preprocess_completion_online(
400 + proxy, prompt_input=proxy.input, prompt_embeds=None
401 + )
402 +
403 + def _batch_render_chat(
404 + self,
405 + request: CohereEmbedRequest,
406 + all_messages: Sequence[list[ChatCompletionMessageParam]],
407 + truncate_prompt_tokens: int | None,
408 + truncation_side: Literal["left", "right"] | None,
409 + ) -> list[ProcessorInputs]:
410 + """Batch-render multiple conversations through the chat template."""
411 + if not all_messages:
412 + return []
413 +
414 + proxy = EmbeddingChatRequest(
415 + model=request.model,
416 + messages=list(all_messages[0]),
417 + dimensions=request.output_dimension,
418 + encoding_format="float",
419 + truncate_prompt_tokens=truncate_prompt_tokens,
420 + truncation_side=truncation_side,
421 + )
422 +
423 + renderer = self.renderer
424 + mm_config = self.model_config.multimodal_config
425 +
426 + tok_params = proxy.build_tok_params(self.model_config)
427 + chat_params = proxy.build_chat_params(
428 + self.chat_template,
429 + self.chat_template_content_format,
430 + ).with_defaults(
431 + merge_kwargs(
432 + None,
433 + dict(
434 + tools=None,
435 + tokenize=is_mistral_tokenizer(renderer.tokenizer),
436 + ),
437 + ),
438 + default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
439 + )
440 +
441 + _, engine_prompts = renderer.render_chat(all_messages, chat_params, tok_params)
442 + return engine_prompts
443 +
444 + def _validate_input_type(self, input_type: str | None) -> None:
445 + """Raise if *input_type* is not supported by this model."""
446 + if input_type is None:
447 + return
448 + if self.task_instructions is None:
449 + raise ValueError(
450 + f"Unsupported input_type {input_type!r}. "
451 + "This model does not define any input_type task instructions."
452 + )
453 + if input_type not in self.task_instructions:
454 + supported = ", ".join(sorted(self.task_instructions))
455 + raise ValueError(
456 + f"Unsupported input_type {input_type!r}. Supported values: {supported}"
457 + )
458 +
459 + def _apply_task_instruction(
460 + self,
461 + texts: list[str],
462 + input_type: str | None,
463 + ) -> list[str]:
464 + """Prepend the task-instruction prefix for *input_type*.
465 +
466 + Returns *texts* unchanged when no matching prefix is configured.
467 + """
468 + prefix = self._get_task_instruction_prefix(input_type)
469 + if not prefix:
470 + return texts
471 + return [prefix + t for t in texts]
472 +
473 + def _get_task_instruction_prefix(self, input_type: str | None) -> str | None:
474 + """Return the task-instruction prefix for *input_type*, or ``None``."""
475 + if not self.task_instructions or input_type is None:
476 + return None
477 + return self.task_instructions.get(input_type) or None
478 +
479 + def _enforce_cohere_max_tokens(self, ctx: PoolingServeContext) -> None:
480 + if isinstance(ctx.request, CohereEmbedRequest):
481 + request = ctx.request
482 + if request.truncate == "NONE" and request.max_tokens is not None:
483 + self._check_cohere_max_tokens(ctx.final_res_batch, request.max_tokens)
vllm/entrypoints/pooling/embed/protocol.py
+167 -3
@@ -1,9 +1,19 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Embedding API protocol models for OpenAI and Cohere formats.
4 +
5 +OpenAI: https://platform.openai.com/docs/api-reference/embeddings
6 +Cohere: https://docs.cohere.com/reference/embed
7 +"""
8 +
9 +import base64
10 +import builtins
11 +import struct
12 import time
4 -from typing import TypeAlias
13 +from collections.abc import Sequence
14 +from typing import Literal, TypeAlias
15
6 -from pydantic import Field
16 +from pydantic import BaseModel, Field
17
18 from vllm import PoolingParams
19 from vllm.config import ModelConfig
@@ -17,6 +27,10 @@ from vllm.entrypoints.pooling.base.protocol import (
27 from vllm.renderers import TokenizeParams
28 from vllm.utils import random_uuid
29
30 +# ---------------------------------------------------------------------------
31 +# OpenAI /v1/embeddings — request models
32 +# ---------------------------------------------------------------------------
33 +
34
35 def _get_max_total_output_tokens(
36 model_config: ModelConfig,
@@ -50,6 +64,7 @@ class EmbeddingCompletionRequest(
64 max_total_tokens=max_total_tokens,
65 max_output_tokens=max_output_tokens,
66 truncate_prompt_tokens=self.truncate_prompt_tokens,
67 + truncation_side=self.truncation_side,
68 do_lower_case=encoder_config.get("do_lower_case", False),
69 add_special_tokens=self.add_special_tokens,
70 max_total_tokens_param="max_model_len",
@@ -79,6 +94,7 @@ class EmbeddingChatRequest(
94 max_total_tokens=max_total_tokens,
95 max_output_tokens=max_output_tokens,
96 truncate_prompt_tokens=self.truncate_prompt_tokens,
97 + truncation_side=self.truncation_side,
98 do_lower_case=encoder_config.get("do_lower_case", False),
99 add_special_tokens=self.add_special_tokens,
100 max_total_tokens_param="max_model_len",
@@ -96,6 +112,11 @@ class EmbeddingChatRequest(
112 EmbeddingRequest: TypeAlias = EmbeddingCompletionRequest | EmbeddingChatRequest
113
114
115 +# ---------------------------------------------------------------------------
116 +# OpenAI /v1/embeddings — response models
117 +# ---------------------------------------------------------------------------
118 +
119 +
120 class EmbeddingResponseData(OpenAIBaseModel):
121 index: int
122 object: str = "embedding"
@@ -106,7 +127,7 @@ class EmbeddingResponse(OpenAIBaseModel):
127 id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
128 object: str = "list"
129 created: int = Field(default_factory=lambda: int(time.time()))
109 - model: str
130 + model: str | None = None
131 data: list[EmbeddingResponseData]
132 usage: UsageInfo
133
@@ -115,3 +136,146 @@ class EmbeddingBytesResponse(OpenAIBaseModel):
136 content: list[bytes]
137 headers: dict[str, str] | None = None
138 media_type: str = "application/octet-stream"
139 +
140 +
141 +# ---------------------------------------------------------------------------
142 +# Cohere /v2/embed — request models
143 +# ---------------------------------------------------------------------------
144 +
145 +CohereEmbeddingType = Literal[
146 + "float",
147 + "binary",
148 + "ubinary",
149 + "base64",
150 +]
151 +CohereTruncate = Literal["NONE", "START", "END"]
152 +
153 +
154 +class CohereEmbedContent(BaseModel):
155 + type: Literal["text", "image_url"]
156 + text: str | None = None
157 + image_url: dict[str, str] | None = None
158 +
159 +
160 +class CohereEmbedInput(BaseModel):
161 + content: list[CohereEmbedContent]
162 +
163 +
164 +class CohereEmbedRequest(BaseModel):
165 + model: str | None = None
166 + input_type: str | None = None
167 + texts: list[str] | None = None
168 + images: list[str] | None = None
169 + inputs: list[CohereEmbedInput] | None = None
170 + output_dimension: int | None = None
171 + embedding_types: list[CohereEmbeddingType] | None = None
172 + truncate: CohereTruncate = "END"
173 + max_tokens: int | None = None
174 + priority: int = 0
175 +
176 +
177 +# ---------------------------------------------------------------------------
178 +# Cohere /v2/embed — response models
179 +# ---------------------------------------------------------------------------
180 +
181 +
182 +class CohereApiVersion(BaseModel):
183 + version: str = "2"
184 +
185 +
186 +class CohereBilledUnits(BaseModel):
187 + input_tokens: int | None = None
188 + image_tokens: int | None = None
189 +
190 +
191 +class CohereMeta(BaseModel):
192 + api_version: CohereApiVersion = Field(default_factory=CohereApiVersion)
193 + billed_units: CohereBilledUnits | None = None
194 +
195 +
196 +class CohereEmbedByTypeEmbeddings(BaseModel):
197 + # The field name ``float`` shadows the builtin type, so the annotation
198 + # must use ``builtins.float`` to avoid a self-referential type error.
199 + float: list[list[builtins.float]] | None = None
200 + binary: list[list[int]] | None = None
201 + ubinary: list[list[int]] | None = None
202 + base64: list[str] | None = None
203 +
204 +
205 +class CohereEmbedResponse(BaseModel):
206 + id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
207 + embeddings: CohereEmbedByTypeEmbeddings
208 + texts: list[str] | None = None
209 + meta: CohereMeta | None = None
210 + response_type: Literal["embeddings_by_type"] = "embeddings_by_type"
211 +
212 +
213 +# ---------------------------------------------------------------------------
214 +# Cohere embedding type conversion helpers
215 +# ---------------------------------------------------------------------------
216 +
217 +_UNSIGNED_TO_SIGNED_DIFF = 1 << 7 # 128
218 +
219 +
220 +def _pack_binary_embeddings(
221 + float_embeddings: list[list[float]],
222 + signed: bool,
223 +) -> list[list[int]]:
224 + """Bit-pack float embeddings: positive -> 1, negative -> 0.
225 +
226 + Each bit is shifted left by ``7 - idx%8``, and every 8 bits are packed
227 + into one byte.
228 + """
229 + result: list[list[int]] = []
230 + for embedding in float_embeddings:
231 + dim = len(embedding)
232 + if dim % 8 != 0:
233 + raise ValueError(
234 + "Embedding dimension must be a multiple of 8 for binary "
235 + f"embedding types, but got {dim}."
236 + )
237 + packed_len = dim // 8
238 + packed: list[int] = []
239 + byte_val = 0
240 + for idx, value in enumerate(embedding):
241 + bit = 1 if value >= 0 else 0
242 + byte_val += bit << (7 - idx % 8)
243 + if (idx + 1) % 8 == 0:
244 + if signed:
245 + byte_val -= _UNSIGNED_TO_SIGNED_DIFF
246 + packed.append(byte_val)
247 + byte_val = 0
248 + assert len(packed) == packed_len
249 + result.append(packed)
250 + return result
251 +
252 +
253 +def _encode_base64_embeddings(
254 + float_embeddings: list[list[float]],
255 +) -> list[str]:
256 + """Encode float embeddings as base64 (little-endian float32)."""
257 + result: list[str] = []
258 + for embedding in float_embeddings:
259 + buf = struct.pack(f"<{len(embedding)}f", *embedding)
260 + result.append(base64.b64encode(buf).decode("utf-8"))
261 + return result
262 +
263 +
264 +def build_typed_embeddings(
265 + float_embeddings: list[list[float]],
266 + embedding_types: Sequence[str],
267 +) -> CohereEmbedByTypeEmbeddings:
268 + """Convert float embeddings to all requested Cohere embedding types."""
269 + result = CohereEmbedByTypeEmbeddings()
270 +
271 + for emb_type in embedding_types:
272 + if emb_type == "float":
273 + result.float = float_embeddings
274 + elif emb_type == "binary":
275 + result.binary = _pack_binary_embeddings(float_embeddings, signed=True)
276 + elif emb_type == "ubinary":
277 + result.ubinary = _pack_binary_embeddings(float_embeddings, signed=False)
278 + elif emb_type == "base64":
279 + result.base64 = _encode_base64_embeddings(float_embeddings)
280 +
281 + return result
vllm/entrypoints/pooling/embed/serving.py
+53 -11
@@ -5,7 +5,7 @@ from collections.abc import Callable
5 from functools import partial
6 from typing import Literal, TypeAlias, cast
7
8 -from fastapi.responses import JSONResponse, StreamingResponse
8 +from fastapi.responses import JSONResponse, Response, StreamingResponse
9 from typing_extensions import assert_never
10
11 from vllm.config import ModelConfig
@@ -14,10 +14,15 @@ from vllm.entrypoints.openai.engine.protocol import UsageInfo
14 from vllm.entrypoints.pooling.base.serving import PoolingServing
15 from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
16 from vllm.entrypoints.pooling.embed.protocol import (
17 + CohereBilledUnits,
18 + CohereEmbedRequest,
19 + CohereEmbedResponse,
20 + CohereMeta,
21 EmbeddingBytesResponse,
22 EmbeddingRequest,
23 EmbeddingResponse,
24 EmbeddingResponseData,
25 + build_typed_embeddings,
26 )
27 from vllm.entrypoints.pooling.typing import PoolingServeContext
28 from vllm.entrypoints.pooling.utils import (
@@ -26,24 +31,23 @@ from vllm.entrypoints.pooling.utils import (
31 encode_pooling_output_float,
32 get_json_response_cls,
33 )
34 +from vllm.logger import init_logger
35 from vllm.outputs import PoolingRequestOutput
36 from vllm.renderers import BaseRenderer
37 from vllm.utils.serial_utils import EmbedDType, Endianness
38
39 +logger = init_logger(__name__)
40 +
41 JSONResponseCLS = get_json_response_cls()
42
43 EmbeddingServeContext: TypeAlias = PoolingServeContext[EmbeddingRequest]
44
45
46 class ServingEmbedding(PoolingServing):
39 - """
40 - Embedding API similar to OpenAI's API.
41 -
42 - See https://platform.openai.com/docs/api-reference/embeddings/create
43 - for the API specification. This API mimics the OpenAI Embedding API.
44 - """
47 + """Embedding API supporting both OpenAI and Cohere formats."""
48
49 request_id_prefix = "embd"
50 + io_processor: EmbedIOProcessor
51
52 def init_io_processor(
53 self,
@@ -58,6 +62,14 @@ class ServingEmbedding(PoolingServing):
62 )
63
64 async def _build_response(
65 + self,
66 + ctx: PoolingServeContext,
67 + ) -> Response:
68 + if isinstance(ctx.request, CohereEmbedRequest):
69 + return self._build_cohere_response_from_ctx(ctx)
70 + return await self._build_openai_response(ctx)
71 +
72 + async def _build_openai_response(
73 self,
74 ctx: EmbeddingServeContext,
75 ) -> JSONResponse | StreamingResponse:
@@ -66,7 +78,7 @@ class ServingEmbedding(PoolingServing):
78 endianness = ctx.request.endianness
79
80 if encoding_format == "float" or encoding_format == "base64":
69 - return self._request_output_to_embed_json_response(
81 + return self._openai_json_response(
82 ctx.final_res_batch,
83 ctx.request_id,
84 ctx.created_time,
@@ -77,7 +89,7 @@ class ServingEmbedding(PoolingServing):
89 )
90
91 if encoding_format == "bytes" or encoding_format == "bytes_only":
80 - return self._request_output_to_to_embed_bytes_response(
92 + return self._openai_bytes_response(
93 ctx.final_res_batch,
94 ctx.request_id,
95 ctx.created_time,
@@ -89,7 +101,7 @@ class ServingEmbedding(PoolingServing):
101
102 assert_never(encoding_format)
103
92 - def _request_output_to_embed_json_response(
104 + def _openai_json_response(
105 self,
106 final_res_batch: list[PoolingRequestOutput],
107 request_id: str,
@@ -139,7 +151,7 @@ class ServingEmbedding(PoolingServing):
151 )
152 return JSONResponseCLS(content=response.model_dump())
153
142 - def _request_output_to_to_embed_bytes_response(
154 + def _openai_bytes_response(
155 self,
156 final_res_batch: list[PoolingRequestOutput],
157 request_id: str,
@@ -177,3 +189,33 @@ class ServingEmbedding(PoolingServing):
189 headers=response.headers,
190 media_type=response.media_type,
191 )
192 +
193 + @staticmethod
194 + def _build_cohere_response_from_ctx(
195 + ctx: PoolingServeContext,
196 + ) -> JSONResponse:
197 + request = ctx.request
198 + assert isinstance(request, CohereEmbedRequest)
199 +
200 + all_floats = [encode_pooling_output_float(out) for out in ctx.final_res_batch]
201 + total_tokens = sum(len(out.prompt_token_ids) for out in ctx.final_res_batch)
202 +
203 + image_tokens = total_tokens if request.images is not None else 0
204 + texts_echo = request.texts
205 +
206 + embedding_types = request.embedding_types or ["float"]
207 + embeddings_obj = build_typed_embeddings(all_floats, embedding_types)
208 +
209 + input_tokens = total_tokens - image_tokens
210 + response = CohereEmbedResponse(
211 + id=ctx.request_id,
212 + embeddings=embeddings_obj,
213 + texts=texts_echo,
214 + meta=CohereMeta(
215 + billed_units=CohereBilledUnits(
216 + input_tokens=input_tokens,
217 + image_tokens=image_tokens,
218 + ),
219 + ),
220 + )
221 + return JSONResponse(content=response.model_dump(exclude_none=True))
vllm/entrypoints/pooling/pooling/protocol.py
+3
@@ -36,6 +36,7 @@ class PoolingCompletionRequest(
36 max_total_tokens=model_config.max_model_len,
37 max_output_tokens=0,
38 truncate_prompt_tokens=self.truncate_prompt_tokens,
39 + truncation_side=self.truncation_side,
40 do_lower_case=encoder_config.get("do_lower_case", False),
41 add_special_tokens=self.add_special_tokens,
42 max_total_tokens_param="max_model_len",
@@ -61,6 +62,7 @@ class PoolingChatRequest(
62 max_total_tokens=model_config.max_model_len,
63 max_output_tokens=0,
64 truncate_prompt_tokens=self.truncate_prompt_tokens,
65 + truncation_side=self.truncation_side,
66 do_lower_case=encoder_config.get("do_lower_case", False),
67 add_special_tokens=self.add_special_tokens,
68 max_total_tokens_param="max_model_len",
@@ -88,6 +90,7 @@ class IOProcessorRequest(PoolingBasicRequestMixin, EncodingRequestMixin, Generic
90 max_total_tokens=model_config.max_model_len,
91 max_output_tokens=0,
92 truncate_prompt_tokens=self.truncate_prompt_tokens,
93 + truncation_side=self.truncation_side,
94 do_lower_case=encoder_config.get("do_lower_case", False),
95 add_special_tokens=not model_config.is_encoder_decoder,
96 max_total_tokens_param="max_model_len",
vllm/entrypoints/pooling/score/protocol.py
+2
@@ -30,6 +30,7 @@ class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin):
30 max_total_tokens=model_config.max_model_len,
31 max_output_tokens=0,
32 truncate_prompt_tokens=self.truncate_prompt_tokens,
33 + truncation_side=self.truncation_side,
34 do_lower_case=encoder_config.get("do_lower_case", False),
35 max_total_tokens_param="max_model_len",
36 )
@@ -105,6 +106,7 @@ class RerankRequest(PoolingBasicRequestMixin, ClassifyRequestMixin):
106 max_total_tokens=model_config.max_model_len,
107 max_output_tokens=0,
108 truncate_prompt_tokens=self.truncate_prompt_tokens,
109 + truncation_side=self.truncation_side,
110 do_lower_case=encoder_config.get("do_lower_case", False),
111 max_total_tokens_param="max_model_len",
112 )
vllm/entrypoints/pooling/typing.py
+2
@@ -15,6 +15,7 @@ from vllm.entrypoints.pooling.classify.protocol import (
15 ClassificationResponse,
16 )
17 from vllm.entrypoints.pooling.embed.protocol import (
18 + CohereEmbedRequest,
19 EmbeddingBytesResponse,
20 EmbeddingChatRequest,
21 EmbeddingCompletionRequest,
@@ -50,6 +51,7 @@ AnyPoolingRequest: TypeAlias = (
51 | IOProcessorRequest
52 | RerankRequest
53 | ScoreRequest
54 + | CohereEmbedRequest
55 )
56
57 AnyPoolingResponse: TypeAlias = (
vllm/renderers/params.py
+24 -2
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import dataclass, field
4 -from typing import TYPE_CHECKING, Any, TypeVar
4 +from typing import TYPE_CHECKING, Any, Literal, TypeVar
5
6 from vllm.exceptions import VLLMValidationError
7 from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt
@@ -153,6 +153,14 @@ class TokenizeParams:
153 - `-1` maps to `max_input_tokens`.
154 """
155
156 + truncation_side: Literal["left", "right"] | None = None
157 + """
158 + Which side to truncate from when ``truncate_prompt_tokens`` is active:
159 + - ``"right"`` keeps the first N tokens (truncate from the end).
160 + - ``"left"`` keeps the last N tokens (truncate from the start).
161 + - ``None`` falls back to the tokenizer default.
162 + """
163 +
164 do_lower_case: bool = False
165 """Whether to normalize text to lower case before tokenization."""
166
@@ -271,6 +279,7 @@ class TokenizeParams:
279 ),
280 pad_prompt_tokens=pad_prompt_tokens,
281 truncate_prompt_tokens=truncate_prompt_tokens,
282 + truncation_side=self.truncation_side,
283 do_lower_case=do_lower_case,
284 add_special_tokens=add_special_tokens,
285 needs_detokenization=needs_detokenization,
@@ -286,6 +295,16 @@ class TokenizeParams:
295 # while still failing `self._token_len_check` as expected by users
296 max_length = self.max_input_tokens + 1
297
298 + # Left-side truncation requires the full token sequence so we can
299 + # slice from the end in _token_truncation. Disable HF-level
300 + # truncation (which would incorrectly truncate from the right for
301 + # pooling models) and let _token_truncation handle it.
302 + if self.truncation_side == "left":
303 + return dict(
304 + truncation=False,
305 + add_special_tokens=self.add_special_tokens,
306 + )
307 +
308 return dict(
309 truncation=max_length is not None,
310 max_length=max_length,
@@ -375,7 +394,10 @@ class TokenizeParams:
394 if max_length == 0:
395 return tokens[:0]
396
378 - if getattr(tokenizer, "truncation_side", "left") == "left":
397 + side = self.truncation_side or (
398 + tokenizer.truncation_side if tokenizer is not None else None
399 + )
400 + if side == "left":
401 return tokens[-max_length:]
402
403 return tokens[:max_length]