Update app.py

app.py

@@ -1,3 +1,6 @@
+pip install scispacy
+pip install https://s3-us-west-2.amazonaws.com/ai2-s2-scispacy/releases/v0.5.1/en_core_sci_sm-0.5.1.tar.gz
+
 import gradio as gr
 import shap
 import numpy as np
@@ -12,6 +15,7 @@ import csv
 import io
 import base64
 
+
 # Increase CSV field size limit
 csv.field_size_limit(sys.maxsize)
 
@@ -26,12 +30,41 @@ model = AutoModelForSequenceClassification.from_pretrained("TyHamil/ADRv2025").t
 pred = transformers.pipeline("text-classification", model=model, tokenizer=tokenizer, top_k=None, device=device)
 
 # SHAP explainer
-explainer = shap.Explainer(pred)
+#explainer = shap.Explainer(pred)
+import shap
+
+def predict_prob(texts):
+    encoded = tokenizer(texts, return_tensors='pt', padding=True, truncation=True).to(device)
+    with torch.no_grad():
+        outputs = model(**encoded)
+        probs = torch.nn.functional.softmax(outputs.logits, dim=1)
+    return probs.cpu().numpy()
+
+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")
+#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
 
 # SHAP Plotting Function
 def generate_shap_plot(shap_values):
@@ -59,13 +92,13 @@ def adr_predict(x):
         local_plot = "<p>SHAP explanation not available.</p>"
 
     # NER Processing
-    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'
-        }
+     try:
+    htext = scispacy_ner(text_input)
+    except Exception as e:
+    print(f"NER processing failed: {e}")
+    htext = "<p>NER processing not available.</p>"
+
+        
         htext = "<div style='line-height: 1.5; font-family: Poppins;'>"
         prev_end = 0
         res = sorted(res, key=lambda x: x['start'])