Description
DistilBERT Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks.
distilbert_base_token_classifier_ontonotes is a fine-tuned DistilBERT model that is ready to use for Named Entity Recognition and achieves state-of-the-art performance for the NER task. This model has been trained to recognize four types of entities: CARDINAL, DATE, EVENT, FAC, GPE, LANGUAGE, LAW, LOC, MONEY, NORP, ORDINAL, ORG, PERCENT, PERSON, PRODUCT, QUANTITY, TIME, and WORK_OF_ART.
We used TFDistilBertForTokenClassification to train this model and used DistilBertForTokenClassification annotator in Spark NLP 🚀 for prediction at scale!
Predicted Entities
CARDINAL, DATE, EVENT, FAC, GPE, LANGUAGE, LAW, LOC, MONEY, NORP, ORDINAL, ORG, PERCENT, PERSON, PRODUCT, QUANTITY, TIME, WORK_OF_ART
How to use
document_assembler = DocumentAssembler() \
.setInputCol('text') \
.setOutputCol('document')
tokenizer = Tokenizer() \
.setInputCols(['document']) \
.setOutputCol('token')
tokenClassifier = DistilBertForTokenClassification \
.pretrained('distilbert_base_token_classifier_ontonotes', 'en') \
.setInputCols(['token', 'document']) \
.setOutputCol('ner') \
.setCaseSensitive(True) \
.setMaxSentenceLength(512)
# since output column is IOB/IOB2 style, NerConverter can extract entities
ner_converter = NerConverter() \
.setInputCols(['document', 'token', 'ner']) \
.setOutputCol('entities')
pipeline = Pipeline(stages=[
document_assembler,
tokenizer,
tokenClassifier,
ner_converter
])
example = spark.createDataFrame([['My name is John!']]).toDF("text")
result = pipeline.fit(example).transform(example)
val document_assembler = DocumentAssembler()
.setInputCol("text")
.setOutputCol("document")
val tokenizer = Tokenizer()
.setInputCols("document")
.setOutputCol("token")
val tokenClassifier = DistilBertForTokenClassification.pretrained("distilbert_base_token_classifier_ontonotes", "en")
.setInputCols("document", "token")
.setOutputCol("ner")
.setCaseSensitive(true)
.setMaxSentenceLength(512)
// since output column is IOB/IOB2 style, NerConverter can extract entities
val ner_converter = NerConverter()
.setInputCols("document", "token", "ner")
.setOutputCol("entities")
val pipeline = new Pipeline().setStages(Array(document_assembler, tokenizer, tokenClassifier, ner_converter))
val example = Seq.empty["My name is John!"].toDS.toDF("text")
val result = pipeline.fit(example).transform(example)
import nlu
nlu.load("en.ner.distil_bert.ontonotes.base").predict("""My name is John!""")
Results
+------------------------------------------------------------------------------------+
|result |
+------------------------------------------------------------------------------------+
|[B-PERSON, I-PERSON, O, O, O, B-LOC, O, O, O, B-LOC, O, O, O, O, B-PERSON, O, O, O, O, B-LOC]|
+------------------------------------------------------------------------------------+
Model Information
| Model Name: | distilbert_base_token_classifier_ontonotes |
| Compatibility: | Spark NLP 3.2.0+ |
| License: | Open Source |
| Edition: | Official |
| Input Labels: | [token, document] |
| Output Labels: | [ner] |
| Language: | en |
| Case sensitive: | true |
| Max sentense length: | 512 |
Data Source
https://catalog.ldc.upenn.edu/LDC2013T19
Benchmarking
Test:
precision recall f1-score support
B-CARDINAL 0.85 0.86 0.86 935
B-DATE 0.87 0.90 0.88 1602
B-EVENT 0.73 0.59 0.65 63
B-FAC 0.72 0.71 0.71 135
B-GPE 0.96 0.91 0.94 2240
B-LANGUAGE 0.92 0.55 0.69 22
B-LAW 0.80 0.60 0.69 40
B-LOC 0.78 0.78 0.78 179
B-MONEY 0.89 0.89 0.89 314
B-NORP 0.90 0.93 0.91 841
B-ORDINAL 0.81 0.90 0.86 195
B-ORG 0.87 0.89 0.88 1795
B-PERCENT 0.94 0.93 0.93 349
B-PERSON 0.93 0.94 0.94 1988
B-PRODUCT 0.59 0.74 0.65 76
B-QUANTITY 0.80 0.78 0.79 105
B-TIME 0.71 0.64 0.67 212
B-WORK_OF_ART 0.63 0.67 0.65 166
I-CARDINAL 0.80 0.91 0.85 331
I-DATE 0.88 0.91 0.90 2011
I-EVENT 0.71 0.71 0.71 130
I-FAC 0.80 0.86 0.83 213
I-GPE 0.92 0.90 0.91 628
I-LAW 0.93 0.61 0.74 106
I-LOC 0.83 0.81 0.82 180
I-MONEY 0.93 0.96 0.95 685
I-NORP 0.95 0.76 0.84 160
I-ORDINAL 0.00 0.00 0.00 4
I-ORG 0.89 0.93 0.91 2406
I-PERCENT 0.93 0.97 0.95 523
I-PERSON 0.95 0.93 0.94 1412
I-PRODUCT 0.64 0.78 0.70 69
I-QUANTITY 0.80 0.88 0.84 206
I-TIME 0.71 0.79 0.75 255
I-WORK_OF_ART 0.69 0.63 0.65 337
O 0.99 0.99 0.99 131815
accuracy 0.98 152728
macro avg 0.81 0.79 0.80 152728
weighted avg 0.98 0.98 0.98 152728
processed 152728 tokens with 11257 phrases; found: 11586 phrases; correct: 9747.
accuracy: 89.24%; (non-O)
accuracy: 97.80%; precision: 84.13%; recall: 86.59%; FB1: 85.34
CARDINAL: precision: 82.90%; recall: 85.03%; FB1: 83.95 959
DATE: precision: 80.79%; recall: 86.64%; FB1: 83.61 1718
EVENT: precision: 57.63%; recall: 53.97%; FB1: 55.74 59
FAC: precision: 63.01%; recall: 68.15%; FB1: 65.48 146
GPE: precision: 93.96%; recall: 90.31%; FB1: 92.10 2153
LANGUAGE: precision: 92.31%; recall: 54.55%; FB1: 68.57 13
LAW: precision: 55.56%; recall: 50.00%; FB1: 52.63 36
LOC: precision: 71.35%; recall: 76.54%; FB1: 73.85 192
MONEY: precision: 84.00%; recall: 86.94%; FB1: 85.45 325
NORP: precision: 87.96%; recall: 91.20%; FB1: 89.55 872
ORDINAL: precision: 81.48%; recall: 90.26%; FB1: 85.64 216
ORG: precision: 81.64%; recall: 85.46%; FB1: 83.51 1879
PERCENT: precision: 88.42%; recall: 89.68%; FB1: 89.05 354
PERSON: precision: 90.78%; recall: 92.10%; FB1: 91.44 2017
PRODUCT: precision: 50.00%; recall: 68.42%; FB1: 57.78 104
QUANTITY: precision: 66.95%; recall: 75.24%; FB1: 70.85 118
TIME: precision: 56.95%; recall: 59.91%; FB1: 58.39 223
WORK_OF_ART: precision: 46.53%; recall: 56.63%; FB1: 51.09 202