| |
| from pathlib import Path |
| from typing import Iterable, List, Union |
|
|
| import sentencepiece as spm |
|
|
|
|
| class SentencepiecesTokenizer: |
| def __init__(self, bpemodel: Union[Path, str], **kwargs): |
| self.bpemodel = str(bpemodel) |
| self.sp = None |
| self._build() |
|
|
| def __repr__(self): |
| return f'{self.__class__.__name__}(model="{self.bpemodel}")' |
|
|
| def _build(self): |
| if self.sp is None: |
| self.sp = spm.SentencePieceProcessor() |
| self.sp.load(self.bpemodel) |
|
|
| def text2tokens(self, line: str) -> List[str]: |
| self._build() |
| return self.sp.EncodeAsPieces(line) |
|
|
| def tokens2text(self, tokens: Iterable[str]) -> str: |
| self._build() |
| return self.sp.DecodePieces(list(tokens)) |
|
|
| def encode(self, line: str) -> List[int]: |
| self._build() |
| return self.sp.EncodeAsIds(line) |
|
|
| def decode(self, line: List[int]): |
| self._build() |
| return self.sp.DecodeIds(line) |
|
|
| def get_vocab_size(self): |
| self._build() |
| return self.sp.GetPieceSize() |
|
|
| def ids2tokens(self, *args, **kwargs): |
| return self.decode(*args, **kwargs) |
|
|
| def tokens2ids(self, *args, **kwargs): |
| return self.encode(*args, **kwargs) |
|
|
