Encoder based on a user-supplied vocabulary (file or list).
| 155 | |
| 156 | |
| 157 | class TokenTextEncoder(TextEncoder): |
| 158 | """Encoder based on a user-supplied vocabulary (file or list).""" |
| 159 | |
| 160 | def __init__(self, |
| 161 | vocab_filename, |
| 162 | reverse=False, |
| 163 | vocab_list=None, |
| 164 | replace_oov=None, |
| 165 | num_reserved_ids=NUM_RESERVED_TOKENS): |
| 166 | """Initialize from a file or list, one token per line. |
| 167 | |
| 168 | Handling of reserved tokens works as follows: |
| 169 | - When initializing from a list, we add reserved tokens to the vocab. |
| 170 | - When initializing from a file, we do not add reserved tokens to the vocab. |
| 171 | - When saving vocab files, we save reserved tokens to the file. |
| 172 | |
| 173 | Args: |
| 174 | vocab_filename: If not None, the full filename to read vocab from. If this |
| 175 | is not None, then vocab_list should be None. |
| 176 | reverse: Boolean indicating if tokens should be reversed during encoding |
| 177 | and decoding. |
| 178 | vocab_list: If not None, a list of elements of the vocabulary. If this is |
| 179 | not None, then vocab_filename should be None. |
| 180 | replace_oov: If not None, every out-of-vocabulary token seen when |
| 181 | encoding will be replaced by this string (which must be in vocab). |
| 182 | num_reserved_ids: Number of IDs to save for reserved tokens like <EOS>. |
| 183 | """ |
| 184 | super(TokenTextEncoder, self).__init__(num_reserved_ids=num_reserved_ids) |
| 185 | self._reverse = reverse |
| 186 | self._replace_oov = replace_oov |
| 187 | if vocab_filename: |
| 188 | self._init_vocab_from_file(vocab_filename) |
| 189 | else: |
| 190 | assert vocab_list is not None |
| 191 | self._init_vocab_from_list(vocab_list) |
| 192 | self.pad_index = self._token_to_id[PAD] |
| 193 | self.eos_index = self._token_to_id[EOS] |
| 194 | self.unk_index = self._token_to_id[UNK] |
| 195 | self.seg_index = self._token_to_id[SEG] if SEG in self._token_to_id else self.eos_index |
| 196 | |
| 197 | def encode(self, s): |
| 198 | """Converts a space-separated string of tokens to a list of ids.""" |
| 199 | sentence = s |
| 200 | tokens = sentence.strip().split() |
| 201 | if self._replace_oov is not None: |
| 202 | tokens = [t if t in self._token_to_id else self._replace_oov |
| 203 | for t in tokens] |
| 204 | ret = [self._token_to_id[tok] for tok in tokens] |
| 205 | return ret[::-1] if self._reverse else ret |
| 206 | |
| 207 | def decode(self, ids, strip_eos=False, strip_padding=False): |
| 208 | if strip_padding and self.pad() in list(ids): |
| 209 | pad_pos = list(ids).index(self.pad()) |
| 210 | ids = ids[:pad_pos] |
| 211 | if strip_eos and self.eos() in list(ids): |
| 212 | eos_pos = list(ids).index(self.eos()) |
| 213 | ids = ids[:eos_pos] |
| 214 | return " ".join(self.decode_list(ids)) |
no outgoing calls
no test coverage detected