Update app.py

app.py

@@ -6,30 +6,35 @@ import streamlit as st
 # Use CPU
 device = torch.device("cpu")
 
+# Set page config
+st.set_page_config(page_title="Python Tutor", page_icon="🐍")
+st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
+st.markdown("Ask me any *Python programming* question below:")
+
 # Load tokenizer and base model
-tokenizer = AutoTokenizer.from_pretrained("TinyLLaMA/TinyLLaMA-1.1B-Chat-v1.0")
-base_model = AutoModelForCausalLM.from_pretrained(
-    "TinyLLaMA/TinyLLaMA-1.1B-Chat-v1.0",
-    torch_dtype=torch.float32,
-    device_map=None
-).to(device)
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("TinyLLaMA/TinyLLaMA-1.1B-Chat-v1.0")
+    tokenizer.pad_token = tokenizer.eos_token if tokenizer.pad_token is None else tokenizer.pad_token
 
-# Load LoRA adapter
-model = PeftModel.from_pretrained(base_model, "lora_adapter").to(device)
-model.eval()
+    base_model = AutoModelForCausalLM.from_pretrained(
+        "TinyLLaMA/TinyLLaMA-1.1B-Chat-v1.0",
+        torch_dtype=torch.float32,
+    ).to(device)
 
-# Load zero-shot classifier for topic filtering
-classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli", device=-1)
+    model = PeftModel.from_pretrained(base_model, "lora_adapter").to(device)
+    model.eval()
 
-# Streamlit UI
-st.set_page_config(page_title="Python Tutor", page_icon="🐍")
-st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
-st.markdown("Ask me any *Python programming* question below:")
+    classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli", device=-1)
+
+    return tokenizer, model, classifier
 
-# Input box
+tokenizer, model, classifier = load_model()
+
+# User input
 user_question = st.text_input("Your question:")
 
-# Function to filter Python-related questions
+# Helper to check topic relevance
 def is_python_related(question):
     candidate_labels = [
         "Python programming", "Java programming", "General knowledge",
@@ -38,7 +43,7 @@ def is_python_related(question):
     result = classifier(question, candidate_labels)
     return result['labels'][0] == "Python programming"
 
-# Format code blocks
+# Code block formatter
 def format_code_blocks(text):
     if "```" in text:
         return text
@@ -58,7 +63,7 @@ def format_code_blocks(text):
         formatted.append("```")
     return "\n".join(formatted)
 
-# Handle input
+# Process the question
 if user_question:
     if len(user_question.strip()) < 10:
         st.warning("Please ask a more specific Python question.")
@@ -68,35 +73,29 @@ if user_question:
                 st.error("Sorry, I am a Python tutor. I cannot answer this.")
             else:
                 with st.spinner("Thinking..."):
-                    # Prompt to keep model focused
-                    prompt = f"""
-You are a helpful assistant that only answers questions about Python programming.
 
-Instructions:
-- If a question is not related to Python, do not answer it.
-- Focus strictly on Python concepts, syntax, libraries, and best practices.
-- Answer clearly and include code examples if helpful.
+                    prompt = f"""You are a helpful and knowledgeable Python programming tutor.
+Always provide short, clear, beginner-friendly explanations with examples.
 
 Question: {user_question}
 Answer:"""
 
-                    # Tokenize input
                     inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
-                    # Generate output
                     with torch.inference_mode():
                         output = model.generate(
-                        **inputs,
-                        max_new_tokens=512,  # increase if needed
-                        do_sample=False,
-                        repetition_penalty=1.1,
-                        eos_token_id=tokenizer.eos_token_id,
-                        pad_token_id=tokenizer.pad_token_id
-                    )
-                    # Decode and clean output
+                            **inputs,
+                            max_new_tokens=512,
+                            do_sample=False,
+                            repetition_penalty=1.1,
+                            eos_token_id=tokenizer.eos_token_id,
+                            pad_token_id=tokenizer.pad_token_id
+                        )
+
                     decoded = tokenizer.decode(output[0], skip_special_tokens=True)
-                    answer = decoded.split("Answer:")[-1].strip()
+                    answer = decoded.split("Answer:")[-1].split("Question:")[0].strip()  # clean hallucinated extra question
+                    formatted = format_code_blocks(answer)
 
-                    # Display response
                     st.markdown("### 💡 Answer:")
-                    st.markdown(format_code_blocks(answer))
+                    with st.expander("🔍 Click to view full response"):
+                        st.markdown(formatted)