fix gpt-2 unk token test

2019-07-12 11:46:57 +02:00
parent 762ded9b1c
commit 699bc7e86e
2 changed files with 4 additions and 2 deletions
--- a/pytorch_transformers/tokenization_gpt2.py
+++ b/pytorch_transformers/tokenization_gpt2.py
@@ -177,7 +177,9 @@ class GPT2Tokenizer(PreTrainedTokenizer):

    def _convert_token_to_id(self, token):
        """ Converts a token (str/unicode) in an id using the vocab. """
-        return self.encoder.get(token)
+        if token in self.encoder:
+            return self.encoder.get(token)
+        return self.encoder.get(self.unk_token)

    def _convert_id_to_token(self, index):
        """Converts an index (integer) in a token (string/unicode) using the vocab."""