Update app.py

app.py

@@ -1,24 +1,23 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 import streamlit as st
 
 # Load tokenizer and base model
 base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-lora_path = "./lora_adapter"  # make sure your LoRA adapter folder is named like this
-
-bnb_config = BitsAndBytesConfig(load_in_4bit=True,
-                                bnb_4bit_compute_dtype=torch.bfloat16)
+lora_path = "./lora_adapter"
 
 tokenizer = AutoTokenizer.from_pretrained(base_model)
-model = AutoModelForCausalLM.from_pretrained(base_model,
-                                             quantization_config=bnb_config,
-                                             torch_dtype=torch.bfloat16,
-                                             device_map="auto")
 
+# Load base model normally (for CPU)
+model = AutoModelForCausalLM.from_pretrained(base_model)
 model = PeftModel.from_pretrained(model, lora_path)
 model.eval()
 
+# Move to CPU explicitly
+device = torch.device("cpu")
+model.to(device)
+
 # Streamlit UI
 st.set_page_config(page_title="🧠 TinyLLaMA Python Tutor (LoRA)")
 st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
@@ -27,11 +26,10 @@ st.write("Ask me any **Python programming** question:")
 user_input = st.text_input("Your question", placeholder="e.g. What is a lambda function in Python?")
 
 if user_input:
-    # Filtering logic: Only answer Python-related queries
+    # Check if it's a Python-related question
     if "python" not in user_input.lower() and "py" not in user_input.lower():
         st.warning("❌ Sorry, I can only answer Python programming questions.")
     else:
-        # System prompt for tutor behavior
         system_prompt = (
             "You are a helpful and knowledgeable Python tutor. "
             "Answer the user's Python programming questions clearly and concisely. "
@@ -39,7 +37,7 @@ if user_input:
         )
         prompt = f"<|system|>\n{system_prompt}</s>\n<|user|>\n{user_input}</s>\n<|assistant|>"
 
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
         with torch.no_grad():
             with st.spinner("Thinking..."):
@@ -52,8 +50,7 @@ if user_input:
                     eos_token_id=tokenizer.eos_token_id,
                     pad_token_id=tokenizer.eos_token_id
                 )
-            decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-            # Extract answer only (remove prompt)
+            decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
             answer = decoded_output.split("<|assistant|>")[-1].strip()
             st.success(f"💬 Answer:\n\n{answer}")