Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import gradio as gr
+import shap
 import torch
 import tensorflow as tf
 from transformers import RobertaTokenizer, RobertaModel
@@ -15,12 +16,19 @@ def adr_predict(x):
     output = model(**encoded_input)
     scores = output[0][0].detach().numpy()
     scores = tf.nn.softmax(scores)
-    return {"Severe Reaction": float(scores.numpy()[1]), "Non-severe Reaction": float(scores.numpy()[0])} 
+    
+    # build a pipeline object to do predictions
+    pred = transformers.pipeline("text-classification", model=model, 
+                             tokenizer=tokenizer, device=0, return_all_scores=True)
+    explainer = shap.Explainer(pred)
+    shap_values = explainer([x])
+    shap_plot = shap.plots.text(shap_values)
+    return {"Severe Reaction": float(scores.numpy()[1]), "Non-severe Reaction": float(scores.numpy()[0])}, shap_plot
 
 def main(text):
     text = str(text).lower()
     obj = adr_predict(text)
-    return obj
+    return obj[0],obj[1]
 
 title = "Welcome to **ADR Detector** 🪐"
 description1 = """
@@ -41,18 +49,15 @@ with gr.Blocks(title=title) as demo:
         # color_map={"+++": "royalblue","++": "cornflowerblue",
         #  "+": "lightsteelblue", "NA":"white"})
         # NER = gr.HTML(label = 'NER:')
-        # intp =gr.HighlightedText(label="Word Scores",
-        # combine_adjacent=False).style(color_map={"++": "darkgreen","+": "green", 
-        #                                         "--": "darkred",
-        #                                         "-": "red", "NA":"white"})
+        shap_plot = gr.HighlightedText(label="Word Scores",combine_adjacent=False)
 
     submit_btn.click(
         main,
         [text],
-        [label], api_name="adr"
+        [label,shap_plot], api_name="adr"
     )
 
     gr.Markdown("### Click on any of the examples below to see to what extent they contain resilience messaging:")
-    gr.Examples([["I have minor pain."],["I have severe pain."]], [text], [label], main, cache_examples=True)
+    gr.Examples([["I have minor pain."],["I have severe pain."]], [text], [label,shap_plot], main, cache_examples=True)
     
 demo.launch()