import os
from sentencepiece import SentencePieceProcessor

class Tokenizer:
    def __init__(self, model_path: str):
        assert os.path.isfile(model_path), f"Model file not found: {model_path}"
        self.sp_model = SentencePieceProcessor(model_file=model_path)
        self.n_words = self.sp_model.vocab_size()
        self.bos_id = self.sp_model.bos_id()
        self.eos_id = self.sp_model.eos_id()
        self.pad_id = self.sp_model.pad_id()

    def encode(self, s: str, bos: bool = False, eos: bool = False):
        t = self.sp_model.encode(s)
        if bos:
            t = [self.bos_id] + t
        if eos:
            t = t + [self.eos_id]
        return t

    def decode(self, t):
        return self.sp_model.decode(t)