# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] id="pFGiT_o1FJNH"
# ![JohnSnowLabs](https://nlp.johnsnowlabs.com/assets/images/logo.png)

# + [markdown] id="Df3b6PnyFz-A"
# [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/JohnSnowLabs/spark-nlp-workshop/blob/master/tutorials/Certification_Trainings/Healthcare/10.3.Clinical_RE_SparkNLP_Paper_Reproduce.ipynb)

# + [markdown] id="OdkGRLgcFN1o"
# ## 10.3 Relation Extraction Spark NLP Paper Reproduce
#
# **Note: You'll need trial license to run this notebook. To get trial license, please sign up using this link: https://www.johnsnowlabs.com/install/**

# + [markdown] id="4wdK8WtAF8ES"
# # Colab Setup

# + colab={"base_uri": "https://localhost:8080/", "height": 73, "resources": {"http://localhost:8080/nbextensions/google.colab/files.js": {"data": "<KEY>", "headers": [["content-type", "application/javascript"]], "ok": true, "status": 200, "status_text": ""}}} id="61uNqFkvF-UU" outputId="2f5b836a-3e65-4aff-ae8c-0c853993e52c"
import json

from google.colab import files

license_keys = files.upload()

with open(list(license_keys.keys())[0]) as f:
    license_keys = json.load(f)

# + id="Ot8dIi00DasF"
# Defining license key-value pairs as local variables
locals().update(license_keys)

# Adding license key-value pairs to environment variables
import os
os.environ.update(license_keys)

# + id="_SQ6KcNjDebN"
# Installing pyspark and spark-nlp
# ! pip install --upgrade -q pyspark==3.1.2 spark-nlp==$PUBLIC_VERSION

# Installing Spark NLP Healthcare
# ! pip install --upgrade -q spark-nlp-jsl==$JSL_VERSION  --extra-index-url https://pypi.johnsnowlabs.com/$SECRET

# Installing Spark NLP Display Library for visualization
# ! pip install -q spark-nlp-display

# + id="MNfBuiuJDvL9"
# if you want to start the session with custom params as in start function above
from pyspark.sql import SparkSession

def start(SECRET):
    builder = SparkSession.builder \
        .appName("Spark NLP Licensed") \
        .master("local[*]") \
        .config("spark.driver.memory", "16G") \
        .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer") \
        .config("spark.kryoserializer.buffer.max", "2000M") \
        .config("spark.jars.packages", "com.johnsnowlabs.nlp:spark-nlp_2.12:"+PUBLIC_VERSION) \
        .config("spark.jars", "https://pypi.johnsnowlabs.com/"+SECRET+"/spark-nlp-jsl-"+JSL_VERSION+".jar")
      
    return builder.getOrCreate()

#spark = start(SECRET)


# + colab={"base_uri": "https://localhost:8080/"} id="RpdYuemp0qwP" outputId="cead207c-ba29-4e01-a539-2b428323c40e"
import json
import os
from pyspark.ml import Pipeline,PipelineModel
from pyspark.sql import SparkSession

from sparknlp.annotator import *
from sparknlp_jsl.annotator import *
from sparknlp.base import *
import sparknlp_jsl
import sparknlp
import pandas as pd

import warnings
warnings.filterwarnings('ignore')

params = {"spark.driver.memory":"16G", 
          "spark.kryoserializer.buffer.max":"2000M", 
          "spark.driver.maxResultSize":"2000M"} 

spark = sparknlp_jsl.start(license_keys['SECRET'],params=params)

print ("Spark NLP Version :", sparknlp.version())
print ("Spark NLP_JSL Version :", sparknlp_jsl.version())

# + [markdown] id="93TRFUKO7Q2C"
# # Prepare Data for training
#
# **Note: Due to copyrights, we are not providing the raw data. You can download the data from official websites, and process them to create the following dataframe for training**

