app.py

import streamlit as st
import time
from annotated_text import annotated_text
from io import StringIO
from transformers import AutoTokenizer, AutoModelForTokenClassification
import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'

import plotly.express as px
from streamlit_option_menu import option_menu

st. set_page_config(layout="wide")

from transformers import pipeline
import pandas as pd

@st.cache(allow_output_mutation = True)
def init_text_summarization_model():
    MODEL = 'facebook/bart-large-cnn'
    pipe = pipeline("summarization", model=MODEL)
    return pipe

@st.cache(allow_output_mutation = True)
def init_zsl_topic_classification():
    MODEL = 'facebook/bart-large-mnli'
    pipe = pipeline("zero-shot-classification", model=MODEL)
    template = "This text is about {}."
    return pipe, template

@st.cache(allow_output_mutation = True)
def init_zsl_topic_classification():
    MODEL = 'facebook/bart-large-mnli'
    pipe = pipeline("zero-shot-classification", model=MODEL)
    template = "This text is about {}."
    return pipe, template

@st.cache(allow_output_mutation = True)
def init_ner_pipeline():
    tokenizer = AutoTokenizer.from_pretrained("d4data/biomedical-ner-all")
    model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all")
    pipe = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple") # pass device=0 if using gpu
    return pipe

@st.cache(allow_output_mutation = True)
def init_qa_pipeline():
    question_answerer_pipe = pipeline("question-answering", model='deepset/roberta-base-squad2')
    return question_answerer_pipe

def get_formatted_text_for_annotation(output):
    colour_map = {'Coreference': '#29D93B',
    'Severity':'#FCF3CF',
 'Sex': '#E9F7EF',
 'Sign_symptom': '#EAF2F8',
 'Detailed_description': '#078E8B',
 'Date': '#F5EEF8',
 'History': '#FDEDEC',
 'Medication': '#F4F6F6',
 'Therapeutic_procedure': '#A3E4D7',
 'Age': '#85C1E9',
 'Subject': '#D7BDE2',
 'Biological_structure': '#AF7AC5',
 'Activity': '#B2BABB',
 'Lab_value': '#E6B0AA',
 'Family_history': '#2471A3',
 'Diagnostic_procedure': '#CCD1D1',
 'Other_event': '#239B56',
 'Occupation': '#B3B6B7'}
    
    annotated_texts = []
    next_index = 0
    for entity in output:
        if entity['start'] == next_index:
    #         print("found entity")
            extracted_text = text[entity['start']:entity['end']]
    #         print("annotated",annotated_text)
            annotated_texts.append((extracted_text ,entity['entity_group'],colour_map[entity['entity_group']]))
        else:
            unannotated_text = text[next_index:entity['start']-1]
            annotated_texts.append(unannotated_text)
            extracted_text = text[entity['start']:entity['end']]
            annotated_texts.append((extracted_text ,entity['entity_group'],colour_map[entity['entity_group']]))
            next_index =entity['end'] +1
    
    if next_index < len(text):
        annotated_texts.append(text[next_index-1:len(text)-1])
        
    return tuple(annotated_texts)
    
# Model initialization    
pipeline_summarization = init_text_summarization_model()
pipeline_zsl, template = init_zsl_topic_classification()
pipeline_ner =init_ner_pipeline()
pipeline_qa = init_qa_pipeline()

st.header("Intelligent Document Automation")

def get_paragraphs_for_summaries():
    paras =[]
    paras.append("""Stephen is a 53 year old gentleman who does general duties police work. He is married and has an
18 year old daughter at home. He is right hand dominant. Cigarettes nil, alcohol rare, allergies nil.
DVTe nil. Past medical history hyperlipidemia and reflux testicular cancer in 2000 and right knee
reconstruction in 1987. Medications include Nexium and Crestor.""")
    paras.append("""History presenting complaint: Right knee and right elbow injuries. On 12 January 2020, while at work
he was trying to apprehend a stolen vehicle. The deployed some road spikes onto the road. The
stolen vehicle went over the road spikes. He was attempting to retrieve the spikes to prevent
damage to civilian and police vehicles and while he was doing that, he was hit by a police vehicle
coming from behind. The police vehicle was approximately doing 50 km/hr. The headlight of the
police car stuck him on the lateral aspect of the right knee. He jumped in the air and flicked in the air
with his right elbow also hitting the police car before he flung over a barrier. He was able to mobilise
afterwards when the adrenaline was taking effect. After he cooled down that night he developed
significant pain in both the elbow and the knee. The elbow seemed to have settled with time but he
has got slight discomfort on the lateral epicondyle of the right elbow but otherwise no instability,
clicking, locking or catching of the elbow.""")
    paras.append("""His knee has medial sided pain. It is an annoying type of pain that is present at rear and with activity.
It feels like it is getting slightly better but his main problem is that he cannot still fully flex his knee as
he used to do before the accident. He has noticed some clicking ad a little bit of swelling in the knee.
He has had no instability but he is still a little bit stiff when he first gets up. He has been treated with
4 sessions of physiotherapy. He has had two operations on this knee before when it was
reconstructed which was an open procedure and an anthroscopy later to shave off the cartilage. He
never got back to playing rugby after his reconstruction but he was able to do martial arts and all his
policing duties without any problem.""")
    paras.append("""He is mildly overweight, normal gait and no effusion in the knee and a range of motion of 0 to 110
degrees which cannot be improved passively. His right elbow has full range of motion and is stable.
He has had an Xray on right elbow showing no fracture and an xray of knee showed early medial
compartment osteoarthritis. He has got new clicking in his knee and has been advised to get an MRI
scan.""")
    paras.append("""As far as his right elbow is concerned, he should reduce weight lifting activity and any repetitive right
upper limb activity that causes any discomfort. His right knee requires an MRI scan to look for any
loose bodies and I will see him with the results of the scan. In the meantime, he can continue with
his general policing duties which he has anyways been doing since the accident""")
    return paras
    
