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

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

# Load base model
base_model = AutoModelForCausalLM.from_pretrained(
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    device_map="auto",
    torch_dtype=torch.float32
)

# Load LoRA adapter
model = PeftModel.from_pretrained(base_model, "lora_adapter", device_map="auto")

# Load pipeline
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Streamlit UI
st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
st.write("Ask me any Python programming question:")

user_input = st.text_input("Your question")

if user_input:
    if "python" in user_input.lower() or "list" in user_input.lower() or "tuple" in user_input.lower() or "def " in user_input.lower() or "class" in user_input.lower():
        prompt = f"""You are a helpful and friendly Python tutor. Only answer Python programming questions. Be clear and concise.

Question: {user_input}
Answer:"""
        response = pipe(prompt, max_new_tokens=256, temperature=0.7, do_sample=True)[0]["generated_text"]
        answer = response.split("Answer:")[-1].strip()
        st.markdown(f"💬 **Answer:**\n\n{answer}")
    else:
        st.warning("❌ Sorry, I can only answer Python programming questions.")