# -*- coding: utf-8 -*-
"""NER.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/14VcPCWWSAS7tEIolL_I9iA0xulk2gHuN
"""

#!pip install -q transformers

from transformers import pipeline

#!pip install -q gradio

import gradio as gr

# Corrected examples in the nested list format
examples = [
    ["Date of Service: March 15, 2024\nPatient Name: John Doe\nAge: 45\nDiagnosis: Acute bronchitis\n\nChief Complaint:\nPatient presents with a persistent cough, productive of green sputum, for the past week. Reports accompanying symptoms of low-grade fever, malaise, and mild shortness of breath on exertion.\n\nHistory of Present Illness:\nMr. Doe reports the onset of symptoms approximately 10 days ago with a gradual worsening of his cough and overall condition. Denies any recent travel or sick contacts.\n\nAssessment:\nUpon examination, patient displays signs consistent with acute bronchitis, including rhonchi on auscultation, mild tachypnea, and low-grade fever of 100.4°F.\n\nPlan:\n1. Prescribed 7-day course of azithromycin 500mg once daily.\n2. Encouraged increased fluid intake and rest.\n3. Follow-up appointment scheduled in one week for reassessment.\n\nDr. Jane Smith, MD"],
    ["Date of Service: March 16, 2024\nPatient Name: Jane Johnson\nAge: 32\nDiagnosis: Hypertension\n\nChief Complaint:\nPatient presents for routine follow-up of hypertension.\n\nHistory of Present Illness:\nMs. Johnson was diagnosed with hypertension 6 months ago and has been taking lisinopril 10mg daily. Reports occasional headaches and dizziness, denies any chest pain or shortness of breath.\n\nAssessment:\nBlood pressure today is 140/90 mmHg. Otherwise, general examination is unremarkable.\n\nPlan:\n1. Increased lisinopril dosage to 20mg daily.\n2. Advised lifestyle modifications including low-sodium diet and regular exercise.\n3. Follow-up appointment scheduled in 3 months.\n\nDr. Michael Lee, MD"],
    ["Date of Service: March 17, 2024\nPatient Name: Sarah Adams\nAge: 28\nDiagnosis: Urinary Tract Infection (UTI)\n\nChief Complaint:\nPatient complains of burning sensation on urination and increased frequency.\n\nHistory of Present Illness:\nMs. Adams reports symptoms starting 3 days ago. Denies any fever, flank pain, or hematuria.\n\nAssessment:\nUrinalysis reveals presence of leukocytes and nitrites. Mild suprapubic tenderness on examination.\n\nPlan:\n1. Prescribed 3-day course of ciprofloxacin 500mg twice daily.\n2. Advised increased fluid intake and avoidance of irritants like caffeine.\n3. Follow-up in one week for reassessment or sooner if symptoms worsen.\n\nDr. Emily Davis, MD"]
]

# Define multiple NER models
models = {
    "deid_roberta_i2b2": "obi/deid_roberta_i2b2",
    "bert-base-NER": "dslim/bert-base-NER",
    # Add more models as needed
}

# Create a dictionary of model names to their corresponding pipeline instances
ner_pipelines = {name: pipeline("ner", model=model) for name, model in models.items()}

def ner(text, model_name="deid_roberta_i2b2"):
    # Get the NER pipeline based on the selected model name
    ner_pipeline = ner_pipelines[model_name]
    output = ner_pipeline(text, aggregation_strategy="simple")
    return {"text": text, "entities": output}

# Define the Gradio interface
demo = gr.Interface(
    fn=ner,
    inputs=[
        gr.Textbox(placeholder="Enter a sentence here..", lines=10),
        gr.Dropdown(choices=list(models.keys()), label="Select Model"),
    ],
    outputs=gr.HighlightedText(),
    examples=examples,
    title="Named Entity Recognition (NER) Demo",
    description="Select a model and enter a sentence to extract named entities.",
)

# Launch the Gradio interface
demo.launch(share=True)