Update app.py

app.py

@@ -1,43 +1,61 @@
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 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 tokenizer and model (CPU)
+base_model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+adapter_path = "lora_adapter"  # Your LoRA adapter folder path
 
-# Load base model
-base_model = AutoModelForCausalLM.from_pretrained(
-    "TinyLLaMA/TinyLLaMA-1.1B-Chat-v1.0",  # or your exact model name
-    device_map="auto"
-)
+# 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)
 
-# Load LoRA adapter
-model = PeftModel.from_pretrained(
-    base_model,
-    "lora_adapter",  # folder name
-    device_map="auto"
+# Streamlit UI setup
+st.set_page_config(
+    page_title="Python Tutor Chatbot",
+    page_icon="🐍",
+    layout="centered"
 )
 
-# Load pipeline
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+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}
 
-# Streamlit UI
-st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
-st.write("Ask me any Python programming question:")
+### Response:
+"""
 
-user_input = st.text_input("Your question")
+# Chat interface
+user_input = st.text_input("Your Python 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.")
+    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)