Update src/streamlit_app.py

src/streamlit_app.py

@@ -39,18 +39,31 @@ if st.button("Ask DocMed"):
     if question.strip() == "":
         st.warning("Please enter a question.")
     else:
-        prompt = f"Explain simply for a medical student: {question}"
+        prompt = f"""You are DocMed, a medical study assistant.
+Explain the following clearly and concisely for a medical student.
+Give the answer directly.
 
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+Question:
+{question}
 
-        with torch.no_grad():
-            output = model.generate(
-                **inputs,
-                max_new_tokens=120,
-                temperature=0.7,
-                top_p=0.9
-            )
+Answer:
+"""
+
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+with torch.no_grad():
+    output = model.generate(
+        **inputs,
+        max_new_tokens=120,
+        temperature=0.6,
+        top_p=0.9,
+        repetition_penalty=1.2,
+        do_sample=True
+    )
+
+decoded = tokenizer.decode(output[0], skip_special_tokens=True)
 
-        answer = tokenizer.decode(output[0], skip_special_tokens=True)
+# Remove prompt echo if present
+answer = decoded.split("Answer:")[-1].strip()
         st.markdown("### 🧠 DocMed says:")
         st.write(answer)