import pickle

class CaptionTokenizer:
    def __init__(self):
        self.tok_to_ind = {'<PAD>': 0, '<BOS>': 1, '<EOS>': 2, '<UNK>': 3}
        self.ind_to_tok = {v: k for k, v in self.tok_to_ind.items()}
        
    def encode(self, text):
        return [self.tok_to_ind.get(word, 3) for word in text.split()]
        
    def decode(self, indices):
        return ' '.join([self.ind_to_tok.get(idx, '<UNK>') for idx in indices])
        
    def save(self, path):
        with open(path, 'wb') as f:
            pickle.dump(self.__dict__, f)
            
    @classmethod
    def load(cls, path):
        tokenizer = cls()
        with open(path, 'rb') as f:
            tokenizer.__dict__ = pickle.load(f)
        return tokenizer