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

# Load tokenizer and base model
base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
lora_path = "./lora_adapter"

tokenizer = AutoTokenizer.from_pretrained(base_model)

# Load base model normally (for CPU)
model = AutoModelForCausalLM.from_pretrained(base_model)
model = PeftModel.from_pretrained(model, lora_path)
model.eval()

# Move to CPU explicitly
device = torch.device("cpu")
model.to(device)

# Streamlit UI
st.set_page_config(page_title="🧠 TinyLLaMA Python Tutor (LoRA)")
st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
st.write("Ask me any **Python programming** question:")

user_input = st.text_input("Your question", placeholder="e.g. What is a lambda function in Python?")

if user_input:
    # Check if it's a Python-related question
    if "python" not in user_input.lower() and "py" not in user_input.lower():
        st.warning("❌ Sorry, I can only answer Python programming questions.")
    else:
        system_prompt = (
            "You are an expert Python tutor. Provide clear, concise, and accurate explanations with examples. "
            "If the user's question is not related to Python programming, respond with: "
            "'Sorry, I can only help with Python programming questions.'"
        )
        prompt = f"<|system|>\n{system_prompt}</s>\n<|user|>\n{user_input}</s>\n<|assistant|>"

        inputs = tokenizer(prompt, return_tensors="pt").to(device)

        with torch.no_grad():
            with st.spinner("Thinking..."):
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=150,
                    temperature=0.7,
                    top_p=0.95,
                    do_sample=True,
                    eos_token_id=tokenizer.eos_token_id,
                    pad_token_id=tokenizer.eos_token_id
                )

            decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
            answer = decoded_output.split("<|assistant|>")[-1].strip()
            st.success(f"💬 Answer:\n\n{answer}")