Update Demo.py

Demo.py

@@ -1,228 +1,204 @@
-import streamlit as st
-import sparknlp
-import os
-import pandas as pd
-from sparknlp.base import *
-from sparknlp.annotator import *
-from pyspark.ml import Pipeline
-from sparknlp.pretrained import PretrainedPipeline
-from annotated_text import annotated_text
-
-# Page configuration
-st.set_page_config(layout="wide", initial_sidebar_state="auto")
-
-# CSS for styling
-st.markdown("""
-    <style>
-        .main-title {
-            font-size: 36px;
-            color: #4A90E2;
-            font-weight: bold;
-            text-align: center;
-        }
-        .section {
-            background-color: #f9f9f9;
-            padding: 10px;
-            border-radius: 10px;
-            margin-top: 10px;
-        }
-        .section p, .section ul {
-            color: #666666;
-        }
-    </style>
-""", unsafe_allow_html=True)
-
-# Initialize Spark NLP
-@st.cache_resource
-def init_spark():
-    return sparknlp.start()
-
-# Create the NER pipeline
-@st.cache_resource
-def create_pipeline(model, context_dict):
-    documentAssembler = DocumentAssembler() \
-        .setInputCol("text") \
-        .setOutputCol("document")
-
-    sentenceDetector = SentenceDetector() \
-        .setInputCols(["document"]) \
-        .setOutputCol("sentence")
-
-    tokenizer = Tokenizer() \
-        .setInputCols(["sentence"]) \
-        .setOutputCol("token")
-        
-    zero_shot_ner = ZeroShotNerModel.pretrained(model, "en")\
-        .setInputCols(["sentence", "token"])\
-        .setOutputCol("zero_shot_ner")\
-        .setEntityDefinitions(context_dict)
-
-    ner_converter = NerConverter()\
-        .setInputCols(["sentence", "token", "zero_shot_ner"])\
-        .setOutputCol("ner_chunk")\
-
-    pipeline = Pipeline(stages=[documentAssembler, sentenceDetector, tokenizer, zero_shot_ner, ner_converter])
-    return pipeline
-
-# Fit data using the pipeline
-def fit_data(pipeline, data):
-    empty_df = spark.createDataFrame([['']]).toDF('text')
-    pipeline_model = pipeline.fit(empty_df)
-    model = LightPipeline(pipeline_model)
-    result = model.fullAnnotate(data)
-    return result
-
-# Annotate the text with NER results
-def annotate(data):
-    document, chunks, labels = data["Document"], data["NER Chunk"], data["NER Label"]
-    annotated_words = []
-    for chunk, label in zip(chunks, labels):
-        parts = document.split(chunk, 1)
-        if parts[0]:
-            annotated_words.append(parts[0])
-        annotated_words.append((chunk, label))
-        document = parts[1]
-    if document:
-        annotated_words.append(document)
-    annotated_text(*annotated_words)
-
-def df_to_dict(df):
-    context_dict = {}
-    for col in df.columns:
-        values = df[col].dropna().tolist() 
-        if values:
-            context_dict[col] = values
-    return context_dict
-
-# Sidebar content
-model = st.sidebar.selectbox(
-    "Choose the pretrained model",
-    ["zero_shot_ner_roberta"],
-    help="For more info about the models visit: https://sparknlp.org/models"
-)
-
-# Set up the page layout
-st.markdown('<div class="main-title">Zero-Shot Named Entity Recognition (NER)</div>', unsafe_allow_html=True)
-st.markdown('<div class="section"><p>Explore Zero-Shot Named Entity Recognition (NER)—a state-of-the-art technique that detects and classifies named entities in text without needing specific training on annotated datasets. With our interactive interface, you can modify the context by editing the DataFrame to define custom entity types and examples. Then, input your own text or select from predefined examples to see how the model identifies and categorizes entities in real time.</p></div>', unsafe_allow_html=True)
-
-# Reference notebook link in sidebar
-link = """
-<a href="https://github.com/JohnSnowLabs/spark-nlp-workshop/blob/master/tutorials/Certification_Trainings/Healthcare/1.6.ZeroShot_Clinical_NER.ipynb">
-    <img src="https://colab.research.google.com/assets/colab-badge.svg" style="zoom: 1.3" alt="Open In Colab"/>
-</a>
-"""
-st.sidebar.markdown('Reference notebook:')
-st.sidebar.markdown(link, unsafe_allow_html=True)
-
-# Set examples and create DataFrame
-data = {
-    "PROBLEM": [
-        "What is the disease?",
-        "What are the symptoms of the condition?",
-        "What is the patient's diagnosis?",
-        "What kind of disease is he suffering from?",
-        "What specific medical issue does she have?",
-        "What is the main problem the patient is facing?",
-        "What were the reasons for the patient's hospitalization?"
-    ],
-    "DRUG": [
-        "Which medication was prescribed?",
-        "What is the name of the drug used for treatment?",
-        "Which drug is administered for this condition?",
-        "What medication does he take daily?",
-        "What drugs are used to manage his symptoms?",
-        "Which medicine is recommended for this illness?",
-        "What is the prescription for this medical condition?"
-    ],
-    "ADMISSION_DATE": [
-        "When was the patient admitted to the hospital?",
-        "What is the date of the patient's admission?",
-        "On which date did the patient enter the clinic?",
-        "When did the patient check into the hospital?",
-        "What is the admission date for the patient?"
-    ],
-    "PATIENT_AGE": [
-        "How old is the patient?",
-        "What is the patient's age?",
-        "At what age was the patient diagnosed?",
-        "Can you tell me the age of the patient?",
-        "What is the age of the person receiving treatment?"
-    ],
-    "SYMPTOM": [
-        "What symptoms is the patient experiencing?",
-        "What are the signs of the disease?",
-        "Which symptoms did the patient report?",
-        "What were the initial symptoms observed?",
-        "What specific symptoms are present?"
-    ],
-    "TREATMENT": [
-        "What treatment plan was recommended?",
-        "Which therapies are being used?",
-        "What is the current treatment protocol?",
-        "What type of treatment is the patient undergoing?",
-        "What are the options for treating this condition?"
-    ],
-    "DOCTOR": [
-        "Who is the treating physician?",
-        "Which doctor is handling the case?",
-        "What is the name of the attending doctor?",
-        "Who is the specialist for this illness?",
-        "Can you provide the name of the doctor overseeing the treatment?"
-    ]
-}
-
-# Pad shorter lists with empty strings
-max_length = max(len(v) for v in data.values())
-for key in data.keys():
-    while len(data[key]) < max_length:
-        data[key].append(None)
-
-# Create DataFrame and display
-df = pd.DataFrame(data)
-df.index += 1
-st.write("Context DataFrame (Click To Edit)")
-edited_df = st.data_editor(df)
-
-# Example sentences
-examples = [
-    "Dr. Taylor prescribed Lisinopril to a 68-year-old patient with high blood pressure. The patient was admitted to the hospital on April 15, 2024, after experiencing severe hypertension symptoms.",
-    "The 50-year-old male patient reported persistent back pain. The treatment plan includes physical therapy and a medication called Flexeril to alleviate the discomfort.",
-    "The patient was admitted on June 12, 2024, with symptoms of severe abdominal pain. Dr. Kim diagnosed acute appendicitis and recommended immediate surgical intervention.",
-    "A 25-year-old female patient with a history of asthma was treated with Albuterol. Dr. Patel noted that the patient's symptoms of shortness of breath improved significantly after starting the medication.",
-    "The 72-year-old patient underwent surgery for cataracts on May 30, 2024. Dr. Martinez prescribed eye drops and scheduled a follow-up appointment for post-operative care.",
-    "The patient, aged 40, presented with symptoms of chronic fatigue and joint pain. Dr. Nguyen recommended a combination of lifestyle changes and a new drug called Imunorix.",
-    "Dr. Wilson observed that the 34-year-old patient’s condition, diagnosed as lupus, was managed with hydroxychloroquine and a tailored treatment regimen.",
-    "The 56-year-old patient experienced symptoms of severe dehydration due to gastroenteritis. Dr. Rogers provided intravenous fluids and anti-nausea medication as treatment.",
-    "On July 8, 2024, the patient was admitted with acute respiratory distress. Dr. Green prescribed a regimen including corticosteroids and bronchodilators to manage the symptoms.",
-    "The 29-year-old patient, diagnosed with chronic migraines, was treated with a new medication called MigraRelief. The attending physician, Dr. Lewis, also recommended cognitive behavioral therapy."
-]
-
-selected_text = st.selectbox("Select an example", examples)
-custom_input = st.text_input("Try it with your own Sentence!")
-
-text_to_analyze = custom_input if custom_input else selected_text
-context_dict = df_to_dict(edited_df)
-
-# Display example text
-st.subheader('Full Example Text')
-HTML_WRAPPER = """<div class="scroll entities" style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem; margin-bottom: 2.5rem; white-space:pre-wrap">{}</div>"""
-st.markdown(HTML_WRAPPER.format(text_to_analyze), unsafe_allow_html=True)
-
-# Initialize Spark and create pipeline
-spark = init_spark()
-pipeline = create_pipeline(model, context_dict)
-output = fit_data(pipeline, text_to_analyze)
-
-# Display processed output
-st.subheader("Processed Output:")
-results = {
-    'Document': output[0]['document'][0].result,
-    'NER Chunk': [n.result for n in output[0]['ner_chunk']],
-    "NER Label": [n.metadata['entity'] for n in output[0]['ner_chunk']]
-}
-annotate(results)
-
-with st.expander("View DataFrame"):
-    df = pd.DataFrame({'NER Chunk': results['NER Chunk'], 'NER Label': results['NER Label']})
-    df.index += 1
-    st.dataframe(df)
+import streamlit as st
+import sparknlp
+import os
+import pandas as pd
+from sparknlp.base import *
+from sparknlp.annotator import *
+from pyspark.ml import Pipeline
+from sparknlp.pretrained import PretrainedPipeline
+from annotated_text import annotated_text
+
+# Page configuration
+st.set_page_config(layout="wide", initial_sidebar_state="auto")
+
+# CSS for styling
+st.markdown("""
+    <style>
+        .main-title {
+            font-size: 36px;
+            color: #4A90E2;
+            font-weight: bold;
+            text-align: center;
+        }
+        .section {
+            background-color: #f9f9f9;
+            padding: 10px;
+            border-radius: 10px;
+            margin-top: 10px;
+        }
+        .section p, .section ul {
+            color: #666666;
+        }
+    </style>
+""", unsafe_allow_html=True)
+
+# Initialize Spark NLP
+@st.cache_resource
+def init_spark():
+    return sparknlp.start()
+
+# Create the NER pipeline
+@st.cache_resource
+def create_pipeline(model, context_dict):
+    documentAssembler = DocumentAssembler() \
+        .setInputCol("text") \
+        .setOutputCol("document")
+
+    sentenceDetector = SentenceDetector() \
+        .setInputCols(["document"]) \
+        .setOutputCol("sentence")
+
+    tokenizer = Tokenizer() \
+        .setInputCols(["sentence"]) \
+        .setOutputCol("token")
+        
+    zero_shot_ner = ZeroShotNerModel.pretrained(model, "en")\
+        .setInputCols(["sentence", "token"])\
+        .setOutputCol("zero_shot_ner")\
+        .setEntityDefinitions(context_dict)
+
+    ner_converter = NerConverter()\
+        .setInputCols(["sentence", "token", "zero_shot_ner"])\
+        .setOutputCol("ner_chunk")\
+
+    pipeline = Pipeline(stages=[documentAssembler, sentenceDetector, tokenizer, zero_shot_ner, ner_converter])
+    return pipeline
+
+# Fit data using the pipeline
+def fit_data(pipeline, data):
+    empty_df = spark.createDataFrame([['']]).toDF('text')
+    pipeline_model = pipeline.fit(empty_df)
+    model = LightPipeline(pipeline_model)
+    result = model.fullAnnotate(data)
+    return result
+
+# Annotate the text with NER results
+def annotate(data):
+    document, chunks, labels = data["Document"], data["NER Chunk"], data["NER Label"]
+    annotated_words = []
+    for chunk, label in zip(chunks, labels):
+        parts = document.split(chunk, 1)
+        if parts[0]:
+            annotated_words.append(parts[0])
+        annotated_words.append((chunk, label))
+        document = parts[1]
+    if document:
+        annotated_words.append(document)
+    annotated_text(*annotated_words)
+
+def df_to_dict(df):
+    context_dict = {}
+    for col in df.columns:
+        values = df[col].dropna().tolist() 
+        if values:
+            context_dict[col] = values
+    return context_dict
+
+# Sidebar content
+model = st.sidebar.selectbox(
+    "Choose the pretrained model",
+    ["zero_shot_ner_roberta"],
+    help="For more info about the models visit: https://sparknlp.org/models"
+)
+
+# Set up the page layout
+st.markdown('<div class="main-title">Zero-Shot Named Entity Recognition (NER)</div>', unsafe_allow_html=True)
+st.markdown('<div class="section"><p>Explore Zero-Shot Named Entity Recognition (NER)—a state-of-the-art technique that detects and classifies named entities in text without needing specific training on annotated datasets. With our interactive interface, you can modify the context by editing the DataFrame to define custom entity types and examples. Then, input your own text or select from predefined examples to see how the model identifies and categorizes entities in real time.</p></div>', unsafe_allow_html=True)
+
+# Reference notebook link in sidebar
+link = """
+<a href="https://github.com/JohnSnowLabs/spark-nlp-workshop/blob/master/tutorials/Certification_Trainings/Healthcare/1.6.ZeroShot_Clinical_NER.ipynb">
+    <img src="https://colab.research.google.com/assets/colab-badge.svg" style="zoom: 1.3" alt="Open In Colab"/>
+</a>
+"""
+st.sidebar.markdown('Reference notebook:')
+st.sidebar.markdown(link, unsafe_allow_html=True)
+
+# Set examples and create DataFrame
+data = {
+    "PROBLEM": [
+        "What is the disease?",
+        "What are the symptoms of the condition?",
+        "What is the patient's diagnosis?",
+        "What kind of disease is he suffering from?",
+        "What specific medical issue does she have?",
+        "What is the main problem the patient is facing?",
+        "What were the reasons for the patient's hospitalization?"
+    ],
+    "DRUG": [
+        "Which medication was prescribed?",
+        "What is the name of the drug used for treatment?",
+        "Which drug is administered for this condition?",
+        "What medication does he take daily?",
+        "What drugs are used to manage his symptoms?",
+        "Which medicine is recommended for this illness?",
+        "What is the prescription for this medical condition?"
+    ],
+    "SYMPTOM": [
+        "What symptoms is the patient experiencing?",
+        "What are the signs of the disease?",
+        "Which symptoms did the patient report?",
+        "What were the initial symptoms observed?",
+        "What specific symptoms are present?"
+    ]
+}
+
+# Pad shorter lists with empty strings
+max_length = max(len(v) for v in data.values())
+for key in data.keys():
+    while len(data[key]) < max_length:
+        data[key].append(None)
+
+# Create DataFrame 
+df = pd.DataFrame(data)
+df.index += 1
+
+# Add new key with all values as None
+new_key = st.text_input("Add Lable:")
+if new_key:
+    data = {new_key.upper(): [None] * max_length, **data}
+    df = pd.DataFrame(data)
+    df.index += 1
+    st.success(f"Added '{new_key}' to Context DataFrame")
+
+st.write("Context DataFrame (Click To Edit)")
+edited_df = st.data_editor(df)
+
+# Example sentences
+examples = [
+    "The patient is experiencing severe headache and nausea. The doctor has prescribed Ibuprofen to alleviate the symptoms. The patient has been complaining about these symptoms for the last three days.",
+    "The patient's main problem is chronic back pain. This issue has been affecting their daily activities significantly. The doctor recommended a series of physical therapy sessions to address the problem.",
+    "After the diagnosis of diabetes, the patient was given Metformin as a part of their treatment plan. The medication is intended to help manage blood sugar levels effectively.",
+    "The symptoms reported by the patient include persistent cough and shortness of breath. The doctor has advised some tests to identify the underlying cause of these symptoms.",
+    "The patient has been prescribed Prednisone to manage their severe inflammation. This medication is part of the treatment plan for their chronic condition."
+]
+
+selected_text = st.selectbox("Select an example", examples)
+custom_input = st.text_input("Try it with your own Sentence!")
+
+text_to_analyze = custom_input if custom_input else selected_text
+context_dict = df_to_dict(edited_df)
+
+# Display example text
+st.subheader('Full Example Text')
+HTML_WRAPPER = """<div class="scroll entities" style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem; margin-bottom: 2.5rem; white-space:pre-wrap">{}</div>"""
+st.markdown(HTML_WRAPPER.format(text_to_analyze), unsafe_allow_html=True)
+
+# Initialize Spark and create pipeline
+spark = init_spark()
+pipeline = create_pipeline(model, context_dict)
+output = fit_data(pipeline, text_to_analyze)
+
+# Display processed output
+st.subheader("Processed Output:")
+results = {
+    'Document': output[0]['document'][0].result,
+    'NER Chunk': [n.result for n in output[0]['ner_chunk']],
+    "NER Label": [n.metadata['entity'] for n in output[0]['ner_chunk']]
+}
+annotate(results)
+
+with st.expander("View DataFrame"):
+    df = pd.DataFrame({'NER Chunk': results['NER Chunk'], 'NER Label': results['NER Label']})
+    df.index += 1
+    st.dataframe(df)