File size: 3,933 Bytes
90e4c64 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """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
|