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

# Load tokenizer and base model
base_model_path = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(base_model_path)
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_path,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto" if torch.cuda.is_available() else None
)

# Load LoRA adapter
model = PeftModel.from_pretrained(base_model, "lora_adapter")
model.eval()

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

user_input = st.text_input("Your question")

if user_input:
    # Prompt template that helps model decide to answer or reject
    prompt = f"""
You are a helpful and expert Python programming tutor.
Answer only questions related to Python programming.
If the question is unrelated to Python (like history, math, etc), politely respond:
"Sorry, I can only answer Python-related questions."

### Question:
{user_input}

### Answer:
"""

    inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=True).to(model.device)
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=200,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )

    decoded = tokenizer.decode(output[0], skip_special_tokens=True)

    # Clean the output to only show the answer
    if "### Answer:" in decoded:
        final_answer = decoded.split("### Answer:")[-1].strip()
    else:
        final_answer = decoded.strip()

    st.markdown(f"**Answer:** {final_answer}")