Update app.py

app.py

@@ -3,54 +3,53 @@ 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"
+st.set_page_config(page_title="TinyLLaMA Python Tutor", layout="centered")
 
-tokenizer = AutoTokenizer.from_pretrained(base_model)
-
-# 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)")
+# Title
 st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
-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:
-    # 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 = (
-            "You are an expert Python tutor. Provide clear, concise, and accurate explanations with examples. "
-            "If the user's question is not related to Python programming, respond with: "
-            "'Sorry, I can only help with Python programming questions.'"
+st.write("Ask me any Python programming question:")
+
+# Load base model and LoRA adapter
+@st.cache_resource
+def load_model():
+    base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+    adapter_path = "lora_adapter"
+
+    tokenizer = AutoTokenizer.from_pretrained(base_model)
+    model = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=torch.float32)
+    model = PeftModel.from_pretrained(model, adapter_path)
+    model.eval()
+
+    return tokenizer, model
+
+tokenizer, model = load_model()
+
+# Prompt template
+def build_prompt(question):
+    return (
+        "You are a helpful AI tutor that only answers Python programming questions. "
+        "If the user asks something unrelated to Python, respond with: "
+        "'Sorry, I can only answer Python-related questions.'\n\n"
+        f"Question: {question}\nAnswer:"
+    )
+
+# Input box
+question = st.text_input("Your question")
+
+if question:
+    prompt = build_prompt(question)
+    inputs = tokenizer(prompt, return_tensors="pt")
+
+    with st.spinner("Thinking..."):
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=300,  # 🔼 Increased from 200 to 300
+            temperature=0.6,      # 🔽 More deterministic
+            top_p=0.9,
+            pad_token_id=tokenizer.eos_token_id
         )
-        prompt = f"<|system|>\n{system_prompt}</s>\n<|user|>\n{user_input}</s>\n<|assistant|>"
-
-        inputs = tokenizer(prompt, return_tensors="pt").to(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)
-            answer = decoded_output.split("<|assistant|>")[-1].strip()
-            st.success(f"💬 Answer:\n\n{answer}")
+
+    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    answer = decoded_output.split("Answer:")[-1].strip()
+
+    st.markdown(f"**💬 Answer:**\n\n{answer}")