Vocab
| 28 | |
| 29 | |
| 30 | class Vocab(object): |
| 31 | """Vocab""" |
| 32 | def __init__(self, counter, min_freq=1, specials=None, unk_index=0): |
| 33 | self.itos = list(specials) if specials else [] |
| 34 | self.stoi = defaultdict(lambda: unk_index) |
| 35 | self.stoi.update({token: i for i, token in enumerate(self.itos)}) |
| 36 | self.extend([token for token, freq in counter.items() if freq >= min_freq]) |
| 37 | self.unk_index = unk_index |
| 38 | self.n_init = len(self) |
| 39 | |
| 40 | def __len__(self): |
| 41 | """Returns the size of the vocabulary""" |
| 42 | return len(self.itos) |
| 43 | |
| 44 | def __getitem__(self, key): |
| 45 | """According to the key or index, return the index and key""" |
| 46 | if isinstance(key, six.string_types): |
| 47 | return self.stoi[key] |
| 48 | elif not isinstance(key, Iterable): |
| 49 | return self.itos[key] |
| 50 | elif isinstance(key[0], six.string_types): |
| 51 | return [self.stoi[i] for i in key] |
| 52 | else: |
| 53 | return [self.itos[i] for i in key] |
| 54 | |
| 55 | def __contains__(self, token): |
| 56 | """contains""" |
| 57 | return token in self.stoi |
| 58 | |
| 59 | def __getstate__(self): |
| 60 | """getstate""" |
| 61 | # avoid picking defaultdict |
| 62 | attrs = dict(self.__dict__) |
| 63 | # cast to regular dict |
| 64 | attrs['stoi'] = dict(self.stoi) |
| 65 | return attrs |
| 66 | |
| 67 | def __setstate__(self, state): |
| 68 | """setstate""" |
| 69 | stoi = defaultdict(lambda: self.unk_index) |
| 70 | stoi.update(state['stoi']) |
| 71 | state['stoi'] = stoi |
| 72 | self.__dict__.update(state) |
| 73 | |
| 74 | def extend(self, tokens): |
| 75 | """Update tokens to itos and stoi""" |
| 76 | self.itos.extend(sorted(set(tokens).difference(self.stoi))) |
| 77 | self.stoi.update({token: i for i, token in enumerate(self.itos)}) |