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

# Set CPU device
device = torch.device("cpu")

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

base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32,  # You can try float16 if supported
    device_map={"": device}
)

# Load LoRA adapter (your fine-tuned weights)
model = PeftModel.from_pretrained(base_model, "lora_adapter", device_map={"": device})
model.eval()

# Format the prompt with filtering
def format_prompt(instruction):
    return f"""You are a helpful and expert Python programming tutor.
Only answer questions related to Python programming.
If the question is unrelated to Python, respond with:
"Sorry, I can only answer Python-related questions."

### Instruction:
{instruction}

### Response:
"""

# Generate answer
def chat(instruction):
    prompt = format_prompt(instruction)
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.split("### Response:")[-1].strip()

# Streamlit UI
st.title("🐍 Python Tutor Chatbot")
st.markdown("Ask me Python programming questions!")

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

if st.button("Answer"):
    if question.strip():
        with st.spinner("Thinking..."):
            response = chat(question)
        st.markdown("**Answer:**")
        st.write(response)
    else:
        st.warning("Please enter a question.")