# + [markdown] id="kXvvoDca4cha"
# ## Dataset: i2b2 2010 relation challenge:
# https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3168320/
#
# **Details**
#
# The set of relations defined in the 2010 i2b2 relation challenge
#
# TrIP: A certain treatment has improved or cured a medical problem (eg, ‘infection resolved with antibiotic course’)
#
# TrWP: A patient's medical problem has deteriorated or worsened because of or in spite of a treatment being administered (eg, ‘the tumor was growing despite the drain’)
#
# TrCP: A treatment caused a medical problem (eg, ‘penicillin causes a rash’)
#
# TrAP: A treatment administered for a medical problem (eg, ‘Dexamphetamine for narcolepsy’)
#
# TrNAP: The administration of a treatment was avoided because of a medical problem (eg, ‘Ralafen which is contra-indicated because of ulcers’)
#
# TeRP: A test has revealed some medical problem (eg, ‘an echocardiogram revealed a pericardial effusion’)
#
# TeCP: A test was performed to investigate a medical problem (eg, ‘chest x-ray done to rule out pneumonia’)
#
# PIP: Two problems are related to each other (eg, ‘Azotemia presumed secondary to sepsis’)

# + colab={"base_uri": "https://localhost:8080/", "height": 583} id="BXeSpRYp2qvd" outputId="23a32e95-b52b-46dd-a6eb-e7aac4dc2c38"
training_data = pd.read_csv('i2b2_relations.csv')
important_columns = ['dataset', 'sentence', 'chunk1', 'firstCharEnt1', 'lastCharEnt1', 'label1',
                     'chunk2', 'firstCharEnt2', 'lastCharEnt2', 'label2', 'rel']
training_data = training_data[important_columns]
training_data.head()

# + [markdown] id="XZIaEGO-5haV"
# Column Description:
# - **dataset**: to denote rows for training and testing - since we are using official split, this would be fixed.
# - **sentence**: Sentence(s)/text phrase(s) containing both entities in it.
# - **chunk1**: First entity chunk
# - **firstCharEnt1**: Starting index of the first entity
# - **lastCharEnt1**: Ending index of of the first entity
# - **label1**: Entity type of the first chunk
# - **chunk2**: Second entity chunk
# - **firstCharEnt2**: Starting index of the second entity
# - **lastCharEnt2**: Ending index of of the second entity
# - **label2**: Entity type of the second chunk
# - **rel**: Relation type/class

# + [markdown] id="GyWoJM6t-NSf"
# Check train/test distribution

# + colab={"base_uri": "https://localhost:8080/"} id="QxZ4vIYk3BBq" outputId="f738708f-bba8-4a1d-d777-69eb52ca4284"
training_data['dataset'].value_counts()

# + [markdown] id="tuQYmiCW-Rby"
# Check class distribution

# + colab={"base_uri": "https://localhost:8080/"} id="9Jhhj08L3D48" outputId="f612dce4-ac2a-410f-e01e-d536a938cdaf"
training_data['rel'].value_counts()

# + [markdown] id="l0kUmvzI-U3S"
# Convert to pyspark dataframe for training

# + [markdown] id="GnH2rrNa7eXA"
# # Traing RE model

# + [markdown] id="LUTnOgVw8w72"
# Generate a TF graph

# + colab={"base_uri": "https://localhost:8080/"} id="xsLFCVjBPMB7" outputId="f50d4e7e-76da-4a49-f7d3-4f99b9dffffa"
# if you need to custoize the DL arcitecture (more layers, more features etc.)

from sparknlp_jsl.training import tf_graph

# %tensorflow_version 1.x

n_classes = training_data['rel'].nunique()
print ('Classes: ', n_classes)

tf_graph.build("relation_extraction", build_params={"input_dim": 6000, "output_dim": n_classes, 'batch_norm':1, "hidden_layers": [300, 200], "hidden_act": "relu", 'hidden_act_l2':1}, model_location=".", model_filename="re_with_BN.pb")


# + [markdown] id="tupioJtD9JCU"
# Split the data and convert to a pyspark dataframe for traininng

