File size: 757 Bytes
ba7290b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | 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) |