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

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

# Load base model in CPU-only mode
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    device_map="auto",  # Use 'auto' or manually move later
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True  # Ensures meta device usage
)

# Move the model to CPU safely
model = model.to_empty(device=torch.device("cpu"))

# Now load the LoRA adapter
from peft import PeftModel
model = PeftModel.from_pretrained(model, "lora_adapter", device_map="cpu")
model.eval()

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

### USER:
What is a tuple in Python?

### ASSISTANT:
A tuple in Python is an ordered and immutable collection of elements. Tuples are defined using parentheses, like (1, 2, 3). They are useful when you want to store multiple items and ensure they cannot be changed.

### USER:
{instruction}

### ASSISTANT:
"""

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

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

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

if st.button("Get Answer") and user_input.strip():
    with st.spinner("Thinking..."):
        response = chat(user_input)
        st.markdown("**Answer:**")
        st.write(response)