import gradio as gr

from generic import generate_visual

def generate_viz(text, target_class):
    
    print(f"Input Text: {text}")
    print(f"Target Class: {target_class}")
    text_batch=[]
    text_batch.append(text)
    token_import, html_page = generate_visual(text_batch=text_batch, target_class=target_class)
    return token_import, html_page

title = "Explain BERT 😄"


iface = gr.Interface(fn=generate_viz, inputs=[
        gr.Text(label="Input Text"),
        gr.Dropdown(
            [0,1],
            label="Target_Class",
            info="0 is negative, 1 is positive",
        ),
    ],
    outputs=[ gr.Text(label="Token_Importance"), gr.HTML(label="visualize sentence"),],
    title=title,
    allow_flagging="never",
    cache_examples=True,
    examples=[
        ["This movie was the best movie I have ever seen! some scenes were ridiculous, but acting was great.", 1],
        ["I really didn't like this movie. Some of the actors were good, but overall the movie was boring.", 0],
        ["I hate that I love you.", 0],
        ["I hate that I love you.", 1],
    ],
)

iface.launch(debug=True)