import streamlit as st

from intent_classifier.model import IntentClassifier


@st.cache_resource
def load_model(model_name="Serj/intent-classifier", device=None):
    # st.write(f"model path: {model_name}")
    m = IntentClassifier(model_name=model_name, device=device)
    return m


model = load_model()

st.title("Simple Labels")
# text = "Hey, I want to stop my subscription. Can I do that?"
text = "How do I get my money back?"
input = st.text_area("text", value=text)
prompt_options = st.text_area("prompt_options", value="# renew subscription # account deletion # cancel subscription # resume subscription # refund requests # other # general # item damaged # malfunction # hello # intro # question")
is_clicked = st.button("Submit", key="submit smile")
if is_clicked:
    output = model.predict(input, prompt_options)
    st.write(output)

st.title("Complex labels")
text = "I was expecting my card to get this week and I'm wandering why didn't I receive it yet?"
input = st.text_area("text", value=text)
prompt_options_values = " # how do I locate my card # my card hasn't arrived yet # I would like to reactivate my card # where do I link my card? # How do I re add my card? # how do I add the card to my account? # status of the card order "
prompt_options = st.text_area("prompt_options", value=prompt_options_values)
is_clicked = st.button("Submit", key="submit complex")
if is_clicked:
    output = model.structured_predict(input, prompt_options, print_input=True)
    st.write(output)