app.py

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

# Load tokenizer and model (on CPU)
base_model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
adapter_path = "lora_adapter"

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 with instruction to ignore unrelated queries
def create_prompt(user_input):
    return f"""You are a helpful and expert AI Python tutor.

Your job is to only answer questions strictly related to Python programming (syntax, concepts, libraries, frameworks, tools, errors, etc.).

If the question is unrelated to Python, politely respond:
"Sorry, I can only answer Python programming questions."

### Instruction:
{user_input}

### Response:
"""

# 🔎 User Input
user_input = st.text_input("Your Python Question:")

# 🔄 Inference
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():
            outputs = model.generate(
                **inputs,
                max_new_tokens=200,
                temperature=0.7,
                top_p=0.9,
                top_k=50,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )

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

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