import streamlit as st
from transformers import pipeline

st.title("Hugging Face Model Demo")


@st.cache(allow_output_mutation=True)
def load_model():
    return pipeline("text-generation", model="klyang/MentaLLaMA-chat-7B")


model = load_model()

user_input = st.text_input("Enter your text:")
if user_input:
    with st.spinner("Generating response..."):
        result = model(user_input)
        st.success(result[0]["generated_text"])