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

...
documenter = DocumentAssembler()\
    .setInputCol("text")\
    .setOutputCol("document")

sentencer = SentenceDetector()\
    .setInputCols(["document"])\
    .setOutputCol("sentences")

tokenizer = sparknlp.annotators.Tokenizer()\
    .setInputCols(["sentences"])\
    .setOutputCol("tokens")

pos_tagger = PerceptronModel()\
    .pretrained("pos_clinical", "en", "clinical/models") \
    .setInputCols(["sentences", "tokens"])\
    .setOutputCol("pos_tags")

words_embedder = WordEmbeddingsModel() \
    .pretrained("embeddings_clinical", "en", "clinical/models") \
    .setInputCols(["sentences", "tokens"]) \
    .setOutputCol("embeddings")

ner_tagger = MedicalNerModel() \
    .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(["problem-test", "problem-treatment"])

# 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 . Two weeks prior to presentation , she was treated with a five-day course of amoxicillin for a respiratory tract infection . She was on metformin , glipizide , and dapagliflozin for T2DM and atorvastatin and gemfibrozil for HTG . She had been on dapagliflozin for six months at the time of presentation. Physical examination on presentation was significant for dry oral mucosa ; significantly , her abdominal examination was benign with no tenderness , guarding , or rigidity . Pertinent laboratory findings on admission were : serum glucose 111 mg/dl , bicarbonate 18 mmol/l , anion gap 20 , creatinine 0.4 mg/dL , triglycerides 508 mg/dL , total cholesterol 122 mg/dL , glycated hemoglobin ( HbA1c ) 10% , and venous pH 7.27 . Serum lipase was normal at 43 U/L . Serum acetone levels could not be assessed as blood samples kept hemolyzing due to significant lipemia . The patient was initially admitted for starvation ketosis , as she reported poor oral intake for three days prior to admission . However , serum chemistry obtained six hours after presentation revealed her glucose was 186 mg/dL , the anion gap was still elevated at 21 , serum bicarbonate was 16 mmol/L , triglyceride level peaked at 2050 mg/dL , and lipase was 52 U/L . The β-hydroxybutyrate level was obtained and found to be elevated at 5.29 mmol/L - the original sample was centrifuged and the chylomicron layer removed prior to analysis due to interference from turbidity caused by lipemia again . The patient was treated with an insulin drip for euDKA and HTG with a reduction in the anion gap to 13 and triglycerides to 1400 mg/dL , within 24 hours . Her euDKA was thought to be precipitated by her respiratory tract infection in the setting of SGLT2 inhibitor use . The patient was seen by the endocrinology service and she was discharged on 40 units of insulin glargine at night , 12 units of insulin lispro with meals , and metformin 1000 mg two times a day . It was determined that all SGLT2 inhibitors should be discontinued indefinitely . 
She had close follow-up with endocrinology post discharge .
"""
data = spark.createDataFrame([[text]]).toDF("text")
p_model = pipeline.fit(data)
result = p_model.transform(data)
...
val documenter = DocumentAssembler() 
    .setInputCol("text") 
    .setOutputCol("document")

val sentencer = SentenceDetector()
    .setInputCols("document")
    .setOutputCol("sentences")

val tokenizer = sparknlp.annotators.Tokenizer()
    .setInputCols("sentences")
    .setOutputCol("tokens")

val pos_tagger = PerceptronModel()
    .pretrained("pos_clinical", "en", "clinical/models") 
    .setInputCols(Array("sentences", "tokens"))
    .setOutputCol("pos_tags")

val words_embedder = WordEmbeddingsModel()
    .pretrained("embeddings_clinical", "en", "clinical/models")
    .setInputCols(Array("sentences", "tokens"))
    .setOutputCol("embeddings")

val ner_tagger = MedicalNerModel.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("problem-test", "problem-treatment"))

// 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 data = Seq("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 . Two weeks prior to presentation , she was treated with a five-day course of amoxicillin for a respiratory tract infection . She was on metformin , glipizide , and dapagliflozin for T2DM and atorvastatin and gemfibrozil for HTG . She had been on dapagliflozin for six months at the time of presentation. Physical examination on presentation was significant for dry oral mucosa ; significantly , her abdominal examination was benign with no tenderness , guarding , or rigidity . Pertinent laboratory findings on admission were : serum glucose 111 mg/dl , bicarbonate 18 mmol/l , anion gap 20 , creatinine 0.4 mg/dL , triglycerides 508 mg/dL , total cholesterol 122 mg/dL , glycated hemoglobin ( HbA1c ) 10% , and venous pH 7.27 . Serum lipase was normal at 43 U/L . Serum acetone levels could not be assessed as blood samples kept hemolyzing due to significant lipemia . The patient was initially admitted for starvation ketosis , as she reported poor oral intake for three days prior to admission . However , serum chemistry obtained six hours after presentation revealed her glucose was 186 mg/dL , the anion gap was still elevated at 21 , serum bicarbonate was 16 mmol/L , triglyceride level peaked at 2050 mg/dL , and lipase was 52 U/L . The β-hydroxybutyrate level was obtained and found to be elevated at 5.29 mmol/L - the original sample was centrifuged and the chylomicron layer removed prior to analysis due to interference from turbidity caused by lipemia again . The patient was treated with an insulin drip for euDKA and HTG with a reduction in the anion gap to 13 and triglycerides to 1400 mg/dL , within 24 hours . Her euDKA was thought to be precipitated by her respiratory tract infection in the setting of SGLT2 inhibitor use . The patient was seen by the endocrinology service and she was discharged on 40 units of insulin glargine at night , 12 units of insulin lispro with meals , and metformin 1000 mg two times a day . It was determined that all SGLT2 inhibitors should be discontinued indefinitely . She had close follow-up with endocrinology post discharge .").toDF("text")
val result = pipeline.fit(data).transform(data)

Results

+--------+---------+-------------+-----------+--------------------+---------+-------------+-----------+--------------------+----------+
|relation|  entity1|entity1_begin|entity1_end|              chunk1|  entity2|entity2_begin|entity2_end|              chunk2|confidence|
+--------+---------+-------------+-----------+--------------------+---------+-------------+-----------+--------------------+----------+
|    TrAP|TREATMENT|          512|        522|         amoxicillin|  PROBLEM|          528|        556|a respiratory tra...|0.99796957|
|    TrAP|TREATMENT|          571|        579|           metformin|  PROBLEM|          617|        620|                T2DM|0.99757993|
|    TrAP|TREATMENT|          599|        611|       dapagliflozin|  PROBLEM|          659|        661|                 HTG|  0.996036|
|    TrAP|  PROBLEM|          617|        620|                T2DM|TREATMENT|          626|        637|        atorvastatin| 0.9693424|
|    TrAP|  PROBLEM|          617|        620|                T2DM|TREATMENT|          643|        653|         gemfibrozil|0.99460286|
|    TeRP|     TEST|          739|        758|Physical examination|  PROBLEM|          796|        810|     dry oral mucosa|0.99775106|
|    TeRP|     TEST|          830|        854|her abdominal exa...|  PROBLEM|          875|        884|          tenderness|0.99272937|
|    TeRP|     TEST|          830|        854|her abdominal exa...|  PROBLEM|          888|        895|            guarding| 0.9840321|
|    TeRP|     TEST|          830|        854|her abdominal exa...|  PROBLEM|          902|        909|            rigidity| 0.9883966|
|    TeRP|     TEST|         1246|       1258|       blood samples|  PROBLEM|         1265|       1274|          hemolyzing| 0.9534202|
|    TeRP|     TEST|         1507|       1517|         her glucose|  PROBLEM|         1553|       1566|      still elevated| 0.9464761|
|    TeRP|  PROBLEM|         1553|       1566|      still elevated|     TEST|         1576|       1592|   serum bicarbonate| 0.9428323|
|    TeRP|  PROBLEM|         1553|       1566|      still elevated|     TEST|         1656|       1661|              lipase| 0.9558198|
|    TeRP|  PROBLEM|         1553|       1566|      still elevated|     TEST|         1670|       1672|                 U/L| 0.9214444|
|    TeRP|     TEST|         1676|       1702|The β-hydroxybuty...|  PROBLEM|         1733|       1740|            elevated| 0.9863963|
|    TrAP|TREATMENT|         1937|       1951|     an insulin drip|  PROBLEM|         1957|       1961|               euDKA| 0.9852455|
|       O|  PROBLEM|         1957|       1961|               euDKA|     TEST|         1991|       2003|       the anion gap|0.94141793|
|       O|  PROBLEM|         1957|       1961|               euDKA|     TEST|         2015|       2027|       triglycerides| 0.9622529|
+--------+---------+-------------+-----------+--------------------+---------+-------------+-----------+--------------------+----------+

Model Information

Model Name: redl_clinical_biobert
Compatibility: Spark NLP for Healthcare 3.0.3+
License: Licensed
Edition: Official
Language: en
Case sensitive: true

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