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

st.set_page_config(page_title="TinyLlama Chatbot")

@st.cache_resource
def load_model():
    model_name = "TinyLlama/TinyLlama-1.1B-chat-v1.0"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    model.eval()
    return tokenizer, model

st.title("Turbo Chatbot♑")

user_input = st.text_input("Ask a question")

if user_input:
    with st.spinner("Generating response..."):
        tokenizer, model = load_model()

        prompt = f"""<|system|>
You are a helpful AI assistant.
<|user|>
{user_input}
<|assistant|>
"""

        input_ids = tokenizer(prompt, return_tensors="pt").input_ids

        with torch.no_grad():
            output = model.generate(
                input_ids,
                max_new_tokens=120,
                temperature=0.6,
                top_p=0.9,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )

        decoded = tokenizer.decode(output[0], skip_special_tokens=True)
        response = decoded.split("<|assistant|>")[-1].strip()

        st.subheader("Response")
        st.write(response)