import streamlit as st
from transformers import AutoModelForSequenceClassification, AutoTokenizer

@st.cache(allow_output_mutation=True)
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("./my_model")
    model = AutoModelForSequenceClassification.from_pretrained("./my_model")
    return model, tokenizer

model, tokenizer = load_model()

text_input = st.text_area("Enter text here:")
if st.button("Predict"):
    inputs = tokenizer(text_input, return_tensors="pt")
    outputs = model(**inputs)
    prediction = outputs.logits.argmax(-1).item()
    st.write(f"Prediction: {prediction}")