🐛 Add Non Masking layer to BERT.

BrikerMan · May 22, 2019 · dce5719 · dce5719
1 parent 174d924
commit dce5719
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/kashgari/embeddings/bert_embedding.py b/kashgari/embeddings/bert_embedding.py
@@ -109,7 +109,9 @@ def _build_model(self, **kwargs):
                     else:
                         self.sequence_length = (bert_seq_len,)
             self.embedding_size = int(bert_model.output.shape[-1])
-            self.embed_model = bert_model
+            output_features = NonMaskingLayer()(bert_model.output)
+
+            self.embed_model = tf.keras.Model(bert_model.inputs, output_features)
 
     def analyze_corpus(self,
                        x: Union[Tuple[List[List[str]], ...], List[List[str]]],