def get_paragraphs_for_entities():
    paras =[]
    paras.append("""18 February2020
Dr Christine Fowler""")
    paras.append("""Dear Christine,
Re: Stephen Parrot""")
    paras.append("""Stephen is a 53 year old gentleman who does general duties police work. He is married and has an
18 year old daughter at home. He is right hand dominant. Cigarettes nil, alcohol rare, allergies nil.
DVTe nil. Past medical history hyperlipidemia and reflux testicular cancer in 2000 and right knee
reconstruction in 1987. Medications include Nexium and Crestor.""")
    paras.append("""History presenting complaint: Right knee and right elbow injuries. On 12 January 2020, while at work
he was trying to apprehend a stolen vehicle. The deployed some road spikes onto the road. The
stolen vehicle went over the road spikes. He was attempting to retrieve the spikes to prevent
damage to civilian and police vehicles and while he was doing that, he was hit by a police vehicle
coming from behind. The police vehicle was approximately doing 50 km/hr. The headlight of the
police car stuck him on the lateral aspect of the right knee. He jumped in the air and flicked in the air
with his right elbow also hitting the police car before he flung over a barrier. He was able to mobilise
afterwards when the adrenaline was taking effect. After he cooled down that night he developed
significant pain in both the elbow and the knee. The elbow seemed to have settled with time but he
has got slight discomfort on the lateral epicondyle of the right elbow but otherwise no instability,
clicking, locking or catching of the elbow.""")
    paras.append("""His knee has medial sided pain. It is an annoying type of pain that is present at rear and with activity.
It feels like it is getting slightly better but his main problem is that he cannot still fully flex his knee as
he used to do before the accident. He has noticed some clicking ad a little bit of swelling in the knee.
He has had no instability but he is still a little bit stiff when he first gets up. He has been treated with
4 sessions of physiotherapy. He has had two operations on this knee before when it was
reconstructed which was an open procedure and an anthroscopy later to shave off the cartilage. He
never got back to playing rugby after his reconstruction but he was able to do martial arts and all his
policing duties without any problem.""")
    paras.append("""He is mildly overweight, normal gait and no effusion in the knee and a range of motion of 0 to 110
degrees which cannot be improved passively. His right elbow has full range of motion and is stable.
He has had an Xray on right elbow showing no fracture and an xray of knee showed early medial
compartment osteoarthritis. He has got new clicking in his knee and has been advised to get an MRI
scan.""")
    paras.append("""As far as his right elbow is concerned, he should reduce weight lifting activity and any repetitive right
upper limb activity that causes any discomfort. His right knee requires an MRI scan to look for any
loose bodies and I will see him with the results of the scan. In the meantime, he can continue with
his general policing duties which he has anyways been doing since the accident""")
    paras.append("""Kind regards,
Dr Jason Sanders""")
    return paras
    
