import streamlit as st import time from annotated_text import annotated_text from io import StringIO from transformers import AutoTokenizer, AutoModelForTokenClassification import os from streamlit_text_annotation import text_annotation 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") with st.sidebar: selected_menu = option_menu("Select Option", ["Upload Document", "Extract Text", "Summarize Document", "Extract Entities","Detected Barriers","Get Answers","Annotation Tool"], 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?"} with st.spinner("Detecting Barriers..."): for barrier,question_text in 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) if result['score'] < 0.3: st.text("Not Found") else: 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: with st.spinner("Finding Answer(s)..."): result = pipeline_qa(question=question_text, context=context) st.subheader('Answer') st.text(result['answer']) elif selected_menu == "Annotation Tool": data1 = { "tokens": [ {"text": "He", "labels": ["Person"]}, {"text": "loves"}, {"text": "his"}, {"text": "dog", "labels": ["Animal", "Pet"]}, ], "labels": [ {"text": "Person"}, {"text": "Action"}, {"text": "Animal"}, ] } st.subheader("Display Mode:") left, right = st.columns(2) with left: st.text("Vertical labels:") text_annotation(data1) with right: st.text("Horizontal labels:") data1["labelOrientation"] = "horizontal" text_annotation(data1) data2 = { "allowEditing": True, "tokens": [ {"text": "He", "labels": ["Pronoun", "Person"]}, {"text": "loves", "labels": ["Action"]}, {"text": "his"}, {"text": "dog", "labels": ["Animal"]}, ], "labels": [ {"text": "Pronoun", "style": { "color": "red", "background-color": "white", "font-size": "8px", "border": "3px dashed red", }}, {"text": "Verb", "style": { "color": "green", "background-color": "white", "font-size": "8px", "font-weight": "900", }}, {"text": "Noun", "style": { "color": "blue", "background-color": "white", "font-size": "8px", }}, {"text": "Person"}, {"text": "Animal"}, ] } st.subheader("Edit Mode:") data = text_annotation(data2) if data: "Returned data:", data