Spaces:
Build error
Build error
| 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) | |
| def load(cls, path): | |
| tokenizer = cls() | |
| with open(path, 'rb') as f: | |
| tokenizer.__dict__ = pickle.load(f) | |
| return tokenizer |