Update app.py

app.py

@@ -14,42 +14,34 @@ HF_TOKEN = os.getenv("hf_token")
 csv.field_size_limit(sys.maxsize)
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-# ==========================================
-# 1. Load Models and Tokenizers
-# ==========================================
-# Classification Model
+# Load models and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("willwim/adr_SJM_Notebook-Copy_for_T3", token=HF_TOKEN)
 model = AutoModelForSequenceClassification.from_pretrained("willwim/adr_SJM_Notebook-Copy_for_T3", token=HF_TOKEN).to(device)
 
-# Build a pipeline object for predictions (used for both label and SHAP)
+# Build a pipeline object for predictions
 pred = transformers.pipeline("text-classification", model=model, tokenizer=tokenizer, top_k=None, device=device) 
 
 # SHAP explainer
 explainer = shap.Explainer(pred)
 
-# NER pipeline (Added device mapping for faster inference)
+# NER pipeline
 ner_tokenizer = AutoTokenizer.from_pretrained("d4data/biomedical-ner-all")
-ner_model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all").to(device)
-ner_pipe = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, aggregation_strategy="simple", device=device) 
+ner_model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all")
+ner_pipe = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, aggregation_strategy="simple") 
 
-# ==========================================
-# 2. Prediction Function
-# ==========================================
 def adr_predict(x):
     text_input = str(x).lower()
+    encoded_input = tokenizer(text_input, return_tensors='pt').to(device)
+    output = model(**encoded_input)
     
-    # 1. Classification via the pipeline (reusing the pred object)
-    raw_results = pred(text_input)[0]
-    label_output = {item['label']: float(item['score']) for item in raw_results}
+    scores = torch.softmax(output.logits, dim=-1)[0].detach().cpu().numpy()
     
-    # 2. SHAP Logic
     try:
         shap_values = explainer([text_input])
         local_plot = shap.plots.text(shap_values[0], display=False) 
     except Exception as e:
         local_plot = f"<p style='color:red;'>SHAP explanation error: {e}</p>"
 
-    # 3. NER Logic
     try:
         res = ner_pipe(text_input)
         entity_colors = {
@@ -62,32 +54,27 @@ def adr_predict(x):
             'Biological_structure':'#d9d9d9'
         }
         
-        # Inline styles to force black text and white background for readability
-        htext = "<div style='line-height: 2.0; font-size: 1.1em; color: black; background-color: white; padding: 10px;'>"
+        # FIX: Added inline "color: black;" to force all un-highlighted text to be black
+        htext = "<div style='line-height: 2.0; font-size: 1.1em; color: black;'>"
         prev_end = 0
         res = sorted(res, key=lambda x: x['start'])
-        
         for entity in res:
             start, end = entity['start'], entity['end']
             word = text_input[start:end]
             color = entity_colors.get(entity['entity_group'], '#f3f3f3')
             
-            htext += f"<span>{text_input[prev_end:start]}</span>"
-            htext += (f"<mark style='background-color:{color}; color: black; padding: 2px 4px; "
-                      f"border-radius: 4px; font-weight: 500;'>{word} "
-                      f"<small style='opacity: 0.7; font-size: 0.7em;'>[{entity['entity_group']}]</small></mark>")
+            htext += f"{text_input[prev_end:start]}"
+            # Highlighted text is also explicitly set to black
+            htext += f"<mark style='background-color:{color}; color: black; padding: 2px 4px; border-radius: 4px; font-weight: 500;'>{word} <small style='opacity: 0.7;'>[{entity['entity_group']}]</small></mark>"
             prev_end = end
-            
-        htext += f"<span>{text_input[prev_end:]}</span></div>"
-    except Exception as e:
-        htext = f"<p style='color: black;'>NER processing error: {e}</p>"
+        htext += text_input[prev_end:] + "</div>"
+    except:
+        htext = "<p style='color: black;'>NER processing error.</p>"
         
+    label_output = {"Severe Reaction": float(scores[1]), "Non-severe Reaction": float(scores[0])}
     return label_output, local_plot, htext
 
-# ==========================================
-# 3. Gradio Interface
-# ==========================================
-# CSS ensures Gradio's dark mode doesn't override the white background and black text
+# FIX: Added !important tags to ensure Gradio's dark mode doesn't override the white background and black text
 custom_css = """
 .gradio-container { font-family: 'Inter', system-ui, sans-serif; }
 .main-header { text-align: center; margin-bottom: 2rem; }
@@ -95,7 +82,7 @@ custom_css = """
 footer { visibility: hidden; }
 """
 
-with gr.Blocks(title="ADR Detector", css=custom_css, theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="ADR Detector") as demo:
     with gr.Column(elem_classes="main-header"):
         gr.Markdown("# Adverse Drug Reaction (ADR) Detector")
         gr.Markdown("Analyze clinical text for potential medication-related severity and key medical entities.")
@@ -115,9 +102,7 @@ with gr.Blocks(title="ADR Detector", css=custom_css, theme=gr.themes.Soft()) as
             gr.Examples(
                 examples=[
                     ["A 35 year-old male had severe headache after taking Aspirin. The lab results were normal."],
-                    ["A 35 year-old female had minor pain in upper abdomen after taking Acetaminophen."],
-                    ["A 62-year-old female presented with shortness of breath and anaphylaxis minutes after intravenous Penicillin administration."],
-                    ["Patient felt slight drowsiness and dry mouth after taking 10mg of Cetirizine, but no other symptoms were noted."]
+                    ["A 35 year-old female had minor pain in upper abdomen after taking Acetaminophen."]
                 ],
                 inputs=[prob1]
             )
@@ -141,5 +126,4 @@ with gr.Blocks(title="ADR Detector", css=custom_css, theme=gr.themes.Soft()) as
         outputs=[label, local_plot, htext]
     )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch(css=custom_css, theme=gr.themes.Soft())