| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for renderer-level token-offset behavior. |
| 4 | |
| 5 | These exercise ``_tokenize_prompt`` (offset extraction + capability/MM |
| 6 | gating) and the ``_tokenize_prompt -> _process_tokens -> TokensInput`` |
| 7 | forwarding chain. Endpoint-level coverage lives in |
| 8 | ``tests/entrypoints/scale_out/render/test_render.py``. |
| 9 | """ |
| 10 | |
| 11 | import asyncio |
| 12 | |
| 13 | import pytest |
| 14 | |
| 15 | from vllm.renderers.params import TokenizeParams |
| 16 | |
| 17 | |
| 18 | @pytest.fixture |
| 19 | def fast_tokenizer(): |
| 20 | """gpt2 ships a Fast tokenizer; use it to test the offsets happy path.""" |
| 21 | from transformers import AutoTokenizer |
| 22 | |
| 23 | return AutoTokenizer.from_pretrained("openai-community/gpt2") |
| 24 | |
| 25 | |
| 26 | def _make_base_renderer_with(tokenizer): |
| 27 | """Build a minimal BaseRenderer subclass that exposes the tokenizer so we |
| 28 | can call ``_tokenize_prompt`` directly. BaseRenderer is abstract because of |
| 29 | ``render_messages``; we just need a stub.""" |
| 30 | from vllm.renderers.base import BaseRenderer |
| 31 | |
| 32 | class _StubRenderer(BaseRenderer): |
| 33 | def __init__(self, tok): |
| 34 | # Bypass BaseRenderer.__init__ — we don't need a VllmConfig. |
| 35 | from vllm.utils.async_utils import make_async |
| 36 | |
| 37 | self.tokenizer = tok |
| 38 | self._executor = None |
| 39 | # Mirror BaseRenderer.__init__: the async path offloads the sync |
| 40 | # ``_tokenize_prompt`` to a thread pool. |
| 41 | self._tokenize_prompt_async = make_async(self._tokenize_prompt) |
| 42 | self.mm_processor = None |
| 43 | |
| 44 | def get_tokenizer(self): |
| 45 | return self.tokenizer |
| 46 | |
| 47 | def _can_produce_offsets(self): |
| 48 | # Mirror HfRenderer: offsets only for fast tokenizers. |
| 49 | return self.tokenizer is not None and self.tokenizer.is_fast |
| 50 | |
| 51 | def render_messages(self, messages, params): # pragma: no cover |
| 52 | raise NotImplementedError |
| 53 | |
| 54 | return _StubRenderer(tokenizer) |
| 55 | |
| 56 | |
| 57 | class _OffsetTokenizer: |
| 58 | """Small deterministic fast tokenizer for pre-tokenization tests.""" |
| 59 | |
| 60 | is_fast = True |
| 61 | max_chars_per_token = 1 |
| 62 | truncation_side = "right" |
| 63 | pad_token_id = 0 |
| 64 | |
| 65 | def __call__(self, text, **kwargs): |
| 66 | return { |
| 67 | "input_ids": list(range(len(text))), |
| 68 | "offset_mapping": [(i, i + 1) for i in range(len(text))], |
| 69 | } |
| 70 | |
| 71 | |
| 72 | class TestTokenizePromptOffsets: |
| 73 | def test_fast_tokenizer_with_flag_returns_offsets(self, fast_tokenizer): |
| 74 | renderer = _make_base_renderer_with(fast_tokenizer) |
| 75 | params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) |
| 76 | prompt = {"prompt": "Hello, world."} |
| 77 | |
| 78 | result = renderer._tokenize_prompt(prompt, params) |
| 79 | |
| 80 | assert "prompt_token_ids" in result |
| 81 | offsets = result["prompt_token_offsets"] |
| 82 | assert offsets is not None |
| 83 | # Length must match the token sequence, and each (start, end) is an |
| 84 | # ordered pair within the source text. |
| 85 | assert len(offsets) == len(result["prompt_token_ids"]) |
| 86 | text_len = len("Hello, world.") |
| 87 | for s, e in offsets: |
| 88 | assert isinstance(s, int) and isinstance(e, int) |
| 89 | assert 0 <= s <= e <= text_len |
| 90 | |
| 91 | def test_base_renderer_without_override_yields_no_offsets(self, fast_tokenizer): |
| 92 | """A renderer that does not override ``_can_produce_offsets`` never |
| 93 | emits offsets, even with a fast tokenizer and the flag set. This locks |
| 94 | in the base-default-False / subclass-override design.""" |
| 95 | from vllm.renderers.base import BaseRenderer |
| 96 | |
| 97 | class _BareRenderer(BaseRenderer): |
| 98 | def __init__(self, tok): |
| 99 | self.tokenizer = tok |
| 100 | self._executor = None |
| 101 | self.mm_processor = None |
| 102 | |
| 103 | def get_tokenizer(self): |
| 104 | return self.tokenizer |
| 105 | |
| 106 | def render_messages(self, messages, params): # pragma: no cover |
| 107 | raise NotImplementedError |
| 108 | |
| 109 | renderer = _BareRenderer(fast_tokenizer) |
| 110 | params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) |
| 111 | |
| 112 | result = renderer._tokenize_prompt({"prompt": "Hello, world."}, params) |
| 113 | |
| 114 | assert "prompt_token_offsets" not in result |
| 115 | |
| 116 | def test_default_flag_no_offsets(self, fast_tokenizer): |
| 117 | renderer = _make_base_renderer_with(fast_tokenizer) |
| 118 | params = TokenizeParams(max_total_tokens=None) # flag defaults False |
| 119 | |
| 120 | result = renderer._tokenize_prompt({"prompt": "Hello, world."}, params) |
| 121 | |
| 122 | # Field must be absent (not None) so TokensInput serialization stays |
| 123 | # minimal for existing consumers. |
| 124 | assert "prompt_token_offsets" not in result |
| 125 | |
| 126 | def test_slow_tokenizer_with_flag_no_offsets(self, fast_tokenizer): |
| 127 | """Force is_fast=False to simulate a Slow tokenizer: the flag is set |
| 128 | but offsets must not be returned because it cannot produce them.""" |
| 129 | from unittest.mock import PropertyMock, patch |
| 130 | |
| 131 | renderer = _make_base_renderer_with(fast_tokenizer) |
| 132 | params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) |
| 133 | |
| 134 | with patch.object( |
| 135 | type(fast_tokenizer), |
| 136 | "is_fast", |
| 137 | new_callable=PropertyMock, |
| 138 | return_value=False, |
| 139 | ): |
| 140 | result = renderer._tokenize_prompt({"prompt": "Hello, world."}, params) |
| 141 | |
| 142 | assert "prompt_token_offsets" not in result |
| 143 | |
| 144 | @pytest.mark.parametrize("mm_key", ["multi_modal_data", "multi_modal_uuids"]) |
| 145 | def test_multimodal_with_flag_no_offsets(self, fast_tokenizer, mm_key): |
| 146 | """Offsets index the text prompt, which is meaningless once multimodal |
| 147 | data is interleaved, so they are suppressed when MM inputs are present.""" |
| 148 | renderer = _make_base_renderer_with(fast_tokenizer) |
| 149 | params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) |
| 150 | prompt = {"prompt": "Hello.", mm_key: {"image": ["x"]}} |
| 151 | |
| 152 | result = renderer._tokenize_prompt(prompt, params) |
| 153 | |
| 154 | assert "prompt_token_offsets" not in result |
| 155 | |
| 156 | @pytest.mark.asyncio |
| 157 | async def test_tokenize_prompt_async_returns_offsets(self, fast_tokenizer): |
| 158 | """The async path offloads the sync tokenizer; it must yield the same |
| 159 | offsets as the sync path.""" |
| 160 | renderer = _make_base_renderer_with(fast_tokenizer) |
| 161 | params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) |
| 162 | |
| 163 | result = await renderer._tokenize_prompt_async( |
| 164 | {"prompt": "Hello, world."}, params |
| 165 | ) |
| 166 | |
| 167 | offsets = result["prompt_token_offsets"] |
| 168 | assert offsets is not None |
| 169 | assert len(offsets) == len(result["prompt_token_ids"]) |
| 170 | |
| 171 | |
| 172 | class TestProcessTokensForwardsOffsets: |
| 173 | """Tests that the ``_tokenize_prompt -> _process_tokens -> TokensInput`` |
| 174 | chain carries ``prompt_token_offsets`` through to the engine input. |
| 175 | ``_process_tokens`` rebuilds the engine input from scratch, so it must |
| 176 | copy the field explicitly. The sync and async variants are independent |
| 177 | implementations, so both are checked. |
| 178 | """ |
| 179 | |
| 180 | def test_sync_forwards_offsets_to_engine_input(self, fast_tokenizer): |
| 181 | renderer = _make_base_renderer_with(fast_tokenizer) |
| 182 | params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) |
| 183 | |
| 184 | tokens_prompt = renderer._tokenize_prompt({"prompt": "Hello, world."}, params) |
| 185 | # Sanity: offsets must reach the TokensPrompt, else this guards the |
| 186 | # wrong layer. |
| 187 | expected = tokens_prompt["prompt_token_offsets"] |
| 188 | |
| 189 | engine_input = renderer._process_tokens(tokens_prompt) |
| 190 | |
| 191 | assert engine_input["prompt_token_offsets"] == expected |
| 192 | |
| 193 | @pytest.mark.asyncio |
| 194 | async def test_async_forwards_offsets_to_engine_input(self, fast_tokenizer): |
| 195 | renderer = _make_base_renderer_with(fast_tokenizer) |
| 196 | params = TokenizeParams(max_total_tokens=None, return_token_offsets=True) |
| 197 | |
| 198 | tokens_prompt = await renderer._tokenize_prompt_async( |
| 199 | {"prompt": "Hello, world."}, params |
| 200 | ) |
| 201 | expected = tokens_prompt["prompt_token_offsets"] |
| 202 | |
| 203 | engine_input = await renderer._process_tokens_async(tokens_prompt) |
| 204 | |
| 205 | assert engine_input["prompt_token_offsets"] == expected |
| 206 | |
| 207 | def test_no_offsets_forwarded_when_flag_off(self, fast_tokenizer): |
| 208 | renderer = _make_base_renderer_with(fast_tokenizer) |
| 209 | params = TokenizeParams(max_total_tokens=None) # flag defaults False |
| 210 | |
| 211 | tokens_prompt = renderer._tokenize_prompt({"prompt": "Hello, world."}, params) |
| 212 | assert "prompt_token_offsets" not in tokens_prompt |
| 213 | |
| 214 | engine_input = renderer._process_tokens(tokens_prompt) |
| 215 | |
| 216 | assert "prompt_token_offsets" not in engine_input |
| 217 | |
| 218 | |
| 219 | class TestTruncationKeepsOffsetsAligned: |
| 220 | """``prompt_token_offsets`` runs parallel to ``prompt_token_ids``, and both |
| 221 | ``TokensPrompt`` and the render API's ``GenerateRequest`` document that the |
| 222 | two have equal length. Only ``prompt_token_ids`` was being truncated. |
| 223 | """ |
| 224 | |
| 225 | @pytest.mark.parametrize("side", ["left", "right"]) |
| 226 | def test_explicit_truncation_side_truncates_offsets(self, fast_tokenizer, side): |
| 227 | renderer = _make_base_renderer_with(fast_tokenizer) |
| 228 | text = "The quick brown fox jumps over the lazy dog." |
| 229 | keep = 4 |
| 230 | |
| 231 | untruncated = renderer._tokenize_prompt( |
| 232 | {"prompt": text}, |
| 233 | TokenizeParams(max_total_tokens=1024, return_token_offsets=True), |
| 234 | ) |
| 235 | full_offsets = untruncated["prompt_token_offsets"] |
| 236 | assert len(full_offsets) > keep |
| 237 | |
| 238 | # An explicit truncation_side disables tokenizer-level truncation (see |
| 239 | # get_encode_kwargs), so the tokenizer returns the whole sequence and |
| 240 | # truncation happens in apply_post_tokenization. |
| 241 | params = TokenizeParams( |
| 242 | max_total_tokens=1024, |
| 243 | return_token_offsets=True, |
| 244 | truncate_prompt_tokens=keep, |
| 245 | truncation_side=side, |
| 246 | ) |
| 247 | result = params.apply_post_tokenization( |
| 248 | fast_tokenizer, renderer._tokenize_prompt({"prompt": text}, params) |
| 249 | ) |
| 250 | |
| 251 | offsets = result["prompt_token_offsets"] |
| 252 | assert len(result["prompt_token_ids"]) == keep |
| 253 | assert len(offsets) == keep |
| 254 | |
| 255 | # Equal length is not enough: the surviving offsets must be the ones |
| 256 | # belonging to the surviving tokens. |
| 257 | expected = full_offsets[-keep:] if side == "left" else full_offsets[:keep] |
| 258 | assert offsets == expected |
| 259 | |
| 260 | @pytest.mark.parametrize( |
| 261 | ("side", "expected_offsets"), |
| 262 | [ |
| 263 | ("left", [(36, 37), (37, 38), (38, 39), (39, 40)]), |
| 264 | ("right", [(0, 1), (1, 2), (2, 3), (3, 4)]), |
| 265 | ], |
| 266 | ) |
| 267 | def test_text_pretrim_preserves_source_offsets(self, side, expected_offsets): |
| 268 | renderer = _make_base_renderer_with(_OffsetTokenizer()) |
| 269 | text = "0123456789" * 4 |
| 270 | params = TokenizeParams( |
| 271 | max_total_tokens=16, |
| 272 | return_token_offsets=True, |
| 273 | truncate_prompt_tokens=4, |
| 274 | truncation_side=side, |
| 275 | add_special_tokens=False, |
| 276 | ) |
| 277 | |
| 278 | result = renderer.tokenize_prompt({"prompt": text}, params) |
| 279 | |
| 280 | assert result["prompt_token_offsets"] == expected_offsets |
| 281 | |
| 282 | def test_async_left_text_pretrim_preserves_source_offsets(self): |
| 283 | renderer = _make_base_renderer_with(_OffsetTokenizer()) |
| 284 | text = "0123456789" * 4 |
| 285 | params = TokenizeParams( |
| 286 | max_total_tokens=16, |
| 287 | return_token_offsets=True, |
| 288 | truncate_prompt_tokens=4, |
| 289 | truncation_side="left", |
| 290 | ) |
| 291 | |
| 292 | result = asyncio.run(renderer.tokenize_prompt_async({"prompt": text}, params)) |
| 293 | |
| 294 | assert result["prompt_token_offsets"] == [ |
| 295 | (36, 37), |
| 296 | (37, 38), |
| 297 | (38, 39), |
| 298 | (39, 40), |
| 299 | ] |