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

# Load base model & tokenizer
base_model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
adapter_path = "lora_adapter"  # path to your LoRA adapter directory

@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
    base_model = AutoModelForCausalLM.from_pretrained(base_model_name, device_map="auto")
    model = PeftModel.from_pretrained(base_model, adapter_path)
    model.eval()
    return tokenizer, model

tokenizer, model = load_model()

# Prompt formatting
def format_prompt(user_input):
    return f"""You are a helpful and knowledgeable Python tutor chatbot.

You only answer questions related to Python programming, including:
- Python syntax, functions, loops, and conditionals
- Standard libraries and popular packages (e.g., NumPy, pandas)
- Debugging and code explanation
- Python tools, environments, and tips

If a question is not related to Python, reply with:
"Sorry, I can only answer Python-related questions."

### Instruction:
{user_input}

### Response:"""

# Chat handler
def chat(user_input):
    prompt = format_prompt(user_input)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=200,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id
        )
    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
    return decoded.split("### Response:")[-1].strip()

# Streamlit UI
st.title("🧑‍🏫 Python Tutor Chatbot")
st.write("Ask me anything about Python programming!")

user_input = st.text_area("Your Question", height=150)
if st.button("Ask"):
    if user_input.strip():
        with st.spinner("Thinking..."):
            answer = chat(user_input)
        st.markdown("### 💡 Answer:")
        st.write(answer)