model_utils/tokenizer.py

from .vocabulary import vocabulary

class Tokenizer:
    def __init__(self, vocab=vocabulary):
        self.str_to_int = vocab
        self.int_to_str = {id:token for token, id in vocab.items()}

    def encode(self, text):
        text_tokens = [token if token in self.str_to_int else "<|unk|>" for token in text]
        ids = [self.str_to_int[token] for token in text_tokens]
        return ids

    def decode(self, ids):
        text = "".join([self.int_to_str[id] for id in ids])
        return text