| 3 | import torch |
| 4 | |
| 5 | class Dictionary(object): |
| 6 | def __init__(self): |
| 7 | self.word2idx = {} |
| 8 | self.idx2word = [] |
| 9 | |
| 10 | def add_word(self, word): |
| 11 | if word not in self.word2idx: |
| 12 | self.idx2word.append(word) |
| 13 | self.word2idx[word] = len(self.idx2word) - 1 |
| 14 | return self.word2idx[word] |
| 15 | |
| 16 | def __len__(self): |
| 17 | return len(self.idx2word) |
| 18 | |
| 19 | |
| 20 | class Corpus(object): |