Extract relations between problem, treatment and test entities (ReDL)

Description

Extract relations like TrIP : a certain treatment has improved a medical problem and 7 other such relations between problem, treatment and test entities.

Predicted Entities

PIP, TeCP, TeRP, TrAP, TrCP, TrIP, TrNAP, TrWP

Live Demo Open in Colab Download

How to use

...
words_embedder = WordEmbeddingsModel() \
    .pretrained("embeddings_clinical", "en", "clinical/models") \
    .setInputCols(["sentences", "tokens"]) \
    .setOutputCol("embeddings")
ner_tagger = NerDLModel() \
    .pretrained("ner_clinical", "en", "clinical/models") \
    .setInputCols(["sentences", "tokens", "embeddings"]) \
    .setOutputCol("ner_tags")
ner_converter = NerConverter() \
    .setInputCols(["sentences", "tokens", "ner_tags"]) \
    .setOutputCol("ner_chunks")
dependency_parser = DependencyParserModel() \
    .pretrained("dependency_conllu", "en") \
    .setInputCols(["sentences", "pos_tags", "tokens"]) \
    .setOutputCol("dependencies")

# Set a filter on pairs of named entities which will be treated as relation candidates
re_ner_chunk_filter = RENerChunksFilter() \
    .setInputCols(["ner_chunks", "dependencies"])\
    .setMaxSyntacticDistance(10)\
    .setOutputCol("re_ner_chunks").setRelationPairs(['SYMPTOM-EXTERNAL_BODY_PART_OR_REGION'])

# The dataset this model is trained to is sentence-wise. 
# This model can also be trained on document-level relations - in which case, while predicting, use "document" instead of "sentence" as input.
re_model = RelationExtractionDLModel()\
    .pretrained('redl_clinical_biobert', 'en', "clinical/models") \
    .setPredictionThreshold(0.5)\
    .setInputCols(["re_ner_chunks", "sentences"]) \
    .setOutputCol("relations")
pipeline = Pipeline(stages=[documenter, sentencer, tokenizer, pos_tagger, words_embedder, ner_tagger, ner_converter, dependency_parser, re_ner_chunk_filter, re_model])

text ="""A 28-year-old female with a history of gestational diabetes mellitus diagnosed eight years prior to presentation and subsequent type two diabetes mellitus ( T2DM ), one prior episode of HTG-induced pancreatitis three years prior to presentation,  associated with an acute hepatitis , and obesity with a body mass index ( BMI ) of 33.5 kg/m2 , presented with a one-week history of polyuria , polydipsia , poor appetite , and vomiting .
"""
data = spark.createDataFrame([[text]]).toDF("text")
p_model = pipeline.fit(data)
result = p_model.transform(data)
...
val words_embedder = WordEmbeddingsModel()
    .pretrained("embeddings_clinical", "en", "clinical/models")
    .setInputCols(Array("sentences", "tokens"))
    .setOutputCol("embeddings")
val ner_tagger = NerDLModel()
    .pretrained("ner_clinical", "en", "clinical/models")
    .setInputCols(Array("sentences", "tokens", "embeddings"))
    .setOutputCol("ner_tags")
val ner_converter = NerConverter()
    .setInputCols(Array("sentences", "tokens", "ner_tags"))
    .setOutputCol("ner_chunks")
val dependency_parser = DependencyParserModel()
    .pretrained("dependency_conllu", "en")
    .setInputCols(Array("sentences", "pos_tags", "tokens"))
    .setOutputCol("dependencies")

// Set a filter on pairs of named entities which will be treated as relation candidates
val re_ner_chunk_filter = RENerChunksFilter()
    .setInputCols(Array("ner_chunks", "dependencies"))
    .setMaxSyntacticDistance(10)
    .setOutputCol("re_ner_chunks").setRelationPairs(Array("SYMPTOM-EXTERNAL_BODY_PART_OR_REGION"))

// The dataset this model is trained to is sentence-wise. 
// This model can also be trained on document-level relations - in which case, while predicting, use "document" instead of "sentence" as input.
val re_model = RelationExtractionDLModel()
    .pretrained("redl_clinical_biobert", "en", "clinical/models")
    .setPredictionThreshold(0.5)
    .setInputCols(Array("re_ner_chunks", "sentences"))
    .setOutputCol("relations")
val pipeline = new Pipeline().setStages(Array(documenter, sentencer, tokenizer, pos_tagger, words_embedder, ner_tagger, ner_converter, dependency_parser, re_ner_chunk_filter, re_model))

val result = pipeline.fit(Seq.empty["A 28-year-old female with a history of gestational diabetes mellitus diagnosed eight years prior to presentation and subsequent type two diabetes mellitus ( T2DM ), one prior episode of HTG-induced pancreatitis three years prior to presentation,  associated with an acute hepatitis , and obesity with a body mass index ( BMI ) of 33.5 kg/m2 , presented with a one-week history of polyuria , polydipsia , poor appetite , and vomiting ."].toDS.toDF("text")).transform(data)

Results

|    | relation   | entity1   |   entity1_begin |   entity1_end | chunk1                                | entity2   |   entity2_begin |   entity2_end | chunk2                   |   confidence |
|---:|:-----------|:----------|----------------:|--------------:|:--------------------------------------|:----------|----------------:|--------------:|:-------------------------|-------------:|
|  0 | PIP        | PROBLEM   |              39 |            67 | gestational diabetes mellitus         | PROBLEM   |             157 |           160 | T2DM                     |     0.763447 |
|  1 | PIP        | PROBLEM   |              39 |            67 | gestational diabetes mellitus         | PROBLEM   |             289 |           295 | obesity                  |     0.682246 |
|  2 | PIP        | PROBLEM   |             117 |           153 | subsequent type two diabetes mellitus | PROBLEM   |             187 |           210 | HTG-induced pancreatitis |     0.610396 |
|  3 | PIP        | PROBLEM   |             117 |           153 | subsequent type two diabetes mellitus | PROBLEM   |             264 |           281 | an acute hepatitis       |     0.726894 |

Model Information

Model Name: redl_clinical_biobert
Compatibility: Spark NLP 2.7.3+
License: Licensed
Edition: Official
Language: en

Data Source

Trained on 2010 i2b2 relation challenge.

Benchmarking

Relation           Recall Precision        F1   Support
PIP                 0.859     0.878     0.869      1435
TeCP                0.629     0.782     0.697       337
TeRP                0.903     0.929     0.916      2034
TrAP                0.872     0.866     0.869      1693
TrCP                0.641     0.677     0.659       340
TrIP                0.517     0.796     0.627       151
TrNAP               0.402     0.672     0.503       112
TrWP                0.257     0.824     0.392       109
Avg.                0.635     0.803     0.691