entities extraction bugfix
@@ -7,13 +7,16 @@ class NERDL:
def get_default_model(): \
return NerDLModel.pretrained(name='ner_dl_bert', lang='en') \
.setInputCols(["sentence", "token", "word_embeddings"]) \
- .setOutputCol("ner")
+ .setOutputCol("ner") \
+ .setIncludeConfidence(True)
+
@staticmethod
def get_pretrained_model(name, language):
return NerDLModel.pretrained(name,language) \
.setInputCols(["sentence", "token", "word_embeddings"]) \
- .setOutputCol("ner")
+ .setOutputCol("ner") \
+ .setIncludeConfidence(True)
@staticmethod
def get_default_trainable_model():
@@ -22,4 +25,5 @@ class NERDL:
.setLabelColumn("y") \
.setOutputCol("ner") \
.setMaxEpochs(2) \
- .setVerbose(0)
\ No newline at end of file
+ .setVerbose(0) \
+ .setIncludeConfidence(True)
\ No newline at end of file
