Update app.py

app.py

@@ -1,5 +1,3 @@
-
-
 import gradio as gr
 import shap
 import numpy as np
@@ -14,7 +12,6 @@ import csv
 import io
 import base64
 
-
 # Increase CSV field size limit
 csv.field_size_limit(sys.maxsize)
 
@@ -41,29 +38,10 @@ def predict_prob(texts):
 
 explainer = shap.Explainer(predict_prob, tokenizer)
 
-
 # NER pipeline
-#ner_tokenizer = AutoTokenizer.from_pretrained("d4data/biomedical-ner-all")
-#ner_model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all")
-#ner_pipe = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, aggregation_strategy="simple")
-
-import spacy
-import scispacy
-nlp = spacy.load("en_core_sci_sm")  # Use small SciSpacy model
-
-def scispacy_ner(text_input):
-    doc = nlp(text_input)
-    highlighted = text_input
-    offset = 0
-    for ent in doc.ents:
-        start = ent.start_char + offset
-        end = ent.end_char + offset
-        label = ent.label_
-        color = "#a3e635" if "DISEASE" in label else "#1e3a8a"
-        replacement = f"<mark style='background-color:{color}; border-radius: 4px;'>{ent.text} ({label})</mark>"
-        highlighted = highlighted[:start] + replacement + highlighted[end:]
-        offset += len(replacement) - (end - start)
-    return highlighted
+ner_tokenizer = AutoTokenizer.from_pretrained("d4data/biomedical-ner-all")
+ner_model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all")
+ner_pipe = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, aggregation_strategy="simple")
 
 # SHAP Plotting Function
 def generate_shap_plot(shap_values):
@@ -91,13 +69,13 @@ def adr_predict(x):
         local_plot = "<p>SHAP explanation not available.</p>"
 
     # NER Processing
-     try:
-    htext = scispacy_ner(text_input)
-    except Exception as e:
-    print(f"NER processing failed: {e}")
-    htext = "<p>NER processing not available.</p>"
-
-        
+    try:
+        res = ner_pipe(text_input)
+        entity_colors = {
+            'Severity': '#a3e635', 'Sign_symptom': '#1e3a8a', 'Medication': '#c0c0c0',
+            'Age': '#a3e635', 'Sex': '#a3e635', 'Diagnostic_procedure': '#c0c0c0',
+            'Biological_structure': '#c0c0c0'
+        }
         htext = "<div style='line-height: 1.5; font-family: Poppins;'>"
         prev_end = 0
         res = sorted(res, key=lambda x: x['start'])
@@ -302,4 +280,4 @@ with gr.Blocks(title="AwareRx. Painless Input. Painless Life.") as demo:
             gr.HTML(generate_alerts())
 
     gr.Markdown("Data privacy is our priority. All information is securely stored following HIPAA guidelines.")
-demo.launch()
+demo.launch()