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

st.set_page_config(page_title="TinyLLaMA Python Tutor", layout="centered")

st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
st.write("Ask me any Python programming question:")

@st.cache_resource
def load_model():
    base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
    adapter_path = "lora_adapter"

    tokenizer = AutoTokenizer.from_pretrained(base_model)
    model = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=torch.float32)
    model = PeftModel.from_pretrained(model, adapter_path)
    model.eval()

    return tokenizer, model

tokenizer, model = load_model()

def build_prompt(question):
    return (
        "You are a helpful and concise Python programming tutor. "
        "If the question is not about Python, respond with: "
        "'Sorry, I can only answer Python-related questions.'\n\n"
        f"Question: {question}\nAnswer:"
    )

question = st.text_input("Your question")

if question:
    prompt = build_prompt(question)
    inputs = tokenizer(prompt, return_tensors="pt")

    with st.spinner("Thinking..."):
        outputs = model.generate(
            **inputs,
            max_new_tokens=250,
            temperature=0.6,
            top_p=0.85,
            repetition_penalty=1.2,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
    answer = decoded_output.split("Answer:")[-1].strip()

    st.markdown(f"**💬 Answer:**\n\n{answer}")