# + colab={"base_uri": "https://localhost:8080/"} id="4lVnDHfe-Ikx" outputId="31a9b12c-13e8-4f0b-d98a-154d5e923705"
training_data = spark.createDataFrame(training_data)
training_data.show()

# + id="T1kpBYST9Im3"
#Annotation structure
import pyspark.sql.types as T
from pyspark.sql import functions as F

annotationType = T.StructType([
            T.StructField('annotatorType', T.StringType(), False),
            T.StructField('begin', T.IntegerType(), False),
            T.StructField('end', T.IntegerType(), False),
            T.StructField('result', T.StringType(), False),
            T.StructField('metadata', T.MapType(T.StringType(), T.StringType()), False),
            T.StructField('embeddings', T.ArrayType(T.FloatType()), False)
        ])

#UDF function to convert train data to names entitities

@F.udf(T.ArrayType(annotationType))
def createTrainAnnotations(begin1, end1, begin2, end2, chunk1, chunk2, label1, label2):
    
    entity1 = sparknlp.annotation.Annotation("chunk", begin1, end1, chunk1, {'entity': label1.upper(), 'sentence': '0'}, [])
    entity2 = sparknlp.annotation.Annotation("chunk", begin2, end2, chunk2, {'entity': label2.upper(), 'sentence': '0'}, [])    
        
    entity1.annotatorType = "chunk"
    entity2.annotatorType = "chunk"

    return [entity1, entity2]    

#list of valid relations
#rels = ["TrIP", "TrAP", "TeCP", "TrNAP", "TrCP", "PIP", "TrWP", "TeRP"]

#a query to select list of valid relations
#valid_rel_query = "(" + " OR ".join(["rel = '{}'".format(rel) for rel in rels]) + ")"

training_data = training_data\
  .withColumn("begin1i", F.expr("cast(firstCharEnt1 AS Int)"))\
  .withColumn("end1i", F.expr("cast(lastCharEnt1 AS Int)"))\
  .withColumn("begin2i", F.expr("cast(firstCharEnt2 AS Int)"))\
  .withColumn("end2i", F.expr("cast(lastCharEnt2 AS Int)"))\
  .where("begin1i IS NOT NULL")\
  .where("end1i IS NOT NULL")\
  .where("begin2i IS NOT NULL")\
  .where("end2i IS NOT NULL")\
  .withColumn(
      "train_ner_chunks", 
      createTrainAnnotations(
          "begin1i", "end1i", "begin2i", "end2i", "chunk1", "chunk2", "label1", "label2"
      ).alias("train_ner_chunks", metadata={'annotatorType': "chunk"}))
    

train_data = training_data.where("dataset='train'")
test_data = training_data.where("dataset='test'")


# + [markdown] id="Z4vcJoc6-iOd"
# Define Training Pipeline

# + colab={"base_uri": "https://localhost:8080/"} id="j8CInpWY9Ip7" outputId="9b5034b6-1522-420d-c54d-db8b49f5b851"
documenter = DocumentAssembler()\
    .setInputCol("sentence")\
    .setOutputCol("sentences")

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

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

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

reApproach = RelationExtractionApproach()\
    .setInputCols(["embeddings", "pos_tags", "train_ner_chunks", "dependencies"])\
    .setOutputCol("relations")\
    .setLabelColumn("rel")\
    .setEpochsNumber(50)\
    .setBatchSize(200)\
    .setDropout(0.5)\
    .setLearningRate(0.001)\
    .setModelFile("./re_with_BN.pb")\
    .setFixImbalance(True)\
    .setFromEntity("begin1i", "end1i", "label1")\
    .setToEntity("begin2i", "end2i", "label2")\
    .setOutputLogsPath('/content')

train_pipeline = Pipeline(stages=[
    documenter, 
    tokenizer, 
    words_embedder, 
    pos_tagger, 
    dependency_parser, 
    reApproach
])

# + [markdown] id="pqUWwdnx-8Nd"
# Start Training

# + id="0p_aFe0A-gND"
re_model = train_pipeline.fit(train_data)

