Update app.py

app.py

@@ -45,8 +45,8 @@ def sym_score(x):
 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
-##
+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')
@@ -60,7 +60,26 @@ 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)
+    res = ner_pipe(x)
+    
+    entity_colors = {
+    'Severity': 'red',
+    'Sign_symptom': 'green',
+    'Medication': 'blue'}
+
+    htext = ""
+    prev_end = 0
+
+    for entity in res:
+        start = entity['start']
+        end = entity['end']
+        word = entity['word'].replace("##", "")
+        color = entity_colors[entity['entity_group']]
+        
+        htext += f"{x[prev_end:start]}<mark style='background-color:{color};'>{word}</mark>"
+        prev_end = end
+
+    htext += x[prev_end:]
    
     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
 
@@ -85,7 +104,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="NER", combine_adjacent=True, show_legend=True)
+            htext = gr.HTML(label="NER")
 
         with gr.Column(visible=True) as output_col:
             med = gr.Label(label = "Contains Medication")
@@ -101,7 +120,8 @@ with gr.Blocks(title=title) as demo:
     
     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."]], 
+        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, med, sym,htext], main, cache_examples=True)
     
-demo.launch()
+demo.launch()