The class of vocab, include 2 dicts of token to index and index to token
| 23 | |
| 24 | |
| 25 | class Vocab: |
| 26 | """ |
| 27 | The class of vocab, include 2 dicts of token to index and index to token |
| 28 | """ |
| 29 | def __init__(self, sentences): |
| 30 | """ |
| 31 | Args: |
| 32 | sentences: a 2-dim list |
| 33 | """ |
| 34 | flatten = lambda lst: [item for sublist in lst for item in sublist] |
| 35 | self.sentence = sentences |
| 36 | self.token2index = {'<pad>': 0, '<bos>': 1, '<eos>': 2, '<unk>': 3} |
| 37 | self.token2index.update({ |
| 38 | token: index + 4 |
| 39 | for index, (token, freq) in enumerate( |
| 40 | sorted(Counter(flatten(self.sentence)).items(), key=lambda x: x[1], reverse=True)) |
| 41 | }) |
| 42 | self.index2token = {index: token for token, index in self.token2index.items()} |
| 43 | |
| 44 | def __getitem__(self, query): |
| 45 | if isinstance(query, str): |
| 46 | return self.token2index.get(query, self.token2index.get('<unk>')) |
| 47 | elif isinstance(query, (int, np.int32, np.int64)): |
| 48 | return self.index2token.get(query, '<unk>') |
| 49 | elif isinstance(query, (list, tuple, np.ndarray)): |
| 50 | return [self.__getitem__(item) for item in query] |
| 51 | else: |
| 52 | raise ValueError("The type of query is invalid.") |
| 53 | |
| 54 | def __len__(self): |
| 55 | return len(self.index2token) |
| 56 | |
| 57 | |
| 58 | class CmnDataset: |