| import os | |
| from typing import List | |
| from sentencepiece import SentencePieceProcessor | |
| class Tokenizer(): | |
| """Minimal tokenizer""" | |
| def __init__(self, model_path: str) -> None: | |
| """Loads sentencepiece model""" | |
| spm_model_path = os.path.join(model_path, "sentencepiece.bpe.model") | |
| self.sp = SentencePieceProcessor() | |
| self.sp.load(spm_model_path) | |
| def encode(self, text: str, src_lang_code: str) -> List: | |
| """Transform text to tokens""" | |
| tokens = [src_lang_code] | |
| tokens.extend(self.sp.encode_as_pieces(text)) | |
| tokens.append('</s>') | |
| return tokens | |
| def decode(self, tokens: List[str]): | |
| """Transform tokens to full sentence""" | |
| return self.sp.decode_pieces(tokens) |