import streamlit as st
from keras.models import load_model

st.set_page_config(page_title="Fine-tuned Gemma Chatbot", layout="centered")

# Load the Keras model
@st.cache_resource
def load_keras_model():
    model = load_model("gemma_finetuned.keras")
    return model

model = load_keras_model()

# UI
st.title("💬 Fine-tuned Gemma Code Generator")
prompt = st.text_area("Enter your instruction", value="Write a Python function to reverse a string")

if st.button("Generate"):
    # Simple inference logic - you may need to adapt this based on how your model generates text
    sampler = keras_nlp.samplers.TopKSampler(k=5, seed=2)
    model.compile(sampler=sampler)
    response = model.generate(prompt, max_length=256)

    st.subheader("Response:")
    st.code(response, language="python")