File size: 3,627 Bytes
d7562c8 | 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 | """alpha-er's byte-level BPE tokenizer (vocab 12,288).
GPT-2-style byte-level BPE, with the vocabulary laid out as:
0 - 255 raw bytes, in byte order, under GPT-2's bytes->unicode mapping
256 - 258 <|user|>, <|assistant|>, <|end_of_text|>
259 + learned merges, where merges[i] produces token 259 + i
Being byte-level means every input encodes — there is no unknown token and no
normalisation step to disagree about. Decoding maps each vocabulary character
back to its byte and interprets the result as UTF-8, so a sequence cut mid
multi-byte character degrades to a replacement character rather than throwing.
"""
from __future__ import annotations
import json
from functools import lru_cache
@lru_cache(1)
def _byte_maps():
"""GPT-2's reversible bytes<->unicode table."""
bs = (list(range(ord("!"), ord("~") + 1))
+ list(range(ord("\xa1"), ord("\xac") + 1))
+ list(range(ord("\xae"), ord("\xff") + 1)))
cs = bs[:]
n = 0
for b in range(256):
if b not in bs:
bs.append(b); cs.append(256 + n); n += 1
b2u = {b: chr(c) for b, c in zip(bs, cs)}
return b2u, {v: k for k, v in b2u.items()}
class AlphaErTokenizer:
def __init__(self, artifacts: dict):
self.vocab = artifacts["vocab"]
self.merges = [tuple(m) for m in artifacts["merges"]]
self.specials = artifacts["specialTokens"]
self.n_special = len(self.specials)
self.first_merge_id = 256 + self.n_special
# rank[(a,b)] = the token id the pair becomes; lower id = earlier merge.
self.rank = {pair: self.first_merge_id + i for i, pair in enumerate(self.merges)}
self.special_ids = {s: 256 + i for i, s in enumerate(self.specials)}
@classmethod
def from_file(cls, path: str) -> "AlphaErTokenizer":
return cls(json.load(open(path)))
def _encode_chunk(self, text: str) -> list[int]:
b2u, _ = _byte_maps()
ids = list(text.encode("utf-8")) # ids 0-255 ARE the bytes
while len(ids) > 1:
best, best_at = None, -1
for i in range(len(ids) - 1):
r = self.rank.get((ids[i], ids[i + 1]))
if r is not None and (best is None or r < best):
best, best_at = r, i
if best is None:
break
ids[best_at:best_at + 2] = [best]
return ids
def encode(self, text: str) -> list[int]:
"""Split on special tokens first so they survive as single ids."""
parts, out = [text], None
for s in self.specials:
nxt = []
for p in parts:
if isinstance(p, int):
nxt.append(p); continue
bits = p.split(s)
for i, bit in enumerate(bits):
if i: nxt.append(self.special_ids[s])
if bit: nxt.append(bit)
parts = nxt
out = []
for p in parts:
out.extend([p] if isinstance(p, int) else self._encode_chunk(p))
return out
def decode(self, ids) -> str:
_, u2b = _byte_maps()
buf, text = bytearray(), []
for i in ids:
tok = self.vocab[i]
if i < 256 + self.n_special and i >= 256: # a special token
text.append(buf.decode("utf-8", "replace")); buf = bytearray()
text.append(tok)
continue
for ch in tok:
buf.append(u2b.get(ch, ord("?") if ord(ch) < 256 else 63))
text.append(buf.decode("utf-8", "replace"))
return "".join(text)
|