import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
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"  # make sure your LoRA adapter folder is named like this

bnb_config = BitsAndBytesConfig(load_in_4bit=True,
                                bnb_4bit_compute_dtype=torch.bfloat16)

tokenizer = AutoTokenizer.from_pretrained(base_model)
model = AutoModelForCausalLM.from_pretrained(base_model,
                                             quantization_config=bnb_config,
                                             torch_dtype=torch.bfloat16,
                                             device_map="auto")

model = PeftModel.from_pretrained(model, lora_path)
model.eval()

# 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:
    # Filtering logic: Only answer Python-related queries
    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 for tutor behavior
        system_prompt = (
            "You are a helpful and knowledgeable Python tutor. "
            "Answer the user's Python programming questions clearly and concisely. "
            "If the question is unclear, ask for clarification."
        )
        prompt = f"<|system|>\n{system_prompt}</s>\n<|user|>\n{user_input}</s>\n<|assistant|>"

        inputs = tokenizer(prompt, return_tensors="pt").to(model.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)

            # Extract answer only (remove prompt)
            answer = decoded_output.split("<|assistant|>")[-1].strip()
            st.success(f"💬 Answer:\n\n{answer}")