Update app.py

app.py

@@ -41,6 +41,10 @@ def sym_score(x):
     score_1sym = x['score']
     return round(score_1sym,3)
 
+ner_tokenizer = AutoTokenizer.from_pretrained("d4data/biomedical-ner-all")
+ner_model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all")
+
+ner_pipe = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple") # pass device=0 if using gpu
 ##
 
 def adr_predict(x):
@@ -54,14 +58,16 @@ def adr_predict(x):
 
     med = med_score(classifier(x+str(", There is a medication."))[0])
     sym = sym_score(classifier(x+str(", There is a symptom."))[0])
+
+    htext = ner_pipe(x)
    
-    return {"Severe Reaction": float(scores.numpy()[1]), "Non-severe Reaction": float(scores.numpy()[0])}, local_plot, {"Contains Medication": float(med), "No Medications": float(1-med)} , {"Contains Symptoms": float(sym), "No Symptoms": float(1-sym)}
+    return {"Severe Reaction": float(scores.numpy()[1]), "Non-severe Reaction": float(scores.numpy()[0])}, local_plot, {"Contains Medication": float(med), "No Medications": float(1-med)} , {"Contains Symptoms": float(sym), "No Symptoms": float(1-sym)},htext
 
 
 def main(prob1):
     text = str(prob1).lower()
     obj = adr_predict(text)
-    return obj[0],obj[1],obj[2],obj[3]
+    return obj[0],obj[1],obj[2],obj[3],obj[4]
 
 title = "Welcome to **ADR Detector** 🪐"
 description1 = """This app takes text (up to a few sentences) and predicts to what extent the text describes severe (or non-severe) adverse reaction to medicaitons. Please do NOT use for medical diagnosis."""
@@ -78,6 +84,7 @@ with gr.Blocks(title=title) as demo:
         with gr.Column(visible=True) as output_col:
             label = gr.Label(label = "Predicted Label")
             local_plot = gr.HTML(label = 'Shap:')
+            htext = gr.HighlightedText(label="Diff", combine_adjacent=True, show_legend=True)
 
         with gr.Column(visible=True) as output_col:
             med = gr.Label(label = "Contains Medication")
@@ -87,14 +94,13 @@ with gr.Blocks(title=title) as demo:
         main,
         [prob1],
         [label
-         ,local_plot, med
-         , sym
+         ,local_plot, med, sym, htext
         ], api_name="adr"
     )
     
     with gr.Row():
         gr.Markdown("### Click on any of the examples below to see how it works:")
-        gr.Examples([["I had severe headache after taking Aspirin."],["I had minor stomachache after taking Acetaminophen."]], [prob1], [label,local_plot, med, sym
-                                                                          ], main, cache_examples=True)
+        gr.Examples([["I had severe headache after taking Aspirin."],["I had minor stomachache after taking Acetaminophen."]], 
+                    [prob1], [label,local_plot, med, sym,htext], main, cache_examples=True)
     
 demo.launch()