| 95 | self.all_special_ids = [self.encoder[t] for t in special_tokens] |
| 96 | |
| 97 | def bpe(self, token): |
| 98 | if token in self.cache: |
| 99 | return self.cache[token] |
| 100 | word = tuple(token[:-1]) + (token[-1] + '</w>',) |
| 101 | pairs = get_pairs(word) |
| 102 | |
| 103 | if not pairs: |
| 104 | return token + '</w>' |
| 105 | |
| 106 | while True: |
| 107 | bigram = min(pairs, key=lambda pair: self.bpe_ranks.get( |
| 108 | pair, float('inf'))) |
| 109 | if bigram not in self.bpe_ranks: |
| 110 | break |
| 111 | first, second = bigram |
| 112 | new_word = [] |
| 113 | i = 0 |
| 114 | while i < len(word): |
| 115 | try: |
| 116 | j = word.index(first, i) |
| 117 | new_word.extend(word[i:j]) |
| 118 | i = j |
| 119 | except: |
| 120 | new_word.extend(word[i:]) |
| 121 | break |
| 122 | |
| 123 | if word[i] == first and i < len(word) - 1 and word[i + 1] == second: |
| 124 | new_word.append(first + second) |
| 125 | i += 2 |
| 126 | else: |
| 127 | new_word.append(word[i]) |
| 128 | i += 1 |
| 129 | new_word = tuple(new_word) |
| 130 | word = new_word |
| 131 | if len(word) == 1: |
| 132 | break |
| 133 | else: |
| 134 | pairs = get_pairs(word) |
| 135 | word = ' '.join(word) |
| 136 | self.cache[token] = word |
| 137 | return word |
| 138 | |
| 139 | def encode(self, text): |
| 140 | bpe_tokens = [] |