File size: 1,401 Bytes
599bb35 | 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 | """Thin wrapper around a SentencePiece model with fixed special-token ids."""
from pathlib import Path
class SPTokenizer:
PAD_ID, UNK_ID, BOS_ID, EOS_ID = 0, 1, 2, 3
def __init__(self, model_path):
import sentencepiece as spm
model_path = Path(model_path)
if not model_path.exists():
raise FileNotFoundError(f"Tokenizer model not found: {model_path}")
self.sp = spm.SentencePieceProcessor()
self.sp.load(str(model_path))
got = (self.sp.pad_id(), self.sp.unk_id(), self.sp.bos_id(), self.sp.eos_id())
want = (self.PAD_ID, self.UNK_ID, self.BOS_ID, self.EOS_ID)
if got != want:
raise ValueError(f"Tokenizer special ids {got} != expected {want}. "
"Re-train the tokenizer with tokenizer/train_tokenizer.py.")
@property
def vocab_size(self) -> int:
return self.sp.get_piece_size()
def encode(self, text, add_bos=False, add_eos=False):
ids = self.sp.encode(text, out_type=int)
if add_bos:
ids = [self.BOS_ID] + ids
if add_eos:
ids = ids + [self.EOS_ID]
return ids
def pieces(self, text):
return self.sp.encode(text, out_type=str)
def decode(self, ids):
ids = [int(i) for i in ids if int(i) not in (self.PAD_ID, self.BOS_ID, self.EOS_ID)]
return self.sp.decode(ids)
|