tiny_gdn/detokenize.py

"""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