# + [markdown] id="-sAThoyr_FqC"
# # Evaluate the model

# + id="tFio_UG5-gT_"
predictions = re_model.transform(train_data).toPandas()
predictions['relations'] = predictions['relations'].apply(lambda x : x[0].result if len(x)>= 1 else 'O')

# + colab={"base_uri": "https://localhost:8080/"} id="6K7fx0a__io4" outputId="36abfbea-15df-44ea-8543-3f8dbe70af84"
preds = predictions['relations'].astype(str).values
labels = predictions['rel'].astype(str).values

from sklearn.metrics import classification_report
print(classification_report(labels, preds))

# + [markdown] id="qOZk2h0QRGW2"
# # Save the trained model

# + id="8r4vHMmJRJHO"
re_model.stages[-1].save('re_model')

# + [markdown] id="fj0uX0vEQpBn"
# # Run model on sample data

# + colab={"base_uri": "https://localhost:8080/"} id="Xvx1qi-0YJCe" outputId="3613eab5-04da-4b2f-b5d2-602b9f3e3f54"
documenter = DocumentAssembler()\
    .setInputCol("text")\
    .setOutputCol("document")

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

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

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

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

ner_tagger = MedicalNerModel()\
    .pretrained("ner_clinical", "en", "clinical/models")\
    .setInputCols("sentences", "tokens", "embeddings")\
    .setOutputCol("ner_tags") 

ner_chunker = NerConverterInternal()\
    .setInputCols(["sentences", "tokens", "ner_tags"])\
    .setOutputCol("ner_chunks")

dependency_parser = DependencyParserModel()\
    .pretrained("dependency_conllu", "en")\
    .setInputCols(["sentences", "pos_tags", "tokens"])\
    .setOutputCol("dependencies")

reModel = RelationExtractionModel()\
    .load("re_model")\
    .setInputCols(["embeddings", "pos_tags", "ner_chunks", "dependencies"])\
    .setOutputCol("relations")\
    .setMaxSyntacticDistance(4)\
    .setRelationPairs(["problem-test", "problem-treatment"])

pipeline = Pipeline(stages=[
    documenter,
    sentencer,
    tokenizer, 
    words_embedder, 
    pos_tagger, 
    ner_tagger,
    ner_chunker,
    dependency_parser,
    reModel
])

empty_data = spark.createDataFrame([[""]]).toDF("text")

model = pipeline.fit(empty_data)


# + colab={"base_uri": "https://localhost:8080/", "height": 260} id="gwjzZQpoZaNm" outputId="184c1813-f177-4c01-9555-060b1b31a881"
import pandas as pd

def get_relations_df (results, col='relations'):
  rel_pairs=[]
  for rel in results[0][col]:
      rel_pairs.append((
          rel.result, 
          rel.metadata['entity1'], 
          rel.metadata['entity1_begin'],
          rel.metadata['entity1_end'],
          rel.metadata['chunk1'], 
          rel.metadata['entity2'],
          rel.metadata['entity2_begin'],
          rel.metadata['entity2_end'],
          rel.metadata['chunk2'], 
          rel.metadata['confidence']
      ))

  rel_df = pd.DataFrame(rel_pairs, columns=['relation','entity1','entity1_begin','entity1_end','chunk1','entity2','entity2_begin','entity2_end','chunk2', 'confidence'])

  return rel_df

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 .
"""

lmodel = LightPipeline(model)
annotations = lmodel.fullAnnotate(text)

rel_df = get_relations_df (annotations)

rel_df[rel_df.relation!="O"]


# + colab={"base_uri": "https://localhost:8080/", "height": 1000} id="za6Pqr-2Zk3T" outputId="4ebc166f-923f-4d08-e88e-2a96e34d57c3"
from sparknlp_display import RelationExtractionVisualizer


vis = RelationExtractionVisualizer()
vis.display(annotations[0], 'relations', show_relations=True) # default show_relations: True


StarcoderdataJupyterScriptsDedupFiltered