Problem 102 · Tokenization & Sampling

Byte-Fallback Detokenizer with Streaming

MediumAI engineering

01Problem

A SentencePiece-style vocab maps ids to pieces. Normal pieces use `▁` for a leading space; byte-fallback pieces look like `<0xE2>` and stand for one raw byte. Write `decode(ids, vocab)` that returns the text (dropping the single leading space the first piece introduces), and a `StreamDecoder` whose `push(id)` returns only the text that is now complete, holding back partial UTF-8 sequences, plus `flush()` for the end of stream.

02Key insight

Tokens are not characters. An emoji or a CJK character can arrive as three or four separate byte tokens, so decoding each token independently and concatenating strings prints replacement characters mid-stream. The fix is to concatenate **bytes** and decode once, or, for streaming, use an incremental UTF-8 decoder that buffers an incomplete trailing sequence. `flush()` must still surface a truncated sequence as U+FFFD instead of silently dropping it, so a cut-off generation is visible.

03Solution

solution.pypython37 lines

import codecs
import re

_BYTE = re.compile(r"^<0x([0-9A-Fa-f]{2})>$")
SPACE = "▁"


def piece_bytes(piece: str) -> bytes:
    m = _BYTE.match(piece)
    if m:
        return bytes([int(m.group(1), 16)])
    return piece.replace(SPACE, " ").encode("utf-8")


def decode(ids: list[int], vocab: dict[int, str]) -> str:
    text = b"".join(piece_bytes(vocab[i]) for i in ids).decode("utf-8", errors="replace")
    return text[1:] if text.startswith(" ") else text


class StreamDecoder:
    def __init__(self, vocab: dict[int, str]) -> None:
        self.vocab = vocab
        self._dec = codecs.getincrementaldecoder("utf-8")(errors="replace")
        self._started = False

    def _emit(self, s: str) -> str:
        if s and not self._started:
            self._started = True
            if s.startswith(" "):
                s = s[1:]
        return s

    def push(self, token_id: int) -> str:
        return self._emit(self._dec.decode(piece_bytes(self.vocab[token_id])))

    def flush(self) -> str:
        return self._emit(self._dec.decode(b"", final=True))

Try it first.

TimeO(total bytes)

SpaceO(output) for decode; O(1) buffered bytes (at most 3) for streaming

tests.py · 8 passing

tests.pypython12 lines

vocab = {0: "▁Hello", 1: "▁wor", 2: "ld", 3: "<0xE2>", 4: "<0x82>", 5: "<0xAC>", 6: "!", 7: "<0xF0>", 8: "<0x41>"}
assert decode([0, 1, 2, 3, 4, 5, 6], vocab) == "Hello world€!"
assert decode([8, 6], vocab) == "A!"
assert decode([0, 7], vocab) == "Hello�"
s = StreamDecoder(vocab)
outs = [s.push(i) for i in [0, 1, 2, 3, 4, 5, 6]]
assert outs == ["Hello", " wor", "ld", "", "", "€", "!"], outs
assert s.flush() == ""
s2 = StreamDecoder(vocab)
assert s2.push(0) == "Hello" and s2.push(7) == ""
assert s2.flush() == "�"
assert "".join(outs) == decode([0, 1, 2, 3, 4, 5, 6], vocab)

04Staff-level follow-up

Stop sequences interact with this: you must match stop strings on decoded text, but a stop string can straddle tokens, so hold back the longest suffix that is a prefix of any stop string before emitting. Also, many tokenizers are not round-trip stable (normalisation, added prefix spaces), so never re-tokenise decoded text to count tokens for billing; keep the ids.