| import pickle |
| import re |
| from typing import List |
|
|
|
|
| class TAIMGANTokenizer: |
| def __init__(self, captions_path): |
| with open(captions_path, "rb") as ckpt_file: |
| captions = pickle.load(ckpt_file) |
| self.ix_to_word = captions[2] |
| self.word_to_ix = captions[3] |
| self.token_regex = r'\w+' |
| self.pad_token_id = self.word_to_ix["<end>"] |
| self.pad_repr = "[PAD]" |
|
|
| def encode(self, text: str) -> List[int]: |
| return [self.word_to_ix.get(word, self.pad_token_id) |
| for word in re.findall(self.token_regex, text.lower())] |
|
|
| def decode(self, tokens: List[int]) -> str: |
| return ' '.join([self.ix_to_word[token] |
| if token != self.pad_token_id else self.pad_repr |
| for token in tokens]) |
|
|
