Update app.py

app.py

@@ -12,7 +12,7 @@ base_model = AutoModelForCausalLM.from_pretrained(
     device_map="auto" if torch.cuda.is_available() else None
 )
 
-# Load LoRA adapter
+# Load LoRA Adapter
 model = PeftModel.from_pretrained(base_model, "lora_adapter")
 model.eval()
 
@@ -23,35 +23,36 @@ st.write("Ask me any **Python programming** question:")
 user_input = st.text_input("Your question")
 
 if user_input:
-    # Prompt template that helps model decide to answer or reject
-    prompt = f"""
-You are a helpful and expert Python programming tutor.
-Answer only questions related to Python programming.
-If the question is unrelated to Python (like history, math, etc), politely respond:
+    # 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}
+Question: {user_input}
+Answer:"""
 
-### Answer:
-"""
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-    inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=True).to(model.device)
     with torch.no_grad():
-        output = model.generate(
+        outputs = model.generate(
             **inputs,
-            max_new_tokens=200,
+            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 = tokenizer.decode(output[0], skip_special_tokens=True)
+    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Clean the output to only show the answer
-    if "### Answer:" in decoded:
-        final_answer = decoded.split("### Answer:")[-1].strip()
+    # Extract answer after 'Answer:' line
+    answer_start = decoded_output.find("Answer:")
+    if answer_start != -1:
+        final_answer = decoded_output[answer_start + len("Answer:"):].strip()
     else:
-        final_answer = decoded.strip()
+        final_answer = decoded_output.strip()
 
     st.markdown(f"**Answer:** {final_answer}")