Update app.py

app.py

@@ -65,10 +65,17 @@ llm = HuggingFacePipeline(pipeline=pipe)
 # --- Step 4: Setup memory and QA chain ---
 memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 
-prompt = PromptTemplate.from_template(
-    """
-You are a helpful assistant at the University of Hertfordshire. Use only the context below to answer clearly and factually.
-If the answer is not in the context, say you don't know.
+prompt = PromptTemplate.from_template("""
+You are a helpful assistant at the University of Hertfordshire.
+Use the provided context to create a **clear, concise, step-by-step answer** in friendly, student-friendly language.
+Do not copy the context verbatim—paraphrase where possible.  
+Remove any irrelevant details.  
+If the answer is not in the context, reply: "I don't know."
+
+Format the answer like this:
+1) Step one...
+2) Step two...
+(Use numbered steps where possible.)
 
 Context:
 {context}
@@ -76,9 +83,19 @@ Context:
 Question:
 {question}
 
-Helpful answer:
-""")
+Final refined answer:
+""".strip()
+)
 
+def refine_answer(raw_answer):
+    # Remove extra markers
+    text = raw_answer.strip()
+    text = text.replace("Helpful answer:", "").strip()
+    # Capitalise first letter if missing
+    if text and not text[0].isupper():
+        text = text[0].upper() + text[1:]
+    return text
+    
 qa_chain = ConversationalRetrievalChain.from_llm(
     llm=llm,
     retriever=vectordb.as_retriever(search_kwargs={"k": 3}),
@@ -90,11 +107,21 @@ qa_chain = ConversationalRetrievalChain.from_llm(
 UH_LOGO = "images/UH.png"
 
 # --- Step 5: Define chatbot logic ---
+
+def refine_answer(raw_answer: str) -> str:
+    """Clean and polish raw model output."""
+    text = raw_answer.strip()
+    # Remove prompt artifacts
+    for marker in ["Helpful answer:", "<|assistant|>", "Refined helpful answer:"]:
+        text = text.replace(marker, "")
+    # Normalise spaces
+    text = " ".join(text.split())
+    return text
+    
 def chat(message, history):
     result = qa_chain.invoke({"question": message})
-    # ConversationalRetrievalChain returns {"answer": "...", "source_documents": ..., "chat_history": ...}
-    response = result.get("answer", "")
-    response = response.split("Answer: ")[-1].replace("<|assistant|>", "")
+    response = refine_answer(result.get("answer", ""))
+    
     # keep GPU clean between turns (helps on Spaces)
     if torch.cuda.is_available():
         torch.cuda.empty_cache()