| 266 | EOS_token = 2 # End-of-sentence token |
| 267 | |
| 268 | class Voc: |
| 269 | def __init__(self, name): |
| 270 | self.name = name |
| 271 | self.trimmed = False |
| 272 | self.word2index = {} |
| 273 | self.word2count = {} |
| 274 | self.index2word = {PAD_token: "PAD", SOS_token: "SOS", EOS_token: "EOS"} |
| 275 | self.num_words = 3 # Count SOS, EOS, PAD |
| 276 | |
| 277 | def addSentence(self, sentence): |
| 278 | for word in sentence.split(' '): |
| 279 | self.addWord(word) |
| 280 | |
| 281 | def addWord(self, word): |
| 282 | if word not in self.word2index: |
| 283 | self.word2index[word] = self.num_words |
| 284 | self.word2count[word] = 1 |
| 285 | self.index2word[self.num_words] = word |
| 286 | self.num_words += 1 |
| 287 | else: |
| 288 | self.word2count[word] += 1 |
| 289 | |
| 290 | # Remove words below a certain count threshold |
| 291 | def trim(self, min_count): |
| 292 | if self.trimmed: |
| 293 | return |
| 294 | self.trimmed = True |
| 295 | |
| 296 | keep_words = [] |
| 297 | |
| 298 | for k, v in self.word2count.items(): |
| 299 | if v >= min_count: |
| 300 | keep_words.append(k) |
| 301 | |
| 302 | print('keep_words {} / {} = {:.4f}'.format( |
| 303 | len(keep_words), len(self.word2index), len(keep_words) / len(self.word2index) |
| 304 | )) |
| 305 | |
| 306 | # Reinitialize dictionaries |
| 307 | self.word2index = {} |
| 308 | self.word2count = {} |
| 309 | self.index2word = {PAD_token: "PAD", SOS_token: "SOS", EOS_token: "EOS"} |
| 310 | self.num_words = 3 # Count default tokens |
| 311 | |
| 312 | for word in keep_words: |
| 313 | self.addWord(word) |
| 314 | |
| 315 | |
| 316 | ###################################################################### |