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

# Load tokenizer and model (CPU)
base_model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
adapter_path = "lora_adapter"  # Your LoRA adapter folder path

# Force CPU usage
device = torch.device("cpu")

tokenizer = AutoTokenizer.from_pretrained(base_model_name, use_fast=True)
base_model = AutoModelForCausalLM.from_pretrained(base_model_name).to(device)
model = PeftModel.from_pretrained(base_model, adapter_path).to(device)

# Streamlit UI setup
st.set_page_config(
    page_title="Python Tutor Chatbot",
    page_icon="🐍",
    layout="centered"
)

st.title("🐍 Python Tutor Chatbot")
st.markdown("Ask me anything about Python programming!")

# Prompt template
def create_prompt(user_input):
    return f"""
You are a helpful and knowledgeable AI Python Tutor. Your job is to answer only Python-related programming questions.
If the question is unrelated to Python, kindly respond with: "Sorry, I can only answer Python programming questions."

### Instruction:
{user_input}

### Response:
"""

# Chat interface
user_input = st.text_input("Your Python Question:")

if user_input:
    with st.spinner("Generating response..."):
        prompt = create_prompt(user_input)
        inputs = tokenizer(prompt, return_tensors="pt").to(device)

        with torch.no_grad():
            output = model.generate(
                **inputs,
                max_new_tokens=200,
                temperature=0.7,
                do_sample=True,
                top_p=0.9,
                top_k=50
            )

        response = tokenizer.decode(output[0], skip_special_tokens=True)
        final_response = response.split("### Response:")[-1].strip()

        st.markdown("**Answer:**")
        st.write(final_response)