| """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) |
|
|
