Update app.py

app.py

@@ -5,6 +5,7 @@ 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
@@ -18,14 +19,17 @@ device = "cuda:0" if torch.cuda.is_available() else "cpu"
 tokenizer = AutoTokenizer.from_pretrained("jschwaller/ADRv2024")  
 model = AutoModelForSequenceClassification.from_pretrained("jschwaller/ADRv2024")
 
-pred = transformers.pipeline("text-classification", model=model, tokenizer=tokenizer, return_all_scores=True)
+# Build a pipeline object for predictions
+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")
+ner_pipe = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, aggregation_strategy="simple") # pass device=0 if using gpu
+#
 
 def adr_predict(x):
     encoded_input = tokenizer(x, return_tensors='pt')
@@ -38,13 +42,13 @@ def adr_predict(x):
 
     res = ner_pipe(x)
     entity_colors = {
-        'Severity': '#E63946',
-        'Sign_symptom': '#2A9D8F',
-        'Medication': '#457B9D',
-        'Age': '#F4A261',
-        'Sex': '#F4A261',
-        'Diagnostic_procedure': '#9C6644',
-        'Biological_structure': '#BDB2FF',
+        '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
         }
 
     htext = ""
@@ -71,7 +75,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; }
+h1, h2, h3, h4, h5, h6, p, label, .markdown { color: #87CEEB; } /* Ensuring that all text elements are consistently light blue */
 .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; }
 """
@@ -89,7 +93,15 @@ with gr.Blocks(css=css) as demo:
         with gr.Column(visible=True):
             local_plot = gr.HTML(label='Shap:')
             htext = gr.HTML(label="NER")
+
+    submit_btn.click(
+        main,
+        [prob1],
+        [label, local_plot, htext],
+        api_name="adr"
+    )
     
+    gr.Markdown("### Legend")
     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> " +
@@ -97,13 +109,6 @@ with gr.Blocks(css=css) as demo:
                             "<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],
-        [label, local_plot, htext],
-        api_name="adr"
-    )
-    
     gr.Row([legend])
 
     with gr.Row():
@@ -111,5 +116,3 @@ with gr.Blocks(css=css) as demo:
         gr.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."]], 
                     [prob1], [label, local_plot, htext], main, cache_examples=True)
-
-demo.launch()