fix padding error

supercoderhawk · Jul 15, 2020 · ac1c338 · ac1c338
1 parent 0bca198
commit ac1c338
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/deep_keyphrase/copy_rnn/predict_tf.py b/deep_keyphrase/copy_rnn/predict_tf.py
@@ -102,8 +102,9 @@ def generate_input(self, tokens):
         sent_len = len(token_ids)
 
         if len(token_ids) < self.max_src_len:
-            token_ids.extend([PAD_WORD] * (self.max_src_len - len(token_ids)))
-            token_ids_with_oov.extend([PAD_WORD] * (self.max_src_len - len(token_ids)))
+            pad_tokens = [self.vocab2id[PAD_WORD]] * (self.max_src_len - len(token_ids))
+            token_ids.extend(pad_tokens)
+            token_ids_with_oov.extend(pad_tokens)
         elif len(token_ids) > self.max_src_len:
             token_ids = token_ids[:self.max_src_len]
             token_ids_with_oov = token_ids_with_oov[:self.max_src_len]