Update app.py

app.py

@@ -5,7 +5,6 @@ import numpy as np
 import scipy as sp
 import torch
 import transformers
-from transformers import pipeline
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification
 
 import matplotlib.pyplot as plt
@@ -19,45 +18,33 @@ device = "cuda:0" if torch.cuda.is_available() else "cpu"
 tokenizer = AutoTokenizer.from_pretrained("jschwaller/ADRv2024")  
 model = AutoModelForSequenceClassification.from_pretrained("jschwaller/ADRv2024")
 
-# Build a pipeline object for predictions
-pred = transformers.pipeline("text-classification", model=model, 
-                             tokenizer=tokenizer, return_all_scores=True)
+pred = transformers.pipeline("text-classification", model=model, tokenizer=tokenizer, return_all_scores=True)
 
 explainer = shap.Explainer(pred)
 
 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") # pass device=0 if using gpu
-
-# fixing the colors
-# Create a custom color map
-cmap = {'0': '#457B9D',  # Non-severe reactions in blue
-        '1': '#E63946'}  # Severe reactions in red
-
+ner_pipe = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, aggregation_strategy="simple")
 
 def adr_predict(x):
     encoded_input = tokenizer(x, return_tensors='pt')
     output = model(**encoded_input)
     scores = output[0][0].detach()
     scores = torch.nn.functional.softmax(scores)
-
-    # Generate SHAP values and use the custom color map
+   
     shap_values = explainer([str(x).lower()])
-    
-    # Ensure the color depends on the output class; customize as needed
-    base_colors = {label: cmap[str(label)] for label in range(len(scores))} 
-    shap.plots.text(shap_values[0], color=base_colors, display=False)
+    local_plot = shap.plots.text(shap_values[0], display=False)
 
     res = ner_pipe(x)
     entity_colors = {
-        'Severity': '#E63946',  # a vivid red
-        'Sign_symptom': '#2A9D8F',  # a deep teal
-        'Medication': '#457B9D',  # a dusky blue
-        'Age': '#F4A261',  # a sandy orange
-        'Sex': '#F4A261',  # same sandy orange for consistency with 'Age'
-        'Diagnostic_procedure': '#9C6644',  # a brown
-        'Biological_structure': '#BDB2FF',  # a light pastel purple
+        'Severity': '#E63946',
+        'Sign_symptom': '#2A9D8F',
+        'Medication': '#457B9D',
+        'Age': '#F4A261',
+        'Sex': '#F4A261',
+        'Diagnostic_procedure': '#9C6644',
+        'Biological_structure': '#BDB2FF',
         }
 
     htext = ""
@@ -84,7 +71,7 @@ description1 = "This app takes text (up to a few sentences) and predicts to what
 
 css = """
 body { font-family: 'Roboto', sans-serif; background-color: #333; color: #87CEEB; }
-h1, h2, h3, h4, h5, h6, p, label, .markdown { color: #87CEEB; } /* Ensuring that all text elements are consistently light blue */
+h1, h2, h3, h4, h5, h6, p, label, .markdown { color: #87CEEB; }
 .textbox { width: 100%; border-radius: 10px; border: 1px solid #ccc; background-color: white; color: black; }
 .button { background-color: #FF6347; color: white; border: none; border-radius: 10px; padding: 10px 20px; cursor: pointer; }
 """
@@ -102,7 +89,14 @@ with gr.Blocks(css=css) as demo:
         with gr.Column(visible=True):
             local_plot = gr.HTML(label='Shap:')
             htext = gr.HTML(label="NER")
-
+    
+    legend = gr.HTML(value="<div style='margin-top: 20px;'><strong>Legend:</strong><br>" +
+                            "<mark style='background-color:#E63946;'>Severity</mark> " +
+                            "<mark style='background-color:#2A9D8F;'>Sign/Symptom</mark> " +
+                            "<mark style='background-color:#457B9D;'>Medication</mark> " +
+                            "<mark style='background-color:#F4A261;'>Age/Sex</mark> " +
+                            "<mark style='background-color:#9C6644;'>Diagnostic Procedure</mark> " +
+                            "<mark style='background-color:#BDB2FF;'>Biological Structure</mark></div>")
     submit_btn.click(
         main,
         [prob1],
@@ -110,6 +104,8 @@ with gr.Blocks(css=css) as demo:
         api_name="adr"
     )
     
+    gr.Row([legend])
+
     with gr.Row():
         gr.Markdown("### Click on any of the examples below to see how it works:")
         gr.Examples([["A 35 year-old male had severe headache after taking Aspirin. The lab results were normal."],