File size: 3,947 Bytes
5f73143 | 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 | """Uniform interface over the tokenizer flavors.
`unigram` loads a SentencePiece model; `bpe`/`pa` load a HuggingFace byte-level
BPE. Ids 0..3 are <unk>/<bos>/<eos>/<pad> in every flavor, so packing, training
and eval code never branches on flavor.
"""
import json
from functools import cached_property
from pathlib import Path
UNK_ID, BOS_ID, EOS_ID, PAD_ID = 0, 1, 2, 3
class Tok:
def __init__(self, path: str | Path):
self.dir = Path(path)
self.meta = json.loads((self.dir / "meta.json").read_text())
self.flavor = self.meta["flavor"]
self.condition = self.meta["condition"]
self.name = f"{self.flavor}_{self.condition}"
if self.flavor == "unigram":
import sentencepiece as spm
self._sp = spm.SentencePieceProcessor(model_file=str(self.dir / "sp.model"))
self.vocab_size = self._sp.get_piece_size()
else:
from tokenizers import Tokenizer
self._hf = Tokenizer.from_file(str(self.dir / "tokenizer.json"))
self.vocab_size = self._hf.get_vocab_size()
def encode(self, text: str, bos: bool = False, eos: bool = False) -> list[int]:
if self.flavor == "unigram":
ids = self._sp.encode(text)
else:
ids = self._hf.encode(text).ids
if bos:
ids = [BOS_ID] + ids
if eos:
ids = ids + [EOS_ID]
return ids
def encode_batch(self, texts: list[str], bos: bool = False, eos: bool = False):
if self.flavor == "unigram":
batch = self._sp.encode(texts)
else:
batch = [e.ids for e in self._hf.encode_batch(texts)]
if bos or eos:
batch = [([BOS_ID] if bos else []) + ids + ([EOS_ID] if eos else [])
for ids in batch]
return batch
def decode(self, ids: list[int]) -> str:
if self.flavor == "unigram":
return self._sp.decode([i for i in ids if i > PAD_ID])
return self._hf.decode([i for i in ids if i > PAD_ID], skip_special_tokens=True)
# --- introspection for the analysis gate ---
def piece(self, idx: int) -> str:
"""Raw vocabulary piece as stored (SP: '▁'-form / '<0xNN>'; BL: bytelevel-mapped)."""
if self.flavor == "unigram":
return self._sp.id_to_piece(idx)
return self._id_to_piece_bl[idx]
@cached_property
def _id_to_piece_bl(self) -> list[str]:
vocab = self._hf.get_vocab()
pieces = [""] * self.vocab_size
for p, i in vocab.items():
pieces[i] = p
return pieces
@cached_property
def _bl_byte_map(self) -> dict[str, int]:
# inverse of the GPT-2 bytes<->unicode table used by ByteLevel
bs = list(range(ord("!"), ord("~") + 1)) + \
list(range(0xA1, 0xAC + 1)) + list(range(0xAE, 0xFF + 1))
cs = bs[:]
n = 0
for b in range(256):
if b not in bs:
bs.append(b)
cs.append(256 + n)
n += 1
return {chr(c): b for b, c in zip(bs, cs)}
def piece_bytes(self, idx: int) -> bytes:
"""The exact bytes a vocab entry emits (specials -> b'')."""
p = self.piece(idx)
if p in ("<unk>", "<bos>", "<eos>", "<pad>"):
return b""
if self.flavor == "unigram":
if len(p) == 6 and p.startswith("<0x") and p.endswith(">"):
return bytes([int(p[3:5], 16)]) # byte-fallback piece
return p.replace("▁", " ").encode("utf-8")
return bytes(self._bl_byte_map[ch] for ch in p)
def is_byte_piece(self, idx: int) -> bool:
"""True if this vocab entry is a raw-byte atom (SP fallback / BL base byte)."""
p = self.piece(idx)
if self.flavor == "unigram":
return len(p) == 6 and p.startswith("<0x") and p.endswith(">")
return idx >= 4 and len(self.piece_bytes(idx)) == 1
|