Update app.py

app.py

@@ -28,7 +28,7 @@ if not GEMINI_API_KEY:
 genai.configure(api_key=GEMINI_API_KEY)
 model = genai.GenerativeModel('gemini-2.0-flash')
 
-# Preprocess text function
+# Preprocess text
 def preprocess_text(text):
     text = text.lower()
     text = text.replace('\n', ' ').replace('\t', ' ')
@@ -36,7 +36,7 @@ def preprocess_text(text):
     text = ' '.join(text.split()).strip()
     return text
 
-# Retrieve documents
+# Retrieve top-k documents
 def retrieve_docs(query, k=5):
     query_embedding = minilm.encode([query], show_progress_bar=False)[0].astype(np.float32)
     distances, indices = index.search(np.array([query_embedding]), k)
@@ -44,24 +44,24 @@ def retrieve_docs(query, k=5):
     retrieved_docs['distance'] = distances[0]
     return retrieved_docs
 
-# Simplified respond function (no history)
+# Generate structured response
 def respond(message, system_message, max_tokens, temperature, top_p):
-    # Preprocess the user message
+    # Preprocess and retrieve
     preprocessed_query = preprocess_text(message)
-    
-    # Retrieve relevant documents
     retrieved_docs = retrieve_docs(preprocessed_query, k=5)
-    context = "\n".join(retrieved_docs['text'].tolist())
-    
-    # Construct the prompt with system message and RAG context
+
+    # Combine retrieved texts
+    context = "\n".join([f"- *{row['label']}* ({row['source']}): {row['text']}" for _, row in retrieved_docs.iterrows()])
+
+    # Build prompt
     prompt = f"{system_message}\n\n"
     prompt += (
         f"Query: {message}\n"
         f"Relevant Context: {context}\n"
         f"Generate a short, concise, and to-the-point response to the query based only on the provided context."
     )
-    
-    # Generate response with Gemini
+
+    # Get Gemini response
     response = model.generate_content(
         prompt,
         generation_config=genai.types.GenerationConfig(
@@ -76,10 +76,26 @@ def respond(message, system_message, max_tokens, temperature, top_p):
             answer = answer[:last_period + 1]
         else:
             answer += "."
-    
-    return answer
 
-# Simple Gradio Interface
+    # Format output with Markdown
+    formatted_answer = f"""
+**🩺 Patient Query:**  
+{message}
+
+---
+
+**📚 Retrieved Context:**  
+{context}
+
+---
+
+**🧠 Diagnosis / Suggestion:**  
+{answer}
+"""
+
+    return formatted_answer.strip()
+
+# Gradio app
 demo = gr.Interface(
     fn=respond,
     inputs=[
@@ -90,18 +106,12 @@ demo = gr.Interface(
         ),
         gr.Slider(minimum=1, maximum=2048, value=150, step=1, label="Max New Tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.75, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",  # Included but not used by Gemini
-        ),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
     ],
-    outputs=gr.Textbox(label="Diagnosis"),
+    outputs=gr.Markdown(label="Diagnosis"),
     title="🏥 Medical Assistant",
     description="A simple medical assistant that diagnoses patient queries using AI and past records."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()