Update model.py

atulkum · web-flow · commit e5608551067c · 2019-04-10T15:12:20.000-07:00
diff --git a/neural_ner/model.py b/neural_ner/model.py
@@ -8,7 +8,7 @@
 import logging
 
 from crf import CRF_Loss
-from model_utils import init_lstm_wt, init_linear_wt, get_word_embd
+from model_utils import init_lstm_wt, init_linear_wt, get_word_embd, get_mask
 
 logging.basicConfig(level=logging.INFO)
 
@@ -158,6 +158,6 @@ def get_loss(self, logits, y, s_lens):
         return loss
 
     def predict(self, emissions, lengths):
-
+        mask = get_mask(lengths, self.config)
         best_scores, pred = self.crf.viterbi_decode(emissions, mask)
         return pred
