| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Scan delta token IDs for special tokens and split the stream into |
| 4 | pre-lexed terminals and plain text chunks.""" |
| 5 | |
| 6 | from __future__ import annotations |
| 7 | |
| 8 | from collections.abc import Sequence |
| 9 | from dataclasses import dataclass |
| 10 | |
| 11 | DROP_TERMINAL = "__DROP__" |
| 12 | |
| 13 | |
| 14 | @dataclass(slots=True) |
| 15 | class TextChunk: |
| 16 | text: str |
| 17 | token_texts: tuple[str, ...] = () |
| 18 | token_count: int = 0 |
| 19 | |
| 20 | |
| 21 | @dataclass(slots=True) |
| 22 | class PreLexedTerminal: |
| 23 | terminal: str |
| 24 | token_id: int |
| 25 | text: str |
| 26 | |
| 27 | |
| 28 | LexerInput = TextChunk | PreLexedTerminal |
| 29 | |
| 30 | |
| 31 | class TokenIDScanner: |
| 32 | """Maps special token IDs in the delta to terminals. |
| 33 | |
| 34 | Before text-based lexing happens, the scanner checks each token ID |
| 35 | in the delta against a mapping of ``{token_id: terminal_name}``. |
| 36 | Matched tokens are emitted as :class:`PreLexedTerminal` items; |
| 37 | everything else is grouped into :class:`TextChunk` items for the |
| 38 | incremental lexer to process. |
| 39 | |
| 40 | When a terminal's text is not yet in ``delta_text`` (held back by |
| 41 | the detokenizer), the terminal is deferred until the text arrives |
| 42 | in a subsequent delta. |
| 43 | """ |
| 44 | |
| 45 | def __init__( |
| 46 | self, |
| 47 | token_id_to_terminal: dict[int, str], |
| 48 | tokenizer, |
| 49 | ) -> None: |
| 50 | self.token_id_to_terminal = token_id_to_terminal |
| 51 | self.tokenizer = tokenizer |
| 52 | self._token_text_cache: dict[int, str] = {} |
| 53 | self._deferred_terminals: list[PreLexedTerminal] = [] |
| 54 | self._deferred_prefix_token_counts: list[int] = [] |
| 55 | self._deferred_trailing_token_count = 0 |
| 56 | self._deferred_post_text: str = "" |
| 57 | |
| 58 | def reset(self) -> None: |
| 59 | """Clear mutable state for reuse. Preserves the token text cache.""" |
| 60 | self._deferred_terminals.clear() |
| 61 | self._deferred_prefix_token_counts.clear() |
| 62 | self._deferred_trailing_token_count = 0 |
| 63 | self._deferred_post_text = "" |
| 64 | |
| 65 | def _decode_token(self, token_id: int) -> str: |
| 66 | if token_id not in self._token_text_cache: |
| 67 | token_text = self.tokenizer.decode([token_id]) |
| 68 | self._token_text_cache[token_id] = ( |
| 69 | token_text if isinstance(token_text, str) else "" |
| 70 | ) |
| 71 | return self._token_text_cache[token_id] |
| 72 | |
| 73 | def _decode_tokens(self, token_ids: Sequence[int]) -> tuple[str, ...]: |
| 74 | if self.tokenizer is None: |
| 75 | return () |
| 76 | return tuple(self._decode_token(tid) for tid in token_ids) |
| 77 | |
| 78 | _EMPTY: tuple[LexerInput, ...] = () |
| 79 | |
| 80 | def scan( |
| 81 | self, |
| 82 | delta_text: str, |
| 83 | delta_token_ids: Sequence[int], |
| 84 | ) -> Sequence[LexerInput]: |
| 85 | prefix_items: list[LexerInput] = [] |
| 86 | effective_text = delta_text |
| 87 | deferred_trailing_count = 0 |
| 88 | |
| 89 | if self._deferred_terminals: |
| 90 | prefix_items, effective_text, deferred_trailing_count = ( |
| 91 | self._resolve_deferred(delta_text, len(delta_token_ids)) |
| 92 | ) |
| 93 | |
| 94 | if not self.token_id_to_terminal: |
| 95 | if effective_text: |
| 96 | token_texts = self._decode_tokens(delta_token_ids) |
| 97 | prefix_items.append( |
| 98 | TextChunk( |
| 99 | effective_text, |
| 100 | token_texts, |
| 101 | deferred_trailing_count + len(token_texts), |
| 102 | ) |
| 103 | ) |
| 104 | return prefix_items |
| 105 | |
| 106 | has_special = False |
| 107 | token_id_to_terminal = self.token_id_to_terminal |
| 108 | for tid in delta_token_ids: |
| 109 | if tid in token_id_to_terminal: |
| 110 | has_special = True |
| 111 | break |
| 112 | |
| 113 | if not has_special: |
| 114 | if effective_text: |
| 115 | if not prefix_items: |
| 116 | token_texts = self._decode_tokens(delta_token_ids) |
| 117 | return [TextChunk(effective_text, token_texts, len(token_texts))] |
| 118 | token_texts = self._decode_tokens(delta_token_ids) |
| 119 | prefix_items.append( |
| 120 | TextChunk( |
| 121 | effective_text, |
| 122 | token_texts, |
| 123 | deferred_trailing_count + len(token_texts), |
| 124 | ) |
| 125 | ) |
| 126 | return prefix_items or self._EMPTY |
| 127 | |
| 128 | decoded_token_texts = [self._decode_token(tid) for tid in delta_token_ids] |
| 129 | |
| 130 | results: list[LexerInput] = [] |
| 131 | text_accum: list[str] = [] |
| 132 | token_text_accum: list[str] = [] |
| 133 | |
| 134 | for idx, tid in enumerate(delta_token_ids): |
| 135 | terminal = self.token_id_to_terminal.get(tid) |
| 136 | if terminal is not None: |
| 137 | if text_accum: |
| 138 | joined = "".join(text_accum) |
| 139 | if joined: |
| 140 | results.append( |
| 141 | TextChunk( |
| 142 | joined, |
| 143 | tuple(token_text_accum), |
| 144 | len(token_text_accum), |
| 145 | ) |
| 146 | ) |
| 147 | text_accum.clear() |
| 148 | token_text_accum.clear() |
| 149 | results.append( |
| 150 | PreLexedTerminal(terminal, tid, decoded_token_texts[idx]) |
| 151 | ) |
| 152 | else: |
| 153 | text_accum.append(decoded_token_texts[idx]) |
| 154 | token_text_accum.append(decoded_token_texts[idx]) |
| 155 | |
| 156 | if text_accum: |
| 157 | joined = "".join(text_accum) |
| 158 | if joined: |
| 159 | results.append( |
| 160 | TextChunk(joined, tuple(token_text_accum), len(token_text_accum)) |
| 161 | ) |
| 162 | |
| 163 | if deferred_trailing_count: |
| 164 | results.insert(0, TextChunk("", token_count=deferred_trailing_count)) |
| 165 | |
| 166 | if effective_text: |
| 167 | results = self._recover_holdback_text(effective_text, results) |
| 168 | else: |
| 169 | # No detokenizer text to validate against — individually-decoded |
| 170 | # TextChunks are unreliable (context-dependent decoding). |
| 171 | # Defer PreLexedTerminals so the state machine doesn't |
| 172 | # transition before the preceding text has arrived. The |
| 173 | # deferred terminals will be resolved against the actual |
| 174 | # delta_text in a subsequent scan() or flushed by finish(). |
| 175 | prefix_token_count = 0 |
| 176 | for r in results: |
| 177 | if isinstance(r, TextChunk): |
| 178 | prefix_token_count += r.token_count |
| 179 | if isinstance(r, PreLexedTerminal): |
| 180 | self._deferred_terminals.append(r) |
| 181 | self._deferred_prefix_token_counts.append(prefix_token_count) |
| 182 | prefix_token_count = 0 |
| 183 | self._deferred_trailing_token_count += prefix_token_count |
| 184 | results = [] |
| 185 | |
| 186 | return prefix_items + results |
| 187 | |
| 188 | def flush_pending(self) -> list[LexerInput]: |
| 189 | if not self._deferred_terminals and not self._deferred_post_text: |
| 190 | return [] |
| 191 | results: list[LexerInput] = [] |
| 192 | if self._deferred_post_text: |
| 193 | prefix_count = ( |
| 194 | self._deferred_prefix_token_counts[0] |
| 195 | if self._deferred_prefix_token_counts |
| 196 | else 0 |
| 197 | ) |
| 198 | results.append( |
| 199 | TextChunk(self._deferred_post_text, token_count=prefix_count) |
| 200 | ) |
| 201 | self._deferred_post_text = "" |
| 202 | results.extend(self._deferred_terminals) |
| 203 | if self._deferred_trailing_token_count: |
| 204 | results.append( |
| 205 | TextChunk("", token_count=self._deferred_trailing_token_count) |
| 206 | ) |
| 207 | self._deferred_terminals.clear() |
| 208 | self._deferred_prefix_token_counts.clear() |
| 209 | self._deferred_trailing_token_count = 0 |
| 210 | return results |
| 211 | |
| 212 | def _resolve_deferred( |
| 213 | self, |
| 214 | delta_text: str, |
| 215 | current_token_count: int = 0, |
| 216 | ) -> tuple[list[LexerInput], str, int]: |
| 217 | """Resolve deferred terminals against new delta_text. |
| 218 | |
| 219 | When a previous ``scan()`` deferred a terminal (its text hadn't |
| 220 | arrived yet), the next delta's text should contain that terminal's |
| 221 | text. Split delta_text at the terminal boundary: text before |
| 222 | belongs to the previous parser state, the terminal triggers the |
| 223 | state transition, and text after belongs to the new state. |
| 224 | |
| 225 | Returns ``(prefix_items, remaining_text)`` where prefix_items |
| 226 | are the resolved deferred terminals (with any preceding text) |
| 227 | and remaining_text is the unconsumed portion of delta_text that |
| 228 | should be scanned with the current delta's token IDs. |
| 229 | """ |
| 230 | deferred = self._deferred_terminals |
| 231 | prefix_token_counts = self._deferred_prefix_token_counts |
| 232 | trailing_token_count = self._deferred_trailing_token_count |
| 233 | self._deferred_terminals = [] |
| 234 | self._deferred_prefix_token_counts = [] |
| 235 | self._deferred_trailing_token_count = 0 |
| 236 | |
| 237 | results: list[LexerInput] = [] |
| 238 | remaining = delta_text |
| 239 | |
| 240 | if self._deferred_post_text: |
| 241 | remaining = self._deferred_post_text + remaining |
| 242 | self._deferred_post_text = "" |
| 243 | |
| 244 | # Duplicate-text deferred terminals resolve left-to-right via |
| 245 | # find(); correct when each terminal text appears once in sequence. |
| 246 | for idx, terminal in enumerate(deferred): |
| 247 | prefix_token_count = prefix_token_counts[idx] |
| 248 | pos = remaining.find(terminal.text) |
| 249 | if pos > 0: |
| 250 | results.append( |
| 251 | TextChunk(remaining[:pos], token_count=prefix_token_count) |
| 252 | ) |
| 253 | results.append(terminal) |
| 254 | remaining = remaining[pos + len(terminal.text) :] |
| 255 | elif pos == 0: |
| 256 | results.append(terminal) |
| 257 | remaining = remaining[len(terminal.text) :] |
| 258 | else: |
| 259 | # Accumulate text until terminal text arrives — |
| 260 | # only the terminal provides a reliable split point. |
| 261 | if remaining: |
| 262 | self._deferred_post_text += remaining |
| 263 | prefix_token_count += current_token_count |
| 264 | remaining = "" |
| 265 | self._deferred_terminals.append(terminal) |
| 266 | self._deferred_prefix_token_counts.append(prefix_token_count) |
| 267 | |
| 268 | if self._deferred_terminals: |
| 269 | self._deferred_trailing_token_count = trailing_token_count |
| 270 | trailing_token_count = 0 |
| 271 | |
| 272 | return results, remaining, trailing_token_count |
| 273 | |
| 274 | def _recover_holdback_text( |
| 275 | self, |
| 276 | delta_text: str, |
| 277 | results: list[LexerInput], |
| 278 | ) -> list[LexerInput]: |
| 279 | """Recover detokenizer hold-back text not in delta_token_ids. |
| 280 | |
| 281 | The detokenizer may flush previously held-back text in |
| 282 | ``delta_text`` that has no corresponding token ID in |
| 283 | ``delta_token_ids``. This hold-back text always appears as a |
| 284 | prefix of ``delta_text``. |
| 285 | """ |
| 286 | if not results: |
| 287 | return [TextChunk(delta_text)] |
| 288 | |
| 289 | reconstructed = self._join_decoded_text(results) |
| 290 | |
| 291 | if not reconstructed: |
| 292 | return [TextChunk(delta_text)] + results |
| 293 | |
| 294 | pos = delta_text.find(reconstructed) |
| 295 | if pos > 0: |
| 296 | return [TextChunk(delta_text[:pos])] + results |
| 297 | if pos == 0: |
| 298 | return results |
| 299 | |
| 300 | # Fallback: SentencePiece context-dependent decoding mismatch. |
| 301 | # Rebuild from delta_text using PreLexedTerminals as split anchors. |
| 302 | return self._rebuild_from_anchors(delta_text, results) |
| 303 | |
| 304 | def _join_decoded_text(self, results: list[LexerInput]) -> str: |
| 305 | """Join TextChunk and PreLexedTerminal text into one string.""" |
| 306 | parts: list[str] = [] |
| 307 | for item in results: |
| 308 | if isinstance(item, (TextChunk, PreLexedTerminal)): |
| 309 | parts.append(item.text) |
| 310 | return "".join(parts) |
| 311 | |
| 312 | def _rebuild_from_anchors( |
| 313 | self, |
| 314 | delta_text: str, |
| 315 | results: list[LexerInput], |
| 316 | ) -> list[LexerInput]: |
| 317 | """Rebuild results from delta_text using terminals as anchors. |
| 318 | |
| 319 | When context-dependent decoding creates a mismatch between |
| 320 | individually-decoded tokens and delta_text, use |
| 321 | PreLexedTerminals as split points and reallocate text from |
| 322 | delta_text. If a terminal's text is not found in delta_text, |
| 323 | it is deferred to the next scan() call. |
| 324 | |
| 325 | Anchors are resolved right-to-left with ``rfind`` so that each |
| 326 | anchor binds to the *rightmost* available occurrence of its |
| 327 | text. This prevents earlier literal lookalikes (e.g. a user |
| 328 | mentioning ``<tool_call>`` in prose) from stealing the position |
| 329 | of a real special-token anchor that appears later. |
| 330 | |
| 331 | If the same anchor text appears multiple times as real special |
| 332 | tokens (not prose), the rightmost-first binding could misalign. |
| 333 | In practice this doesn't happen: each special token ID maps to |
| 334 | a distinct PreLexedTerminal, and duplicates in prose are resolved |
| 335 | by the token-ID filtering layer above. |
| 336 | """ |
| 337 | anchors = [item for item in results if isinstance(item, PreLexedTerminal)] |
| 338 | if not anchors: |
| 339 | return [TextChunk(delta_text)] |
| 340 | |
| 341 | token_groups: list[list[str]] = [[] for _ in range(len(anchors) + 1)] |
| 342 | count_groups = [0] * (len(anchors) + 1) |
| 343 | group_idx = 0 |
| 344 | for item in results: |
| 345 | if isinstance(item, PreLexedTerminal): |
| 346 | group_idx += 1 |
| 347 | else: |
| 348 | token_groups[group_idx].extend(item.token_texts) |
| 349 | count_groups[group_idx] += item.token_count |
| 350 | |
| 351 | # Resolve positions right-to-left: each anchor gets the |
| 352 | # rightmost occurrence that is still before the next anchor. |
| 353 | positions: list[int] = [-1] * len(anchors) |
| 354 | search_end = len(delta_text) |
| 355 | for i in range(len(anchors) - 1, -1, -1): |
| 356 | pos = delta_text.rfind(anchors[i].text, 0, search_end) |
| 357 | if pos >= 0: |
| 358 | positions[i] = pos |
| 359 | search_end = pos |
| 360 | |
| 361 | # Build results left-to-right using the resolved positions. |
| 362 | new_results: list[LexerInput] = [] |
| 363 | consumed = 0 |
| 364 | for i, anchor in enumerate(anchors): |
| 365 | pos = positions[i] |
| 366 | if pos >= consumed: |
| 367 | if pos > consumed: |
| 368 | new_results.append( |
| 369 | TextChunk( |
| 370 | delta_text[consumed:pos], |
| 371 | tuple(token_groups[i]), |
| 372 | count_groups[i], |
| 373 | ) |
| 374 | ) |
| 375 | new_results.append(anchor) |
| 376 | consumed = pos + len(anchor.text) |
| 377 | else: |
| 378 | has_later_valid = any(p >= 0 for p in positions[i + 1 :]) |
| 379 | # DROP anchors (EOS, etc.) may have text that never |
| 380 | # arrives in delta_text (stripped by detokenizer). |
| 381 | # Don't defer remaining content waiting for text |
| 382 | # that will never come. |
| 383 | if ( |
| 384 | not has_later_valid |
| 385 | and consumed < len(delta_text) |
| 386 | and anchor.terminal != DROP_TERMINAL |
| 387 | ): |
| 388 | self._deferred_post_text += delta_text[consumed:] |
| 389 | consumed = len(delta_text) |
| 390 | self._deferred_terminals.append(anchor) |
| 391 | self._deferred_prefix_token_counts.append(count_groups[i]) |
| 392 | if consumed < len(delta_text): |
| 393 | new_results.append( |
| 394 | TextChunk( |
| 395 | delta_text[consumed:], |
| 396 | tuple(token_groups[-1]), |
| 397 | count_groups[-1], |
| 398 | ) |
| 399 | ) |
| 400 | return new_results |