[Bugfix][Frontend] Preserve token offset origins after left text pre-trimming (#54692)
Signed-off-by: waizuichougou <2082431897@qq.com>
waizuichougou committed
Sep 1, 2026 at 16:02 UTC
890563368750fad14864181715d9e1b96478868b
3 files changed
+112
tests/renderers/test_token_offsets.py
+58
@@ -8,6 +8,8 @@ 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
@@ -52,6 +54,21 @@ def _make_base_renderer_with(tokenizer):
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)
@@ -239,3 +256,44 @@ class TestTruncationKeepsOffsetsAligned:
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
+ ]
vllm/renderers/base.py
+29
@@ -468,6 +468,27 @@ class BaseRenderer(ABC, Generic[_T]):
468
)
469
return TokensPrompt(prompt_token_ids=list(token_ids), **prompt)
470
471
+ @staticmethod
472
+ def _apply_prompt_char_offset(
473
+ prompt: TokensPrompt, char_offset: int
474
+ ) -> TokensPrompt:
475
+ """Map token offsets back to the original prompt after a text pre-trim."""
476
+ if char_offset == 0:
477
+ return prompt
478
+
479
+ offsets = prompt.get("prompt_token_offsets")
480
+ if offsets is not None:
481
+ prompt["prompt_token_offsets"] = [
482
+ # Fast tokenizers use (0, 0) for special tokens, which are
483
+ # not character spans in the source prompt.
484
+ (start, end)
485
+ if (start, end) == (0, 0)
486
+ else (start + char_offset, end + char_offset)
487
+ for start, end in offsets
488
+ ]
489
+
490
+ return prompt
491
+
492
def _tokenize_prompt(
493
self,
494
prompt: TextPrompt,
@@ -523,8 +544,12 @@ class BaseRenderer(ABC, Generic[_T]):
544
"Expected prompt['prompt'] to be a string before tokenization; "
545
"use 'prompt_token_ids' for token ID inputs"
546
)
547
+ prompt_char_offset = params._get_text_truncation_offset(
548
+ self.tokenizer, prompt["prompt"]
549
+ )
550
prompt = params.apply_pre_tokenization(self.tokenizer, prompt) # type: ignore[arg-type]
551
prompt = self._tokenize_prompt(prompt, params)
552
+ prompt = self._apply_prompt_char_offset(prompt, prompt_char_offset)
553
554
if params.needs_detokenization and "prompt" not in prompt:
555
if "prompt_token_ids" not in prompt:
@@ -559,8 +584,12 @@ class BaseRenderer(ABC, Generic[_T]):
584
"Expected prompt['prompt'] to be a string before tokenization; "
585
"use 'prompt_token_ids' for token ID inputs"
586
)
587
+ prompt_char_offset = params._get_text_truncation_offset(
588
+ self.tokenizer, prompt["prompt"]
589
+ )
590
prompt = params.apply_pre_tokenization(self.tokenizer, prompt) # type: ignore[arg-type]
591
prompt = await self._tokenize_prompt_async(prompt, params)
592
+ prompt = self._apply_prompt_char_offset(prompt, prompt_char_offset)
593
594
if params.needs_detokenization and "prompt" not in prompt:
595
if "prompt_token_ids" not in prompt:
vllm/renderers/params.py
+25
@@ -381,6 +381,31 @@ class TokenizeParams:
381
382
return text
383
384
+ def _get_text_truncation_offset(
385
+ self, tokenizer: TokenizerLike | None, text: str
386
+ ) -> int:
387
+ """Return the number of source characters removed from the left.
388
+
389
+ ``_text_len_check`` pre-truncates long text before tokenization when
390
+ an explicit truncation side is requested. Fast-tokenizer offsets are
391
+ then relative to that shortened string, so callers need this prefix
392
+ length to map them back to the original prompt.
393
+ """
394
+ max_input_tokens = self.max_input_tokens
395
+ if (
396
+ max_input_tokens is None
397
+ or tokenizer is None
398
+ or self.truncate_prompt_tokens is None
399
+ or self.truncation_side != "left"
400
+ ):
401
+ return 0
402
+
403
+ max_input_chars = max_input_tokens * tokenizer.max_chars_per_token
404
+ if max_input_chars <= 0:
405
+ return 0
406
+
407
+ return max(len(text) - max_input_chars, 0)
408
+
409
def _text_lowercase(self, tokenizer: TokenizerLike | None, text: str) -> str:
410
"""Apply lowercase to prompt text if necessary."""
411
return text.lower() if self.do_lower_case else text