preprocessors/tokenize.py

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)