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

# Load tokenizer and model
base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(base_model)

model = AutoModelForCausalLM.from_pretrained(base_model, device_map="cpu")
model = PeftModel.from_pretrained(model, "lora_adapter")  # Change if needed
model.eval()

# Format prompt (no USER/ASSISTANT lines to confuse model)
def format_prompt(instruction):
    return f"""You are a helpful and expert Python programming tutor.
You only answer questions related to Python programming.
If the question is unrelated to Python, say:
"Sorry, I can only answer Python-related questions."

Question: {instruction}
Answer:"""

# Generate answer
def chat(instruction):
    prompt = format_prompt(instruction)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            do_sample=False,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.1
        )

    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Only return the model's answer
    return full_output.split("Answer:")[-1].strip()

# Streamlit UI
st.set_page_config(page_title="🐍 Python Tutor Chatbot (LoRA)")
st.title("🐍 Python Tutor Chatbot (LoRA)")
st.write("Ask me Python programming questions!")

user_input = st.text_area("Your question:")

if st.button("Get Answer") and user_input.strip():
    with st.spinner("Thinking..."):
        response = chat(user_input)
        st.markdown("**Answer:**")
        if "```" in response:
            st.markdown(response)
        else:
            st.write(response)