Update app.py

app.py

@@ -6,31 +6,39 @@ import nltk
 model = NextUsRegressor()
 model.load_state_dict(torch.load("./nextus_regressor1012.pt"))
 model.eval()
-
+mask = "[MASKED]"
 def shap(txt, tok_level):
     batch = [txt]
     if tok_level == "word":
-        print("word")
+        tokens = nltk.word_tokenize(txt)
+        #print("word")
     elif tok_level == "sentence":
-        print("sentence")
+        #print("sentence")
+        tokens = nltk.sent_tokenize(txt)
     else:
-        print("this token granularity not supported")
+        pass
+        #print("this token granularity not supported")
     #tokens = nltk
+    for i, _ in enumerate(tokens):
+        batch.append(" ".join([s for j, s in enumerate(tokens) if j!=i]))
     with torch.no_grad():
         y_pred = model(txt)
-    return txt
+        y_offs = model(batch)
+        shaps = (y_offs - y_pred).tolist()[0] # convert to list and make tuple to be returned
+    return [token, shap for token, shap in list(zip(tokens, shaps))]
 
 demo = gr.Interface(shap,
                     [
                          gr.Textbox(label="기사", lines=30, placeholder="기사를 입력하세요."),
-                         gr.Radio(["sentence", "word"], value="sentence", info="문장 단위의 해설은 sentence를 단어 단위의 해설은 word를 선택하세요.")
+                         gr.Radio(choices=["sentence", "word"], label="해설 표시 단위", value="sentence", info="문장 단위의 해설은 sentence를 단어 단위의 해설은 word를 선택하세요.")
                     ],
-                    gr.Textbox(label="Slant Score"))
-                    #gr.HighlightedText(
-                    #    label="Diff",
-                    #    combine_adjacent=True,
-                    #    show_legend=True,
-                    #    color_map={"+": "red", "-": "green"}),
+                    #gr.Textbox(label="Slant Score"))
+                    gr.HighlightedText(
+                        label="Diff",
+                        combine_adjacent=True,
+                        show_legend=True,
+                        color_map={"+": "red", "-": "green"})
+                   )
                     #theme=gr.themes.Base())
 
 if __name__ == "__main__":