import streamlit as st

from intent_classifier.model import IntentClassifier
from consts import FLAN_T5_SMALL, FLAN_T5_BASE


@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()

# text = "Hey, I want to stop my subscription. Can I do that?"
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")
if is_clicked:
    output = model.structured_predict(input, prompt_options, print_input=True)
    st.write(output)