| """Byte-accurate incremental decode for GPT-2-style byte-level BPE. |
| |
| `tokenizer.decode(ids)` UTF-8-replaces incomplete sequences. Streaming that |
| as a string prefix then breaks on the next byte: decode([20492]) is a |
| replacement character and decode([20492, 294]) is `` ≈``, so the |
| accumulated text is no longer a prefix and the UI re-emits the whole |
| think block. |
| |
| This decoder concatenates raw BPE bytes and only emits complete UTF-8 |
| characters. The completing byte of ``≈`` yields `` ≈``; nothing is rewound. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Iterable, Mapping, Sequence |
| from typing import Any |
|
|
|
|
| def gpt2_bytes_to_unicode() -> dict[int, str]: |
| raw = ( |
| list(range(ord("!"), ord("~") + 1)) |
| + list(range(ord("¡"), ord("¬") + 1)) |
| + list(range(ord("®"), ord("ÿ") + 1)) |
| ) |
| mapped = raw[:] |
| extra = 0 |
| for byte in range(256): |
| if byte not in raw: |
| raw.append(byte) |
| mapped.append(256 + extra) |
| extra += 1 |
| return dict(zip(raw, [chr(code) for code in mapped], strict=True)) |
|
|
|
|
| def unicode_to_gpt2_bytes() -> dict[str, int]: |
| return {char: byte for byte, char in gpt2_bytes_to_unicode().items()} |
|
|
|
|
| def utf8_char_length(lead: int) -> int | None: |
| if lead < 0x80: |
| return 1 |
| if 0xC2 <= lead <= 0xDF: |
| return 2 |
| if 0xE0 <= lead <= 0xEF: |
| return 3 |
| if 0xF0 <= lead <= 0xF4: |
| return 4 |
| return None |
|
|
|
|
| def split_complete_utf8(buffer: bytes) -> tuple[str, bytes]: |
| """Return (decoded prefix, leftover incomplete or invalid-lead bytes).""" |
| index = 0 |
| end = len(buffer) |
| while index < end: |
| size = utf8_char_length(buffer[index]) |
| if size is None: |
| break |
| if index + size > end: |
| break |
| chunk = buffer[index : index + size] |
| try: |
| chunk.decode("utf-8") |
| except UnicodeDecodeError: |
| break |
| index += size |
| if index == 0: |
| return "", buffer |
| return buffer[:index].decode("utf-8"), buffer[index:] |
|
|
|
|
| def piece_to_bytes(piece: str, unicode_to_byte: Mapping[str, int]) -> bytes: |
| raw = bytearray() |
| for char in piece: |
| mapped = unicode_to_byte.get(char) |
| if mapped is None: |
| raw.extend(char.encode("utf-8")) |
| else: |
| raw.append(mapped) |
| return bytes(raw) |
|
|
|
|
| class IncrementalUtf8Decoder: |
| """Push token ids; read only complete Unicode as it becomes available.""" |
|
|
| def __init__( |
| self, |
| tokenizer: Any, |
| *, |
| skip_special_ids: Iterable[int] | None = None, |
| ) -> None: |
| vocab = tokenizer.get_vocab() |
| self._id_to_piece = {tid: piece for piece, tid in vocab.items()} |
| self._unicode_to_byte = unicode_to_gpt2_bytes() |
| self._skip_special_ids = set(skip_special_ids or ()) |
| self._pending = b"" |
| self._text = "" |
|
|
| @property |
| def text(self) -> str: |
| return self._text |
|
|
| def reset(self) -> None: |
| self._pending = b"" |
| self._text = "" |
|
|
| def token_bytes(self, token_id: int) -> bytes: |
| piece = self._id_to_piece.get(token_id, "") |
| return piece_to_bytes(piece, self._unicode_to_byte) |
|
|
| def push(self, token_id: int) -> str: |
| if token_id in self._skip_special_ids: |
| return "" |
| self._pending += self.token_bytes(token_id) |
| complete, self._pending = split_complete_utf8(self._pending) |
| self._text += complete |
| return complete |
|
|
| def push_many(self, token_ids: Sequence[int]) -> str: |
| delta = [] |
| for token_id in token_ids: |
| piece = self.push(token_id) |
| if piece: |
| delta.append(piece) |
| return "".join(delta) |
|
|
| def finalize(self) -> str: |
| if self._pending: |
| self._text += self._pending.decode("utf-8", errors="replace") |
| self._pending = b"" |
| return self._text |
|
|
