| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Offline unit tests for `prompt_embeds` chat-completion content parts.""" |
| 5 | |
| 6 | from __future__ import annotations |
| 7 | |
| 8 | import inspect |
| 9 | import io |
| 10 | from typing import Final |
| 11 | from unittest import mock |
| 12 | |
| 13 | import pybase64 as base64 |
| 14 | import pytest |
| 15 | import regex as re |
| 16 | import torch |
| 17 | from transformers import AutoTokenizer |
| 18 | |
| 19 | from vllm.entrypoints.chat_utils import ( |
| 20 | _ENABLE_PROMPT_EMBEDS_ERROR, |
| 21 | _PROMPT_EMBEDS_MISSING_DATA_ERROR, |
| 22 | _RESERVED_PLACEHOLDER_IN_TEXT_ERROR, |
| 23 | MM_PARSER_MAP, |
| 24 | MODALITY_PLACEHOLDERS_MAP, |
| 25 | PROMPT_EMBEDS_PLACEHOLDER_TOKEN, |
| 26 | parse_chat_messages, |
| 27 | parse_chat_messages_async, |
| 28 | ) |
| 29 | from vllm.exceptions import VLLMValidationError |
| 30 | from vllm.renderers import TokenizeParams |
| 31 | from vllm.renderers.hf import ( |
| 32 | _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR, |
| 33 | _build_mixed_prompt_embeds, |
| 34 | _build_prompt_embeds_positions, |
| 35 | _build_prompt_embeds_updates, |
| 36 | _ensure_prompt_embeds_placeholder_token, |
| 37 | _expand_prompt_embeds_placeholders, |
| 38 | ) |
| 39 | |
| 40 | # Cover distinct tokenizer families: |
| 41 | # GPT2TokenizerFast (BPE, OpenAI-style) |
| 42 | # Qwen2TokenizerFast (SentencePiece BPE variant) |
| 43 | # BertTokenizerFast (WordPiece) |
| 44 | TOKENIZER_IDS: Final[list[str]] = [ |
| 45 | "openai-community/gpt2", |
| 46 | "Qwen/Qwen2.5-1.5B-Instruct", |
| 47 | "bert-base-uncased", |
| 48 | ] |
| 49 | |
| 50 | |
| 51 | @pytest.fixture(params=TOKENIZER_IDS, ids=TOKENIZER_IDS) |
| 52 | def tokenizer(request): |
| 53 | """A fresh tokenizer instance per tokenizer family.""" |
| 54 | return AutoTokenizer.from_pretrained(request.param) |
| 55 | |
| 56 | |
| 57 | # Minimal chat template that works with any tokenizer. Iterates |
| 58 | # `message.content` as either a string or a list of dicts (openai format). |
| 59 | _SIMPLE_CHAT_TEMPLATE: Final[str] = ( |
| 60 | "{% for m in messages %}" |
| 61 | "{% if m['content'] is string %}{{m['content']}}" |
| 62 | "{% else %}{% for p in m['content'] %}{{p['text']}}{% endfor %}" |
| 63 | "{% endif %}\n{% endfor %}" |
| 64 | ) |
| 65 | |
| 66 | |
| 67 | async def _maybe_await(fn, *args, **kwargs): |
| 68 | """Call *fn* and `await` the result if it's a coroutine.""" |
| 69 | result = fn(*args, **kwargs) |
| 70 | if inspect.iscoroutine(result): |
| 71 | result = await result |
| 72 | return result |
| 73 | |
| 74 | |
| 75 | # Parametrize over sync / async parse paths so every end-to-end test |
| 76 | # exercises both. |
| 77 | _PARSE_FUNCTIONS = [parse_chat_messages, parse_chat_messages_async] |
| 78 | |
| 79 | |
| 80 | @pytest.fixture(params=_PARSE_FUNCTIONS, ids=["sync", "async"]) |
| 81 | def parse_fn(request): |
| 82 | """Either the sync or async `parse_chat_messages` callable.""" |
| 83 | return request.param |
| 84 | |
| 85 | |
| 86 | def _encode_tensor(t: torch.Tensor) -> str: |
| 87 | buf = io.BytesIO() |
| 88 | torch.save(t, buf) |
| 89 | return base64.b64encode(buf.getvalue()).decode("utf-8") |
| 90 | |
| 91 | |
| 92 | _MOCK_HIDDEN_SIZE: Final[int] = 8 |
| 93 | _MOCK_DTYPE: Final[torch.dtype] = torch.float32 |
| 94 | |
| 95 | |
| 96 | def _make_mock_model_config(*, enable_prompt_embeds: bool = True) -> mock.MagicMock: |
| 97 | mc = mock.MagicMock() |
| 98 | mc.enable_prompt_embeds = enable_prompt_embeds |
| 99 | mc.multimodal_config = None |
| 100 | mc.allowed_local_media_path = None |
| 101 | mc.allowed_media_domains = None |
| 102 | # Test text-only code path in `MultiModalItemTracker.resolve_items`. |
| 103 | mc.is_multimodal_model = False |
| 104 | # `safe_load_prompt_embeds` pins each tensor to the model's hidden_size |
| 105 | # and dtype, so the mock must return concrete values. |
| 106 | mc.get_hidden_size.return_value = _MOCK_HIDDEN_SIZE |
| 107 | mc.dtype = _MOCK_DTYPE |
| 108 | return mc |
| 109 | |
| 110 | |
| 111 | def test_prompt_embeds_keys_registered(): |
| 112 | assert "prompt_embeds" in MODALITY_PLACEHOLDERS_MAP |
| 113 | assert MODALITY_PLACEHOLDERS_MAP["prompt_embeds"] == "<##PROMPT_EMBEDS##>" |
| 114 | assert "prompt_embeds" in MM_PARSER_MAP |
| 115 | |
| 116 | |
| 117 | def test_ensure_placeholder_token_is_single_token_and_idempotent(tokenizer): |
| 118 | """Ensure the placeholder token is a single token and that multiple calls to |
| 119 | "ensure" are idempotent, across all tokenizer families.""" |
| 120 | tid1 = _ensure_prompt_embeds_placeholder_token(tokenizer) |
| 121 | tid2 = _ensure_prompt_embeds_placeholder_token(tokenizer) |
| 122 | assert tid1 == tid2 |
| 123 | |
| 124 | ids = tokenizer.encode(PROMPT_EMBEDS_PLACEHOLDER_TOKEN, add_special_tokens=False) |
| 125 | assert ids == [tid1] |
| 126 | |
| 127 | # Repeating it in a string N times must produce exactly that many tokens. |
| 128 | N = 5 |
| 129 | ids_rep = tokenizer.encode( |
| 130 | PROMPT_EMBEDS_PLACEHOLDER_TOKEN * N, add_special_tokens=False |
| 131 | ) |
| 132 | assert ids_rep == [tid1] * N |
| 133 | |
| 134 | |
| 135 | def test_parse_chat_messages_openai_format(): |
| 136 | NUM_TOKENS = 3 |
| 137 | t = torch.randn(NUM_TOKENS, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE) |
| 138 | b64 = _encode_tensor(t) |
| 139 | mc = _make_mock_model_config() |
| 140 | |
| 141 | messages = [ |
| 142 | { |
| 143 | "role": "user", |
| 144 | "content": [ |
| 145 | {"type": "text", "text": "Hello "}, |
| 146 | {"type": "prompt_embeds", "data": b64}, |
| 147 | {"type": "text", "text": " world"}, |
| 148 | ], |
| 149 | } |
| 150 | ] |
| 151 | conv, mm_data, _ = parse_chat_messages( |
| 152 | messages, |
| 153 | mc, |
| 154 | content_format="openai", |
| 155 | ) |
| 156 | # The middle content part is rewritten to a single placeholder-token |
| 157 | # sentinel. |
| 158 | texts = [p["text"] for p in conv[0]["content"]] |
| 159 | assert texts == [ |
| 160 | "Hello ", |
| 161 | PROMPT_EMBEDS_PLACEHOLDER_TOKEN, |
| 162 | " world", |
| 163 | ] |
| 164 | assert mm_data is not None and "prompt_embeds" in mm_data |
| 165 | assert torch.equal(mm_data["prompt_embeds"][0], t) |
| 166 | |
| 167 | |
| 168 | # Each layout entry is one content part: |
| 169 | # ("text", "A") -> {"type": "text", "text": "A"} |
| 170 | # ("embed", N) -> {"type": "prompt_embeds", "data": <base64 of (N, H) tensor>} |
| 171 | @pytest.mark.parametrize( |
| 172 | "layout", |
| 173 | [ |
| 174 | # Case: Single embed only. |
| 175 | [("embed", 2)], |
| 176 | # Case: Embed at the start of the message. |
| 177 | [("embed", 3), ("text", "B")], |
| 178 | # Case: Embed at the end of the message. |
| 179 | [("text", "A"), ("embed", 1)], |
| 180 | # Case: Embed sandwiched between text spans. |
| 181 | [("text", "A"), ("embed", 2), ("text", "B")], |
| 182 | # Case: Multiple embeds with text in between. |
| 183 | [("text", "A"), ("embed", 2), ("text", "B"), ("embed", 3)], |
| 184 | # Case: Adjacent embeds with no separating text. |
| 185 | [("embed", 1), ("embed", 2)], |
| 186 | # Case: Multiple text spans before a trailing embed. |
| 187 | [("text", "A"), ("text", "B"), ("embed", 1)], |
| 188 | # Case: Long-ish run mixing both kinds. |
| 189 | [ |
| 190 | ("text", "head"), |
| 191 | ("embed", 4), |
| 192 | ("text", "mid"), |
| 193 | ("embed", 1), |
| 194 | ("embed", 2), |
| 195 | ("text", "tail"), |
| 196 | ], |
| 197 | ], |
| 198 | ids=[ |
| 199 | "single-embed", |
| 200 | "embed-then-text", |
| 201 | "text-then-embed", |
| 202 | "text-embed-text", |
| 203 | "text-embed-text-embed", |
| 204 | "adjacent-embeds", |
| 205 | "text-text-embed", |
| 206 | "long-mixed-run", |
| 207 | ], |
| 208 | ) |
| 209 | @pytest.mark.parametrize( |
| 210 | "interleave_mm_strings", |
| 211 | # `None`: text-only path where `multimodal_config` is absent. |
| 212 | # `False`: non-interleave multimodal path (the common default). |
| 213 | # `True`: sentinel-substitution interleave path. |
| 214 | # All three must preserve the request ordering of prompt_embeds |
| 215 | # relative to surrounding text because prompt_embeds are spliced at the |
| 216 | # token offset during rendering. |
| 217 | [None, False, True], |
| 218 | ids=["text-only", "interleave-off", "interleave-on"], |
| 219 | ) |
| 220 | def test_parse_chat_messages_string_format_preserves_position( |
| 221 | layout, interleave_mm_strings |
| 222 | ): |
| 223 | mc = _make_mock_model_config() |
| 224 | if interleave_mm_strings is not None: |
| 225 | mm_cfg = mock.MagicMock() |
| 226 | mm_cfg.interleave_mm_strings = interleave_mm_strings |
| 227 | mc.multimodal_config = mm_cfg |
| 228 | |
| 229 | content: list[dict] = [] |
| 230 | expected_parts: list[str] = [] |
| 231 | expected_embeds: list[torch.Tensor] = [] |
| 232 | for kind, value in layout: |
| 233 | if kind == "text": |
| 234 | content.append({"type": "text", "text": value}) |
| 235 | expected_parts.append(value) |
| 236 | else: # prompt embeds |
| 237 | num_tokens = value |
| 238 | t = torch.randn(num_tokens, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE) |
| 239 | expected_embeds.append(t) |
| 240 | content.append({"type": "prompt_embeds", "data": _encode_tensor(t)}) |
| 241 | # Parser emits ONE sentinel per part. |
| 242 | expected_parts.append(PROMPT_EMBEDS_PLACEHOLDER_TOKEN) |
| 243 | |
| 244 | messages = [{"role": "user", "content": content}] |
| 245 | conv, mm_data, _ = parse_chat_messages( |
| 246 | messages, |
| 247 | mc, |
| 248 | content_format="string", |
| 249 | ) |
| 250 | |
| 251 | assert conv[0]["content"] == "\n".join(expected_parts) |
| 252 | assert mm_data is not None and "prompt_embeds" in mm_data |
| 253 | assert len(mm_data["prompt_embeds"]) == len(expected_embeds) |
| 254 | for got, want in zip(mm_data["prompt_embeds"], expected_embeds, strict=True): |
| 255 | assert torch.equal(got, want) |
| 256 | |
| 257 | |
| 258 | def test_parse_chat_messages_requires_flag(): |
| 259 | t = torch.randn(2, 4) |
| 260 | b64 = _encode_tensor(t) |
| 261 | mc = _make_mock_model_config(enable_prompt_embeds=False) |
| 262 | |
| 263 | messages = [ |
| 264 | { |
| 265 | "role": "user", |
| 266 | "content": [{"type": "prompt_embeds", "data": b64}], |
| 267 | } |
| 268 | ] |
| 269 | with pytest.raises(VLLMValidationError, match=_ENABLE_PROMPT_EMBEDS_ERROR): |
| 270 | parse_chat_messages( |
| 271 | messages, |
| 272 | mc, |
| 273 | content_format="openai", |
| 274 | ) |
| 275 | |
| 276 | |
| 277 | def test_parse_chat_messages_rejects_missing_data(): |
| 278 | # `data` is marked `Required` on `ChatCompletionContentPartPromptEmbedsParam`; |
| 279 | # malformed requests without `data` must surface a clear validation error |
| 280 | # rather than being silently dropped. |
| 281 | mc = _make_mock_model_config() |
| 282 | messages = [ |
| 283 | { |
| 284 | "role": "user", |
| 285 | "content": [{"type": "prompt_embeds"}], # no `data` |
| 286 | } |
| 287 | ] |
| 288 | with pytest.raises(VLLMValidationError, match=_PROMPT_EMBEDS_MISSING_DATA_ERROR): |
| 289 | parse_chat_messages( |
| 290 | messages, |
| 291 | mc, |
| 292 | content_format="openai", |
| 293 | ) |
| 294 | |
| 295 | |
| 296 | # Reserved placeholder guard: when `enable_prompt_embeds=True` the tokenizer is |
| 297 | # mutated to make `<prompt_embeds>` a single unsplittable token. Any user text |
| 298 | # containing that literal sequence would tokenize to the same sentinel ID and |
| 299 | # be mistaken for a splice point, so we reject it at parse time. |
| 300 | _PLACEHOLDER_ERROR_PATTERN: Final[str] = re.sub( |
| 301 | r"\\{[^}]*\\}", ".*", re.escape(_RESERVED_PLACEHOLDER_IN_TEXT_ERROR) |
| 302 | ) |
| 303 | |
| 304 | |
| 305 | @pytest.mark.parametrize( |
| 306 | "content", |
| 307 | [ |
| 308 | # Case: Top-level string content (wrapped as a single text part). |
| 309 | f"hello {PROMPT_EMBEDS_PLACEHOLDER_TOKEN} world", |
| 310 | # Case: List with a typed text part containing the placeholder. |
| 311 | [{"type": "text", "text": f"leading {PROMPT_EMBEDS_PLACEHOLDER_TOKEN}"}], |
| 312 | # Case: List with a plain-string part (no wrapping dict). |
| 313 | [f"raw string {PROMPT_EMBEDS_PLACEHOLDER_TOKEN}"], |
| 314 | ], |
| 315 | ids=["top-level-string", "typed-text-part", "plain-string-part"], |
| 316 | ) |
| 317 | def test_parse_chat_messages_rejects_placeholder_in_user_text(content): |
| 318 | mc = _make_mock_model_config() # enable_prompt_embeds=True by default |
| 319 | messages = [{"role": "user", "content": content}] |
| 320 | with pytest.raises(VLLMValidationError, match=_PLACEHOLDER_ERROR_PATTERN): |
| 321 | parse_chat_messages(messages, mc, content_format="openai") |
| 322 | |
| 323 | |
| 324 | def test_parse_chat_messages_allows_placeholder_in_text_when_feature_disabled(): |
| 325 | # When `enable_prompt_embeds=False` the tokenizer is never mutated, so the |
| 326 | # literal `<prompt_embeds>` is just ordinary text and must pass through. |
| 327 | mc = _make_mock_model_config(enable_prompt_embeds=False) |
| 328 | messages = [ |
| 329 | { |
| 330 | "role": "user", |
| 331 | "content": f"benign mention of {PROMPT_EMBEDS_PLACEHOLDER_TOKEN} here", |
| 332 | } |
| 333 | ] |
| 334 | conv, mm_data, _ = parse_chat_messages(messages, mc, content_format="openai") |
| 335 | assert mm_data is None or "prompt_embeds" not in mm_data |
| 336 | # Text reaches the rendered conversation unchanged. |
| 337 | texts = [p["text"] for p in conv[0]["content"]] |
| 338 | assert PROMPT_EMBEDS_PLACEHOLDER_TOKEN in "".join(texts) |
| 339 | |
| 340 | |
| 341 | # Token-stream spec: ints are regular token IDs, tuples `(N,)` expand to |
| 342 | # a placeholder span of length N (creates corresponding `(N, H)` tensor). |
| 343 | # `expected` lists the `(start_idx, length)` pairs that |
| 344 | # `_build_prompt_embeds_positions` should return. |
| 345 | @pytest.mark.parametrize( |
| 346 | "stream, expected", |
| 347 | [ |
| 348 | # Case: Single run in the middle. |
| 349 | ([10, 20, (3,), 30], [(2, 3)]), |
| 350 | # Case: Single run at the start. |
| 351 | ([(2,), 10, 20], [(0, 2)]), |
| 352 | # Case: Single run at the end. |
| 353 | ([10, 20, (4,)], [(2, 4)]), |
| 354 | # Case: Two runs with tokens between. |
| 355 | ([1, (2,), 2, 3, (3,), 4], [(1, 2), (5, 3)]), |
| 356 | # Case: Adjacent runs (no separating tokens). |
| 357 | ([(1,), (2,)], [(0, 1), (1, 2)]), |
| 358 | # Case: Three runs. |
| 359 | ([5, (2,), 6, (1,), 7, (3,), 8], [(1, 2), (4, 1), (6, 3)]), |
| 360 | ], |
| 361 | ids=[ |
| 362 | "single-middle", |
| 363 | "single-start", |
| 364 | "single-end", |
| 365 | "two-runs-separated", |
| 366 | "two-runs-adjacent", |
| 367 | "three-runs", |
| 368 | ], |
| 369 | ) |
| 370 | def test_build_positions(tokenizer, stream, expected): |
| 371 | H = 4 |
| 372 | tid = _ensure_prompt_embeds_placeholder_token(tokenizer) |
| 373 | tensors: list[torch.Tensor] = [] |
| 374 | token_ids: list[int] = [] |
| 375 | for item in stream: |
| 376 | if isinstance(item, tuple): |
| 377 | length = item[0] |
| 378 | tensors.append(torch.randn(length, H)) |
| 379 | token_ids.extend([tid] * length) |
| 380 | else: |
| 381 | token_ids.append(item) |
| 382 | mm_updates = _build_prompt_embeds_updates(tensors, tid) |
| 383 | positions = _build_prompt_embeds_positions(token_ids, len(tensors), mm_updates) |
| 384 | assert positions == expected |
| 385 | |
| 386 | |
| 387 | def test_build_positions_length_mismatch(tokenizer): |
| 388 | N1, H1 = 2, 4 |
| 389 | N2, H2 = 3, 4 |
| 390 | tid = _ensure_prompt_embeds_placeholder_token(tokenizer) |
| 391 | # 2 tensors expected but only a single placeholder run in the token |
| 392 | # stream (simulating dropping the second one). |
| 393 | tensors = [torch.randn(N1, H1), torch.randn(N2, H2)] |
| 394 | token_ids = [1, tid, tid, 2, 3] |
| 395 | mm_updates = _build_prompt_embeds_updates(tensors, tid) |
| 396 | # The error constant is a `str.format` template, escape it and turn |
| 397 | # the `{field}` placeholders into `.*` so it matches any substitution. |
| 398 | pattern = re.sub( |
| 399 | r"\\{[^}]*\\}", ".*", re.escape(_PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR) |
| 400 | ) |
| 401 | with pytest.raises(ValueError, match=pattern): |
| 402 | _build_prompt_embeds_positions(token_ids, len(tensors), mm_updates) |
| 403 | |
| 404 | |
| 405 | # ints = regular token IDs (any value) |
| 406 | # (N,) = embed span of length N |
| 407 | @pytest.mark.parametrize( |
| 408 | "stream", |
| 409 | [ |
| 410 | [10, 20, (3,), 30], |
| 411 | [(2,), 10, 20], |
| 412 | [10, 20, (4,)], |
| 413 | [1, (2,), 2, 3, (3,), 4], |
| 414 | [(1,), (2,)], |
| 415 | [5, (2,), 6, (1,), 7, (3,), 8], |
| 416 | ], |
| 417 | ids=[ |
| 418 | "single-middle", |
| 419 | "single-start", |
| 420 | "single-end", |
| 421 | "two-spans-separated", |
| 422 | "two-spans-adjacent", |
| 423 | "three-spans", |
| 424 | ], |
| 425 | ) |
| 426 | def test_build_mixed_prompt_embeds(stream): |
| 427 | H = 8 |
| 428 | _PLACEHOLDER = 0 # sentinel for embed positions in token_ids |
| 429 | |
| 430 | tensors: list[torch.Tensor] = [] |
| 431 | token_ids: list[int] = [] |
| 432 | positions: list[tuple[int, int]] = [] |
| 433 | cursor = 0 |
| 434 | for item in stream: |
| 435 | if isinstance(item, tuple): |
| 436 | length = item[0] |
| 437 | tensors.append(torch.randn(length, H)) |
| 438 | positions.append((cursor, length)) |
| 439 | token_ids.extend([_PLACEHOLDER] * length) |
| 440 | cursor += length |
| 441 | else: |
| 442 | token_ids.append(item) |
| 443 | cursor += 1 |
| 444 | |
| 445 | embeds, mask = _build_mixed_prompt_embeds(token_ids, tensors, positions) |
| 446 | |
| 447 | assert embeds.shape == (len(token_ids), H) |
| 448 | assert len(mask) == len(token_ids) |
| 449 | |
| 450 | # Mask: False exactly at embed positions, True everywhere else. |
| 451 | expected_mask = torch.ones(len(token_ids), dtype=torch.bool) |
| 452 | for start, length in positions: |
| 453 | expected_mask[start : start + length] = False |
| 454 | assert mask == expected_mask.tolist() |
| 455 | |
| 456 | # Embed rows match input tensors at the right positions. |
| 457 | for tensor, (start, length) in zip(tensors, positions): |
| 458 | assert torch.equal(embeds[start : start + length], tensor) |
| 459 | |
| 460 | # Non-embed positions remain zero-filled. |
| 461 | assert torch.all(embeds[expected_mask] == 0) |
| 462 | |
| 463 | |
| 464 | # End-to-end tests: each runs both sync and async parse paths via the |
| 465 | # `parse_fn` fixture. |
| 466 | |
| 467 | |
| 468 | @pytest.mark.asyncio |
| 469 | @pytest.mark.parametrize("role", ["user", "system"]) |
| 470 | async def test_end_to_end_expand_and_build(tokenizer, parse_fn, role): |
| 471 | """Full renderer pipeline: parse -> chat template -> expand -> locate |
| 472 | -> build mixed prompt, across tokenizers, roles, and sync/async.""" |
| 473 | tokenizer.chat_template = _SIMPLE_CHAT_TEMPLATE |
| 474 | tid = _ensure_prompt_embeds_placeholder_token(tokenizer) |
| 475 | |
| 476 | LEN_A, LEN_B = 3, 2 |
| 477 | t_a = torch.randn(LEN_A, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE) |
| 478 | t_b = torch.randn(LEN_B, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE) |
| 479 | NUM_TENSORS = 2 |
| 480 | |
| 481 | mc = _make_mock_model_config() |
| 482 | |
| 483 | messages = [ |
| 484 | { |
| 485 | "role": role, |
| 486 | "content": [ |
| 487 | {"type": "text", "text": "Hello "}, |
| 488 | {"type": "prompt_embeds", "data": _encode_tensor(t_a)}, |
| 489 | {"type": "text", "text": " world "}, |
| 490 | {"type": "prompt_embeds", "data": _encode_tensor(t_b)}, |
| 491 | {"type": "text", "text": "!"}, |
| 492 | ], |
| 493 | } |
| 494 | ] |
| 495 | |
| 496 | conv, mm_data, _ = await _maybe_await( |
| 497 | parse_fn, messages, mc, content_format="openai" |
| 498 | ) |
| 499 | tensors = list(mm_data["prompt_embeds"]) |
| 500 | assert len(tensors) == NUM_TENSORS |
| 501 | |
| 502 | # Tokenize: each prompt_embeds part becomes 1 placeholder token. |
| 503 | # `return_dict=False` to get a flat `list[int]` on transformers v5 |
| 504 | # (where the default flipped to True and yields a `BatchEncoding` dict). |
| 505 | token_ids = tokenizer.apply_chat_template(conv, tokenize=True, return_dict=False) |
| 506 | assert sum(t == tid for t in token_ids) == NUM_TENSORS |
| 507 | |
| 508 | # Expand, locate, and build. |
| 509 | mm_updates = _build_prompt_embeds_updates(tensors, tid) |
| 510 | expanded = _expand_prompt_embeds_placeholders(token_ids, mm_updates) |
| 511 | assert len(expanded) == len(token_ids) + LEN_A + LEN_B - NUM_TENSORS |
| 512 | |
| 513 | positions = _build_prompt_embeds_positions(expanded, len(tensors), mm_updates) |
| 514 | assert positions[0][1] == LEN_A |
| 515 | assert positions[1][1] == LEN_B |
| 516 | |
| 517 | embeds, mask = _build_mixed_prompt_embeds(expanded, tensors, positions) |
| 518 | assert embeds.shape == (len(expanded), _MOCK_HIDDEN_SIZE) |
| 519 | assert mask.count(False) == LEN_A + LEN_B |
| 520 | assert torch.equal(embeds[positions[0][0] : positions[0][0] + LEN_A], t_a) |
| 521 | assert torch.equal(embeds[positions[1][0] : positions[1][0] + LEN_B], t_b) |
| 522 | |
| 523 | |
| 524 | @pytest.mark.asyncio |
| 525 | async def test_end_to_end_multi_message_conversation(tokenizer, parse_fn): |
| 526 | """Full pipeline with prompt_embeds spread across system + user messages, |
| 527 | verifying ordering and positioning in the final token stream.""" |
| 528 | tokenizer.chat_template = _SIMPLE_CHAT_TEMPLATE |
| 529 | tid = _ensure_prompt_embeds_placeholder_token(tokenizer) |
| 530 | |
| 531 | LEN_SYS, LEN_USR = 4, 3 |
| 532 | t_sys = torch.randn(LEN_SYS, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE) |
| 533 | t_usr = torch.randn(LEN_USR, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE) |
| 534 | NUM_TENSORS = 2 # t_sys and t_usr. |
| 535 | |
| 536 | mc = _make_mock_model_config() |
| 537 | |
| 538 | messages = [ |
| 539 | { |
| 540 | "role": "system", |
| 541 | "content": [ |
| 542 | {"type": "text", "text": "You are helpful."}, |
| 543 | {"type": "prompt_embeds", "data": _encode_tensor(t_sys)}, |
| 544 | ], |
| 545 | }, |
| 546 | { |
| 547 | "role": "user", |
| 548 | "content": [ |
| 549 | {"type": "prompt_embeds", "data": _encode_tensor(t_usr)}, |
| 550 | {"type": "text", "text": "Summarize."}, |
| 551 | ], |
| 552 | }, |
| 553 | ] |
| 554 | |
| 555 | conv, mm_data, _ = await _maybe_await( |
| 556 | parse_fn, messages, mc, content_format="openai" |
| 557 | ) |
| 558 | tensors = list(mm_data["prompt_embeds"]) |
| 559 | assert len(tensors) == NUM_TENSORS |
| 560 | |
| 561 | # Tokenize, expand, locate, and build. |
| 562 | # `return_dict=False` to get a flat `list[int]` on transformers v5 |
| 563 | # (where the default flipped to True and yields a `BatchEncoding` dict). |
| 564 | token_ids = tokenizer.apply_chat_template(conv, tokenize=True, return_dict=False) |
| 565 | mm_updates = _build_prompt_embeds_updates(tensors, tid) |
| 566 | expanded = _expand_prompt_embeds_placeholders(token_ids, mm_updates) |
| 567 | positions = _build_prompt_embeds_positions(expanded, len(tensors), mm_updates) |
| 568 | |
| 569 | assert positions[0][1] == LEN_SYS |
| 570 | assert positions[1][1] == LEN_USR |
| 571 | # System embed must appear before user embed in the token stream. |
| 572 | assert positions[0][0] < positions[1][0] |
| 573 | |
| 574 | embeds, mask = _build_mixed_prompt_embeds(expanded, tensors, positions) |
| 575 | assert embeds.shape == (len(expanded), _MOCK_HIDDEN_SIZE) |
| 576 | assert mask.count(False) == LEN_SYS + LEN_USR |
| 577 | assert torch.equal(embeds[positions[0][0] : positions[0][0] + LEN_SYS], t_sys) |
| 578 | assert torch.equal(embeds[positions[1][0] : positions[1][0] + LEN_USR], t_usr) |
| 579 | |
| 580 | |
| 581 | def test_truncation_keeps_the_mixed_mask_aligned_with_the_prompt(): |
| 582 | """`prompt_is_token_ids` must be truncated with `prompt_token_ids`. |
| 583 | |
| 584 | The mask is positional: the engine pairs it element-wise with the prompt |
| 585 | tokens to decide which positions read a real token id and which read a row |
| 586 | of `prompt_embeds`. Truncating the prompt without truncating the mask |
| 587 | leaves the two describing different positions, so embed placeholders stop |
| 588 | being recognised as such. |
| 589 | """ |
| 590 | hidden_size = 8 |
| 591 | placeholder_token_id = 0 |
| 592 | num_head, embed_len, num_tail = 10, 5, 5 |
| 593 | |
| 594 | token_ids = ( |
| 595 | list(range(1, num_head + 1)) |
| 596 | + [placeholder_token_id] * embed_len |
| 597 | + list(range(1, num_tail + 1)) |
| 598 | ) |
| 599 | tensors = [torch.randn(embed_len, hidden_size)] |
| 600 | positions = [(num_head, embed_len)] |
| 601 | embeds, mask = _build_mixed_prompt_embeds(token_ids, tensors, positions) |
| 602 | |
| 603 | prompt = { |
| 604 | "prompt_token_ids": token_ids, |
| 605 | "prompt_embeds": embeds, |
| 606 | "prompt_is_token_ids": mask, |
| 607 | } |
| 608 | keep = 8 |
| 609 | tok_params = TokenizeParams( |
| 610 | max_total_tokens=100, |
| 611 | truncate_prompt_tokens=keep, |
| 612 | truncation_side="left", |
| 613 | ) |
| 614 | |
| 615 | result = tok_params.apply_post_tokenization(None, prompt) |
| 616 | |
| 617 | # Keeping the last 8 of 20 positions starts inside the embed span, so the |
| 618 | # first 3 surviving positions are embed rows and the rest are real tokens. |
| 619 | # Comparing contents rather than lengths is what pins the alignment: a |
| 620 | # mask reduced from the wrong side has the right length and still |
| 621 | # describes the wrong positions. |
| 622 | assert result["prompt_is_token_ids"] == [False] * 3 + [True] * num_tail |
| 623 | assert result["prompt_token_ids"] == token_ids[-keep:] |
| 624 | assert torch.equal(result["prompt_embeds"], embeds[-keep:]) |