wehub-resource-sync
297 lines
11 KiB
Python
297 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Scan delta token IDs for special tokens and split the stream into
|
|
pre-lexed terminals and plain text chunks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
|
|
DROP_TERMINAL = "__DROP__"
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class TextChunk:
|
|
text: str
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class PreLexedTerminal:
|
|
terminal: str
|
|
token_id: int
|
|
text: str
|
|
|
|
|
|
LexerInput = TextChunk | PreLexedTerminal
|
|
|
|
|
|
class TokenIDScanner:
|
|
"""Maps special token IDs in the delta to terminals.
|
|
|
|
Before text-based lexing happens, the scanner checks each token ID
|
|
in the delta against a mapping of ``{token_id: terminal_name}``.
|
|
Matched tokens are emitted as :class:`PreLexedTerminal` items;
|
|
everything else is grouped into :class:`TextChunk` items for the
|
|
incremental lexer to process.
|
|
|
|
When a terminal's text is not yet in ``delta_text`` (held back by
|
|
the detokenizer), the terminal is deferred until the text arrives
|
|
in a subsequent delta.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
token_id_to_terminal: dict[int, str],
|
|
tokenizer,
|
|
) -> None:
|
|
self.token_id_to_terminal = token_id_to_terminal
|
|
self.tokenizer = tokenizer
|
|
self._token_text_cache: dict[int, str] = {}
|
|
self._deferred_terminals: list[PreLexedTerminal] = []
|
|
self._deferred_post_text: str = ""
|
|
|
|
def reset(self) -> None:
|
|
"""Clear mutable state for reuse. Preserves the token text cache."""
|
|
self._deferred_terminals.clear()
|
|
self._deferred_post_text = ""
|
|
|
|
def _decode_token(self, token_id: int) -> str:
|
|
if token_id not in self._token_text_cache:
|
|
self._token_text_cache[token_id] = self.tokenizer.decode([token_id])
|
|
return self._token_text_cache[token_id]
|
|
|
|
_EMPTY: tuple[LexerInput, ...] = ()
|
|
|
|
def scan(
|
|
self,
|
|
delta_text: str,
|
|
delta_token_ids: Sequence[int],
|
|
) -> Sequence[LexerInput]:
|
|
prefix_items: list[LexerInput] = []
|
|
effective_text = delta_text
|
|
|
|
if self._deferred_terminals:
|
|
prefix_items, effective_text = self._resolve_deferred(delta_text)
|
|
|
|
if not self.token_id_to_terminal:
|
|
if effective_text:
|
|
prefix_items.append(TextChunk(effective_text))
|
|
return prefix_items
|
|
|
|
has_special = False
|
|
token_id_to_terminal = self.token_id_to_terminal
|
|
for tid in delta_token_ids:
|
|
if tid in token_id_to_terminal:
|
|
has_special = True
|
|
break
|
|
|
|
if not has_special:
|
|
if effective_text:
|
|
if not prefix_items:
|
|
return [TextChunk(effective_text)]
|
|
prefix_items.append(TextChunk(effective_text))
|
|
return prefix_items or self._EMPTY
|
|
|
|
token_texts = [self._decode_token(tid) for tid in delta_token_ids]
|
|
|
|
results: list[LexerInput] = []
|
|
text_accum: list[str] = []
|
|
|
|
for idx, tid in enumerate(delta_token_ids):
|
|
terminal = self.token_id_to_terminal.get(tid)
|
|
if terminal is not None:
|
|
if text_accum:
|
|
joined = "".join(text_accum)
|
|
if joined:
|
|
results.append(TextChunk(joined))
|
|
text_accum.clear()
|
|
results.append(PreLexedTerminal(terminal, tid, token_texts[idx]))
|
|
else:
|
|
text_accum.append(token_texts[idx])
|
|
|
|
if text_accum:
|
|
joined = "".join(text_accum)
|
|
if joined:
|
|
results.append(TextChunk(joined))
|
|
|
|
if effective_text:
|
|
results = self._recover_holdback_text(effective_text, results)
|
|
else:
|
|
# No detokenizer text to validate against — individually-decoded
|
|
# TextChunks are unreliable (context-dependent decoding).
|
|
# Defer PreLexedTerminals so the state machine doesn't
|
|
# transition before the preceding text has arrived. The
|
|
# deferred terminals will be resolved against the actual
|
|
# delta_text in a subsequent scan() or flushed by finish().
|
|
for r in results:
|
|
if isinstance(r, PreLexedTerminal):
|
|
self._deferred_terminals.append(r)
|
|
results = []
|
|
|
|
return prefix_items + results
|
|
|
|
def flush_pending(self) -> list[LexerInput]:
|
|
if not self._deferred_terminals and not self._deferred_post_text:
|
|
return []
|
|
results: list[LexerInput] = []
|
|
if self._deferred_post_text:
|
|
results.append(TextChunk(self._deferred_post_text))
|
|
self._deferred_post_text = ""
|
|
results.extend(self._deferred_terminals)
|
|
self._deferred_terminals.clear()
|
|
return results
|
|
|
|
def _resolve_deferred(
|
|
self,
|
|
delta_text: str,
|
|
) -> tuple[list[LexerInput], str]:
|
|
"""Resolve deferred terminals against new delta_text.
|
|
|
|
When a previous ``scan()`` deferred a terminal (its text hadn't
|
|
arrived yet), the next delta's text should contain that terminal's
|
|
text. Split delta_text at the terminal boundary: text before
|
|
belongs to the previous parser state, the terminal triggers the
|
|
state transition, and text after belongs to the new state.
|
|
|
|
Returns ``(prefix_items, remaining_text)`` where prefix_items
|
|
are the resolved deferred terminals (with any preceding text)
|
|
and remaining_text is the unconsumed portion of delta_text that
|
|
should be scanned with the current delta's token IDs.
|
|
"""
|
|
deferred = self._deferred_terminals
|
|
self._deferred_terminals = []
|
|
|
|
results: list[LexerInput] = []
|
|
remaining = delta_text
|
|
|
|
if self._deferred_post_text:
|
|
remaining = self._deferred_post_text + remaining
|
|
self._deferred_post_text = ""
|
|
|
|
# Duplicate-text deferred terminals resolve left-to-right via
|
|
# find(); correct when each terminal text appears once in sequence.
|
|
for terminal in deferred:
|
|
pos = remaining.find(terminal.text)
|
|
if pos > 0:
|
|
results.append(TextChunk(remaining[:pos]))
|
|
results.append(terminal)
|
|
remaining = remaining[pos + len(terminal.text) :]
|
|
elif pos == 0:
|
|
results.append(terminal)
|
|
remaining = remaining[len(terminal.text) :]
|
|
else:
|
|
# Accumulate text until terminal text arrives —
|
|
# only the terminal provides a reliable split point.
|
|
if remaining:
|
|
self._deferred_post_text += remaining
|
|
remaining = ""
|
|
self._deferred_terminals.append(terminal)
|
|
|
|
return results, remaining
|
|
|
|
def _recover_holdback_text(
|
|
self,
|
|
delta_text: str,
|
|
results: list[LexerInput],
|
|
) -> list[LexerInput]:
|
|
"""Recover detokenizer hold-back text not in delta_token_ids.
|
|
|
|
The detokenizer may flush previously held-back text in
|
|
``delta_text`` that has no corresponding token ID in
|
|
``delta_token_ids``. This hold-back text always appears as a
|
|
prefix of ``delta_text``.
|
|
"""
|
|
if not results:
|
|
return [TextChunk(delta_text)]
|
|
|
|
reconstructed = self._join_decoded_text(results)
|
|
|
|
if not reconstructed:
|
|
return [TextChunk(delta_text)] + results
|
|
|
|
pos = delta_text.find(reconstructed)
|
|
if pos > 0:
|
|
return [TextChunk(delta_text[:pos])] + results
|
|
if pos == 0:
|
|
return results
|
|
|
|
# Fallback: SentencePiece context-dependent decoding mismatch.
|
|
# Rebuild from delta_text using PreLexedTerminals as split anchors.
|
|
return self._rebuild_from_anchors(delta_text, results)
|
|
|
|
def _join_decoded_text(self, results: list[LexerInput]) -> str:
|
|
"""Join TextChunk and PreLexedTerminal text into one string."""
|
|
parts: list[str] = []
|
|
for item in results:
|
|
if isinstance(item, (TextChunk, PreLexedTerminal)):
|
|
parts.append(item.text)
|
|
return "".join(parts)
|
|
|
|
def _rebuild_from_anchors(
|
|
self,
|
|
delta_text: str,
|
|
results: list[LexerInput],
|
|
) -> list[LexerInput]:
|
|
"""Rebuild results from delta_text using terminals as anchors.
|
|
|
|
When context-dependent decoding creates a mismatch between
|
|
individually-decoded tokens and delta_text, use
|
|
PreLexedTerminals as split points and reallocate text from
|
|
delta_text. If a terminal's text is not found in delta_text,
|
|
it is deferred to the next scan() call.
|
|
|
|
Anchors are resolved right-to-left with ``rfind`` so that each
|
|
anchor binds to the *rightmost* available occurrence of its
|
|
text. This prevents earlier literal lookalikes (e.g. a user
|
|
mentioning ``<tool_call>`` in prose) from stealing the position
|
|
of a real special-token anchor that appears later.
|
|
|
|
If the same anchor text appears multiple times as real special
|
|
tokens (not prose), the rightmost-first binding could misalign.
|
|
In practice this doesn't happen: each special token ID maps to
|
|
a distinct PreLexedTerminal, and duplicates in prose are resolved
|
|
by the token-ID filtering layer above.
|
|
"""
|
|
anchors = [item for item in results if isinstance(item, PreLexedTerminal)]
|
|
if not anchors:
|
|
return [TextChunk(delta_text)]
|
|
|
|
# Resolve positions right-to-left: each anchor gets the
|
|
# rightmost occurrence that is still before the next anchor.
|
|
positions: list[int] = [-1] * len(anchors)
|
|
search_end = len(delta_text)
|
|
for i in range(len(anchors) - 1, -1, -1):
|
|
pos = delta_text.rfind(anchors[i].text, 0, search_end)
|
|
if pos >= 0:
|
|
positions[i] = pos
|
|
search_end = pos
|
|
|
|
# Build results left-to-right using the resolved positions.
|
|
new_results: list[LexerInput] = []
|
|
consumed = 0
|
|
for i, anchor in enumerate(anchors):
|
|
pos = positions[i]
|
|
if pos >= consumed:
|
|
if pos > consumed:
|
|
new_results.append(TextChunk(delta_text[consumed:pos]))
|
|
new_results.append(anchor)
|
|
consumed = pos + len(anchor.text)
|
|
else:
|
|
has_later_valid = any(p >= 0 for p in positions[i + 1 :])
|
|
# DROP anchors (EOS, etc.) may have text that never
|
|
# arrives in delta_text (stripped by detokenizer).
|
|
# Don't defer remaining content waiting for text
|
|
# that will never come.
|
|
if (
|
|
not has_later_valid
|
|
and consumed < len(delta_text)
|
|
and anchor.terminal != DROP_TERMINAL
|
|
):
|
|
self._deferred_post_text += delta_text[consumed:]
|
|
consumed = len(delta_text)
|
|
self._deferred_terminals.append(anchor)
|
|
if consumed < len(delta_text):
|
|
new_results.append(TextChunk(delta_text[consumed:]))
|
|
return new_results
|