Encoder based on a user-supplied vocabulary (file or list).
| 119 | |
| 120 | |
| 121 | class TokenTextEncoder(TextEncoder): |
| 122 | """Encoder based on a user-supplied vocabulary (file or list).""" |
| 123 | |
| 124 | def __init__(self, |
| 125 | vocab_filename, |
| 126 | reverse=False, |
| 127 | vocab_list=None, |
| 128 | replace_oov=None, |
| 129 | num_reserved_ids=NUM_RESERVED_TOKENS): |
| 130 | """Initialize from a file or list, one token per line. |
| 131 | |
| 132 | Handling of reserved tokens works as follows: |
| 133 | - When initializing from a list, we add reserved tokens to the vocab. |
| 134 | - When initializing from a file, we do not add reserved tokens to the vocab. |
| 135 | - When saving vocab files, we save reserved tokens to the file. |
| 136 | |
| 137 | Args: |
| 138 | vocab_filename: If not None, the full filename to read vocab from. If this |
| 139 | is not None, then vocab_list should be None. |
| 140 | reverse: Boolean indicating if tokens should be reversed during encoding |
| 141 | and decoding. |
| 142 | vocab_list: If not None, a list of elements of the vocabulary. If this is |
| 143 | not None, then vocab_filename should be None. |
| 144 | replace_oov: If not None, every out-of-vocabulary token seen when |
| 145 | encoding will be replaced by this string (which must be in vocab). |
| 146 | num_reserved_ids: Number of IDs to save for reserved tokens like <EOS>. |
| 147 | """ |
| 148 | super(TokenTextEncoder, self).__init__(num_reserved_ids=num_reserved_ids) |
| 149 | self._reverse = reverse |
| 150 | self._replace_oov = replace_oov |
| 151 | if vocab_filename: |
| 152 | self._init_vocab_from_file(vocab_filename) |
| 153 | else: |
| 154 | assert vocab_list is not None |
| 155 | self._init_vocab_from_list(vocab_list) |
| 156 | self.pad_index = self.token_to_id[PAD] |
| 157 | self.eos_index = self.token_to_id[EOS] |
| 158 | self.unk_index = self.token_to_id[UNK] |
| 159 | self.seg_index = self.token_to_id[SEG] if SEG in self.token_to_id else self.eos_index |
| 160 | |
| 161 | def encode(self, s): |
| 162 | """Converts a space-separated string of tokens to a list of ids.""" |
| 163 | if isinstance(s, str): |
| 164 | sentence = s |
| 165 | tokens = sentence.strip().split() |
| 166 | else: |
| 167 | tokens = s |
| 168 | if self._replace_oov is not None: |
| 169 | tokens = [t if t in self.token_to_id else self._replace_oov |
| 170 | for t in tokens] |
| 171 | ret = [self.token_to_id[tok] for tok in tokens] |
| 172 | return ret[::-1] if self._reverse else ret |
| 173 | |
| 174 | def decode(self, ids, strip_eos=False, strip_padding=False): |
| 175 | if strip_padding and self.pad() in list(ids): |
| 176 | pad_pos = list(ids).index(self.pad()) |
| 177 | ids = ids[:pad_pos] |
| 178 | if strip_eos and self.eos() in list(ids): |
no outgoing calls
no test coverage detected