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

MODEL_PATH = "./tinyllama-python-tutor-lora"

st.title("TinyLLaMA Python Tutor 💬")

@st.cache_resource
def load_model():
    bnb_config = BitsAndBytesConfig(
        load_in_8bit=True,
        llm_int8_threshold=6.0,
        llm_int8_skip_modules=None,
        llm_int8_enable_fp32_cpu_offload=True
    )

    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        quantization_config=bnb_config,
        device_map="auto"
    )
    return tokenizer, model

tokenizer, model = load_model()

prompt = st.text_area("Ask me about Python:", height=200)

if st.button("Generate Response"):
    if prompt.strip():
        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
        output = model.generate(**inputs, max_new_tokens=200, do_sample=True)
        response = tokenizer.decode(output[0], skip_special_tokens=True)
        st.write("### Response")
        st.success(response)
    else:
        st.warning("Please enter a prompt!")