import streamlit as st
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load fine-tuned model
model_path = "./tinyllama_lora_finetuned"
st.write("Loading fine-tuned TinyLlama... (CPU Mode)")

model = AutoModelForCausalLM.from_pretrained(model_path, device_map="cpu")
tokenizer = AutoTokenizer.from_pretrained(model_path)

st.title("Fine-Tuned TinyLlama Chatbot (LoRA)")
st.write("🚀 Chatbot trained with LoRA on CPU.")

user_input = st.text_area("Enter your prompt:", "")

if st.button("Generate Response"):
    if user_input:
        with st.spinner("Generating response..."):
            inputs = tokenizer(user_input, return_tensors="pt").to("cpu")
            output = model.generate(**inputs, max_length=100)
            response = tokenizer.decode(output[0], skip_special_tokens=True)
            st.write("**Response:**")
            st.write(response)
    else:
        st.warning("Please enter a prompt!")

st.write("✅ Running on CPU - May be slow.")