Update app.py

app.py

@@ -42,10 +42,10 @@ explainer = shap.Explainer(pred)
 #     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_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
+ner_pipe = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, aggregation_strategy="simple") # pass device=0 if using gpu
 #
 
 def adr_predict(x):
@@ -60,41 +60,39 @@ 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])
 
-    # res = ner_pipe(x)
+    res = ner_pipe(x)
     
-    # entity_colors = {
-    # 'Severity': 'red',
-    # 'Sign_symptom': 'green',
-    # 'Medication': 'blue',
-    # 'Age': 'yellow',
-    # 'Sex':'yellow',
-    # 'Diagnostic_procedure':'gray',
-    # 'Biological_structure':'silver'}
-
-    # htext = ""
-    # prev_end = 0
-
-    # for entity in res:
-    #     start = entity['start']
-    #     end = entity['end']
-    #     word = entity['word'].replace("##", "")
-    #     color = entity_colors[entity['entity_group']]
+    entity_colors = {
+    'Severity': 'red',
+    'Sign_symptom': 'green',
+    'Medication': 'blue',
+    'Age': 'yellow',
+    'Sex':'yellow',
+    'Diagnostic_procedure':'gray',
+    'Biological_structure':'silver'}
+
+    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 += f"{x[prev_end:start]}<mark style='background-color:{color};'>{word}</mark>"
+        prev_end = end
 
-    # htext += x[prev_end:]
+    htext += x[prev_end:]
    
-    return {"Severe Reaction": float(scores.numpy()[1]), "Non-severe Reaction": float(scores.numpy()[0])}, local_plot
-    # ,htext
+    return {"Severe Reaction": float(scores.numpy()[1]), "Non-severe Reaction": float(scores.numpy()[0])}, local_plot,htext
     # ,{"Contains Medication": float(med), "No Medications": float(1-med)} , {"Contains Symptoms": float(sym), "No Symptoms": float(1-sym)}
 
 
 def main(prob1):
     text = str(prob1).lower()
     obj = adr_predict(text)
-    return obj[0],obj[1]
-    # ,obj[2]
+    return obj[0],obj[1],obj[2]
 
 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."""
@@ -114,7 +112,7 @@ with gr.Blocks(title=title) as demo:
 
         with gr.Column(visible=True) as output_col:
             local_plot = gr.HTML(label = 'Shap:')
-            # htext = gr.HTML(label="NER")
+            htext = gr.HTML(label="NER")
             # med = gr.Label(label = "Contains Medication")
             # sym = gr.Label(label = "Contains Symptoms")
             
@@ -122,8 +120,7 @@ with gr.Blocks(title=title) as demo:
         main,
         [prob1],
         [label
-         ,local_plot
-         # , htext
+         ,local_plot, htext
          # , med, sym
         ], api_name="adr"
     )
@@ -132,8 +129,7 @@ with gr.Blocks(title=title) as demo:
         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."],
                      ["A 35 year-old female had minor pain in upper abdomen after taking Acetaminophen."]], 
-                    [prob1], [label,local_plot
-                              # , htext
+                    [prob1], [label,local_plot, htext
          # , med, sym
                              ], main, cache_examples=True)