Update app.py

app.py

@@ -1,58 +1,59 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from peft import PeftModel
 import streamlit as st
 
 # Load tokenizer and base model
-base_model_path = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-tokenizer = AutoTokenizer.from_pretrained(base_model_path)
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_path,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto" if torch.cuda.is_available() else None
-)
-
-# Load LoRA Adapter
-model = PeftModel.from_pretrained(base_model, "lora_adapter")
+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)
+
+tokenizer = AutoTokenizer.from_pretrained(base_model)
+model = AutoModelForCausalLM.from_pretrained(base_model,
+                                             quantization_config=bnb_config,
+                                             torch_dtype=torch.bfloat16,
+                                             device_map="auto")
+
+model = PeftModel.from_pretrained(model, lora_path)
 model.eval()
 
 # Streamlit UI
+st.set_page_config(page_title="🧠 TinyLLaMA Python Tutor (LoRA)")
 st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
 st.write("Ask me any **Python programming** question:")
 
-user_input = st.text_input("Your question")
+user_input = st.text_input("Your question", placeholder="e.g. What is a lambda function in Python?")
 
 if user_input:
-    # Better Prompt Template
-    prompt = f"""You are a helpful Python programming tutor.
-
-You will ONLY answer questions related to Python programming.
-If the question is unrelated to Python, reply:
-"Sorry, I can only answer Python-related questions."
-
-Question: {user_input}
-Answer:"""
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=300,
-            temperature=0.7,
-            do_sample=True,
-            top_p=0.95,
-            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 after 'Answer:' line
-    answer_start = decoded_output.find("Answer:")
-    if answer_start != -1:
-        final_answer = decoded_output[answer_start + len("Answer:"):].strip()
+    # Filtering logic: Only answer Python-related queries
+    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:
-        final_answer = decoded_output.strip()
-
-    st.markdown(f"**Answer:** {final_answer}")
+        # 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. "
+            "If the question is unclear, ask for clarification."
+        )
+        prompt = f"<|system|>\n{system_prompt}</s>\n<|user|>\n{user_input}</s>\n<|assistant|>"
+
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+        with torch.no_grad():
+            with st.spinner("Thinking..."):
+                outputs = model.generate(
+                    **inputs,
+                    max_new_tokens=150,
+                    temperature=0.7,
+                    top_p=0.95,
+                    do_sample=True,
+                    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)
+            answer = decoded_output.split("<|assistant|>")[-1].strip()
+            st.success(f"💬 Answer:\n\n{answer}")