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

# Load model and tokenizer
model_name = "lora_adapter"  # Update this to your LoRA model path

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")

# Chat function with prompt-based filtering
def chat(instruction):
    prompt = f"""You are a helpful and expert Python programming tutor.

Only answer questions that are clearly related to Python programming. Do not try to answer any questions that are greetings, general knowledge, or unrelated topics.

If the question is not related to Python, simply respond with:
"Sorry, I can only answer Python-related questions."

### Instruction:
{instruction}

### Response:
"""
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            temperature=0.7,
            top_p=0.95,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.split("### Response:")[-1].strip()

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

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

if user_input:
    with st.spinner("Generating response..."):
        response = chat(user_input)
        st.markdown("**Answer:**")
        st.markdown(response)