adjust predictor call

supercoderhawk · supercoderhawk · commit 629a4fc7e746 · 2020-02-21T00:32:35.000+08:00
diff --git a/deep_keyphrase/copy_rnn/predict.py b/deep_keyphrase/copy_rnn/predict.py
@@ -8,7 +8,7 @@
 from deep_keyphrase.copy_rnn.beam_search import BeamSearch
 from deep_keyphrase.dataloader import KeyphraseDataLoader, RAW_BATCH, TOKENS, INFERENCE_MODE, EVAL_MODE
 from deep_keyphrase.utils.vocab_loader import load_vocab
-from deep_keyphrase.utils.constants import BOS_WORD
+from deep_keyphrase.utils.constants import BOS_WORD, UNK_WORD
 from deep_keyphrase.utils.tokenizer import token_char_tokenize
 
 
@@ -47,6 +47,7 @@ def __init__(self, model_info, vocab_info, beam_size, max_target_len, max_src_le
                                         max_target_len=self.max_target_len,
                                         id2vocab=self.id2vocab,
                                         bos_idx=self.vocab2id[BOS_WORD],
+                                        unk_idx=self.vocab2id[UNK_WORD],
                                         args=self.config)
         self.pred_base_config = {'max_oov_count': self.config.max_oov_count,
                                  'max_src_len': self.max_src_len,
@@ -74,7 +75,7 @@ def predict(self, text_list, batch_size=10, delimiter=None, tokenized=False):
             text_list = [{TOKENS: i} for i in text_list]
         else:
             text_list = [{TOKENS: token_char_tokenize(i)} for i in text_list]
-        args = Munch({'batch_size': batch_size, **self.pred_base_config})
+        args = Munch({'batch_size': batch_size, **self.config._asdict(), **self.pred_base_config})
         loader = KeyphraseDataLoader(data_source=text_list,
                                      vocab2id=self.vocab2id,
                                      mode=INFERENCE_MODE,
@@ -106,6 +107,7 @@ def eval_predict(self, src_filename, dest_filename, args,
                                             max_target_len=self.max_target_len,
                                             id2vocab=self.id2vocab,
                                             bos_idx=self.vocab2id[BOS_WORD],
+                                            unk_idx=self.vocab2id[UNK_WORD],
                                             args=self.config)
 
         for batch in loader:
