Refine Vocab.to_indice (#641)

Author: guoshengCS

paddlenlp/data/vocab.py

@@ -92,13 +92,19 @@ def __init__(self,
                 assert special_token in token_to_idx, '{} is not in token_to_idx'.format(
                     special_token)
             self._token_to_idx = token_to_idx
-            self._idx_to_token = {idx: token for token, idx in token_to_idx.items()}
+            self._idx_to_token = {
+                idx: token
+                for token, idx in token_to_idx.items()
+            }
             if unk_token:
                 unk_index = self._token_to_idx[unk_token]
                 self._token_to_idx = collections.defaultdict(lambda: unk_index)
                 self._token_to_idx.update(token_to_idx)
         else:
-            self._idx_to_token = {idx: special_token for idx, special_token in enumerate(special_tokens)}
+            self._idx_to_token = {
+                idx: special_token
+                for idx, special_token in enumerate(special_tokens)
+            }
             self._token_to_idx = collections.defaultdict()
             self._token_to_idx.update(
                 (token, idx) for idx, token in self._idx_to_token.items())
@@ -136,7 +142,8 @@ def _index_counter_keys(self, counter, special_tokens, max_size, min_freq):
             if freq < min_freq or len(self._idx_to_token) == max_size:
                 break
             if token not in special_tokens:
-                self._idx_to_token[max(list(self._idx_to_token.keys()) + [-1]) + 1] = token
+                self._idx_to_token[max(list(self._idx_to_token.keys()) + [-1]) +
+                                   1] = token
                 self._token_to_idx[token] = max(self._idx_to_token.keys())
 
     def _sort_index_according_to_user_specification(self, token_to_idx):
@@ -206,20 +213,20 @@ def to_tokens(self, indices):
                 'Token indices is invalid. Expected 1D array, but received {}D array. '.
                 format(len(indices.shape)))
 
-        max_idx = max(self._idx_to_token.keys())
-
         tokens = []
         for idx in indices:
             if not isinstance(idx, int):
                 warnings.warn(
                     "The type of `to_tokens()`'s input `indices` is not `int` which will be forcibly transfered to `int`. "
                 )
                 idx = int(idx)
-            if idx > max_idx:
+
+            try:
+                tokens.append(self._idx_to_token[idx])
+            except KeyError:
                 raise ValueError(
                     'Token index {} in the provided `indices` is invalid.'.
                     format(idx))
-            tokens.append(self._idx_to_token[idx])
 
         return tokens[0] if to_reduce else tokens