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

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("TinyLLaMA/TinyLLaMA-1.1B-Chat-v1.0")

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

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

# Set title
st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
st.markdown("Ask me any **Python programming** question:")

# User input
user_question = st.text_input("Your question")

if user_question:
    with st.spinner("Thinking..."):

        # Clean prompt
        prompt = f"""
You are a helpful and expert Python programming tutor.
If the question is about Python, explain clearly with examples.
If the question is unrelated to Python, respond with "Sorry, I can only answer Python-related questions."

Question: {user_question}
Answer:"""

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

        output = model.generate(
        **inputs,
        max_new_tokens=512,   # allow longer answers
        do_sample=True,
        top_p=0.9,
        temperature=0.7,
        repetition_penalty=1.1

        )

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

        # Extract only the generated answer after "Answer:"
        answer_start = decoded_output.find("Answer:")
        answer = decoded_output[answer_start + len("Answer:"):].strip() if answer_start != -1 else decoded_output.strip()

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