app.py

import gradio as gr
import torch
from nextus_regressor_class import *
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":
        tokens = nltk.word_tokenize(txt)
        #print("word")
    elif tok_level == "sentence":
        #print("sentence")
        tokens = nltk.sent_tokenize(txt)
    else:
        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)
        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]
    return list(zip(tokens, labels))

demo = gr.Interface(shap,
                    [
                         gr.Textbox(label="기사", lines=30, placeholder="기사를 입력하세요."),
                         gr.Radio(choices=["sentence", "word"], label="해설 표시 단위", value="sentence", info="문장 단위의 해설은 sentence를, 단어 단위의 해설은 word를 선택하세요.")
                    ],
                    #gr.Textbox(label="Test Output"),
                    gr.HighlightedText(
                        label="Diff",
                        combine_adjacent=True,
                        show_legend=True,
                        color_map={"+": "red", "-": "green"}),
                   
                    theme=gr.themes.Base())

demo.launch()