Update app.py

app.py

@@ -2,11 +2,13 @@ import gradio as gr
 import torch
 from nextus_regressor_class import *
 import nltk
+from pprint import pprint
 
 model = NextUsRegressor()
 model.load_state_dict(torch.load("./nextus_regressor1012.pt"))
 model.eval()
 mask = "[MASKED]"
+threshold = 0.05
 def shap(txt, tok_level):
     batch = [txt]
     if tok_level == "word":
@@ -26,9 +28,21 @@ def shap(txt, tok_level):
         y_offs = model(batch)
         shaps = (y_offs - y_pred).tolist() # convert to list and make tuple to be returned
     shapss = [s[0] for s in shaps]
-    labels = ["+" if s < 0 else "-" for s in shapss]
+    labels = list()
+    for s in shapss:
+        if s <= -1.0*threshold:
+            labels.append("+")
+        elif s >= threshold:
+            labels.append("-")
+        else:
+            labels.append(None)
+    # labels = ["+" if s < -1.0*threshold "-" elif s > threshold else " " for s in shapss]
+    # print(len(tokens), len(labels))
+    # print(list(zip(tokens, labels)))
+    # pprint(list(zip(tokens, shapss)))
+    # return str(list(zip(tokens, labels)))
     return list(zip(tokens, labels))
-    #return txt
+    # return txt
 
 demo = gr.Interface(shap,
                     [
@@ -43,4 +57,4 @@ demo = gr.Interface(shap,
                         color_map={"+": "red", "-": "green"}),
                     theme=gr.themes.Base())
 
-demo.launch()
+demo.launch()