def get_text_from_ocr_engine():

    return """18 February2020
Dr Christine Fowler
Dear Christine,
Re: Stephen Parrot
Stephen is a 53 year old gentleman who does general duties police work. He is married and has an
18 year old daughter at home. He is right hand dominant. Cigarettes nil, alcohol rare, allergies nil.
DVTe nil. Past medical history hyperlipidemia and reflux testicular cancer in 2000 and right knee
reconstruction in 1987. Medications include Nexium and Crestor.
History presenting complaint: Right knee and right elbow injuries. On 12 January 2020, while at work
he was trying to apprehend a stolen vehicle. The deployed some road spikes onto the road. The
stolen vehicle went over the road spikes. He was attempting to retrieve the spikes to prevent
damage to civilian and police vehicles and while he was doing that, he was hit by a police vehicle
coming from behind. The police vehicle was approximately doing 50 km/hr. The headlight of the
police car stuck him on the lateral aspect of the right knee. He jumped in the air and flicked in the air
with his right elbow also hitting the police car before he flung over a barrier. He was able to mobilise
afterwards when the adrenaline was taking effect. After he cooled down that night he developed
significant pain in both the elbow and the knee. The elbow seemed to have settled with time but he
has got slight discomfort on the lateral epicondyle of the right elbow but otherwise no instability,
clicking, locking or catching of the elbow.
His knee has medial sided pain. It is an annoying type of pain that is present at rear and with activity.
It feels like it is getting slightly better but his main problem is that he cannot still fully flex his knee as
he used to do before the accident. He has noticed some clicking ad a little bit of swelling in the knee.
He has had no instability but he is still a little bit stiff when he first gets up. He has been treated with
4 sessions of physiotherapy. He has had two operations on this knee before when it was
reconstructed which was an open procedure and an anthroscopy later to shave off the cartilage. He
never got back to playing rugby after his reconstruction but he was able to do martial arts and all his
policing duties without any problem.
He is mildly overweight, normal gait and no effusion in the knee and a range of motion of 0 to 110
degrees which cannot be improved passively. His right elbow has full range of motion and is stable.
He has had an Xray on right elbow showing no fracture and an xray of knee showed early medial
compartment osteoarthritis. He has got new clicking in his knee and has been advised to get an MRI
scan.
As far as his right elbow is concerned, he should reduce weight lifting activity and any repetitive right
upper limb activity that causes any discomfort. His right knee requires an MRI scan to look for any
loose bodies and I will see him with the results of the scan. In the meantime, he can continue with
his general policing duties which he has anyways been doing since the accident.
Kind regards,
Dr Jason Sanders"""

with st.sidebar:
    selected_menu = option_menu("Select Option", 
    ["Upload Document", "Extract Text", "Summarize Document", "Extract Entities","Detected Barriers","Get Answers"], 
        menu_icon="cast", default_index=0)
    

if selected_menu == "Upload Document":
    uploaded_file = st.file_uploader("Choose a file")        
    if uploaded_file is not None:
        ocr_text  = get_text_from_ocr_engine()
        st.write("Upload Successful")
        
elif selected_menu == "Extract Text":
    with st.spinner("Extracting Text..."):
        time.sleep(6)
        st.write(get_text_from_ocr_engine())
        
elif selected_menu == "Summarize Document":
    paragraphs= get_paragraphs_for_summaries()
    
    with st.spinner("Finding Topics..."):
        tags_found = ["Injury Details", "Past Medical Conditions", "Injury Management Plan", "GP Correspondence"]
        time.sleep(5)
        st.write("This document is about:")
        st.markdown(";".join(["#" + tag + " "  for tag in tags_found]) + "**")
        st.markdown("""---""")
        
    with st.spinner("Summarizing Document..."):
        
        
        for text in paragraphs:
            summary_text = pipeline_summarization(text, max_length=130, min_length=30, do_sample=False)
            # Show output
            st.write(summary_text[0]['summary_text'])
            st.markdown("""---""")
     
        
elif selected_menu == "Extract Entities":
    paragraphs= get_paragraphs_for_entities()
    
    with st.spinner("Extracting Entities..."):
        for text in paragraphs:
            output = pipeline_ner (text)
            entities_text =get_formatted_text_for_annotation(output)
            annotated_text(*entities_text)
            st.markdown("""---""")
            
elif selected_menu == "Detected Barriers":
    st.subheader('Barriers Detected')
    barriers_to_detect = {"Chronic Pain":"Is the patint experiencing chronic pain?",
                          "Mental Health Issues":"Does he have any mental issues?",
                          "Prior History":"What is prior medical history?",
                          "Smoking":"Does he smoke?",
                          "Drinking":"Does he drink?",
                          "Comorbidities":"Does he have any comorbidities?"}
                          
    for barrier,question_text from barriers_to_detect.items():
    
        context = get_text_from_ocr_engine()
        if question_text:
            result = pipeline_qa(question=question_text,  context=context)
            st.subheader(barrier)
            st.text(result['answer']) 

elif selected_menu == "Get Answers":
    st.subheader('Question')
    question_text = st.text_input("Type your question")
    context = get_text_from_ocr_engine()
    if question_text:
        result = pipeline_qa(question=question_text,  context=context)
        st.subheader('Answer')
        st.text(result['answer'])