import streamlit as st
from transformers import pipeline

# --------------------------------
# Streamlit Page Settings
# --------------------------------
st.set_page_config(
    page_title="AI Chatbot",
    layout="centered"
)

st.title("🤖 AI Chatbot")

st.write("Chat with MiniCPM 1B model")

# --------------------------------
# Load Model
# --------------------------------
@st.cache_resource
def load_model():

    generator = pipeline(
        "text-generation",
        model="openbmb/MiniCPM5-1B"
    )

    return generator

# Load the model
generator = load_model()

# --------------------------------
# User Input
# --------------------------------
user_input = st.text_area(
    "Enter your prompt:"
)

# --------------------------------
# Generate Response
# --------------------------------
if st.button("Generate Response"):

    if user_input.strip() != "":

        with st.spinner("Generating response..."):

            try:

                result = generator(
                    user_input,
                    max_new_tokens=100,
                    do_sample=True,
                    temperature=0.7
                )

                answer = result[0]["generated_text"]

                st.success(answer)

            except Exception as e:

                st.error(f"Error: {str(e)}")

    else:
        st.warning("Please enter a prompt.")