Update app.py

app.py

@@ -9,7 +9,7 @@ import os
 
 # Load data and FAISS index
 def load_data_and_index():
-    docs_df = pd.read_pickle("data.pkl")  # Adjust path for HF Spaces
+    docs_df = pd.read_pickle("docs_with_embeddings (1).pkl")  # Adjust path for HF Spaces
     embeddings = np.array(docs_df['embeddings'].tolist(), dtype=np.float32)
     dimension = embeddings.shape[1]
     index = faiss.IndexFlatL2(dimension)
@@ -44,7 +44,38 @@ def retrieve_docs(query, k=5):
     retrieved_docs['distance'] = distances[0]
     return retrieved_docs
 
-# Respond function with HTML formatting
+# Parse response into structured sections
+def parse_response(response_text):
+    sections = {
+        "Symptoms": [],
+        "Signs": [],
+        "Risk Factors": [],
+        "Diagnostic Criteria": [],
+        "Other": []
+    }
+    
+    # Simple regex-based parsing (adjust based on your Gemini output format)
+    lines = response_text.split('\n')
+    current_section = "Other"
+    
+    for line in lines:
+        line = line.strip()
+        if line.lower().startswith("symptoms:"):
+            current_section = "Symptoms"
+        elif line.lower().startswith("signs:"):
+            current_section = "Signs"
+        elif line.lower().startswith("risk factors") or line.lower().startswith("past medical history:"):
+            current_section = "Risk Factors"
+        elif line.lower().startswith("diagnostic criteria:"):
+            current_section = "Diagnostic Criteria"
+        elif line and not line.startswith((' ', '\t')) and ':' in line:
+            current_section = "Other"
+        if line and not line.endswith(':'):
+            sections[current_section].append(line)
+    
+    return sections
+
+# Respond function with generic HTML formatting
 def respond(message, system_message, max_tokens, temperature, top_p):
     # Preprocess the user message
     preprocessed_query = preprocess_text(message)
@@ -58,7 +89,8 @@ def respond(message, system_message, max_tokens, temperature, top_p):
     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. Format the response with clear sections like Symptoms, Signs, Risk Factors, and Diagnostic Criteria where applicable."
+        f"Generate a concise response to the query based only on the provided context. "
+        f"Structure the response with clear sections like 'Symptoms:', 'Signs:', 'Risk Factors:', and 'Diagnostic Criteria:' where applicable."
     )
     
     # Generate response with Gemini
@@ -71,48 +103,34 @@ def respond(message, system_message, max_tokens, temperature, top_p):
     )
     answer = response.text.strip()
 
+    # Parse the response into sections
+    sections = parse_response(answer)
+
     # Format the response into HTML with CSS styling
     html_response = """
     <style>
-        .diagnosis-container { font-family: Arial, sans-serif; line-height: 1.6; padding: 10px; }
-        h2 { color: #2c3e50; font-size: 20px; margin-bottom: 10px; }
-        h3 { color: #2980b9; font-size: 16px; margin-top: 15px; margin-bottom: 5px; }
-        ul { margin: 0; padding-left: 20px; }
-        li { margin-bottom: 5px; }
-        p { margin: 5px 0; }
+        .diagnosis-container { font-family: Arial, sans-serif; line-height: 1.6; padding: 15px; max-width: 800px; margin: auto; }
+        h2 { color: #2c3e50; font-size: 24px; margin-bottom: 15px; border-bottom: 2px solid #2980b9; padding-bottom: 5px; }
+        h3 { color: #2980b9; font-size: 18px; margin-top: 15px; margin-bottom: 8px; }
+        ul { margin: 0; padding-left: 25px; }
+        li { margin-bottom: 6px; color: #34495e; }
+        p { margin: 5px 0; color: #34495e; }
     </style>
     <div class="diagnosis-container">
-        <h2>Diagnosis</h2>
+        <h2>AI Response</h2>
+        <p>Based on the provided context, here is the information relevant to your query:</p>
     """
 
-    # Parse the response and structure it (this is a simple example; adjust based on actual output)
-    if "heart failure" in message.lower():
-        html_response += """
-        <p>Based on the provided context, the following information supports the query "heart failure":</p>
-        <h3>Symptoms</h3>
-        <ul>
-            <li>Breathlessness (dyspnea on exertion, progressive SOB)</li>
-            <li>Reduced exercise tolerance</li>
-            <li>Ankle swelling (edema in legs)</li>
-        </ul>
-        <h3>Signs</h3>
-        <ul>
-            <li>Elevated jugular venous pressure (markedly elevated JVP)</li>
-        </ul>
-        <h3>Risk Factors/Past Medical History</h3>
-        <ul>
-            <li>Coronary artery disease (CAD s/p CABG)</li>
-            <li>Arrhythmias (Paroxysmal atrial fibrillation)</li>
-            <li>Hypertension</li>
-        </ul>
-        <h3>Diagnostic Criteria</h3>
-        <ul>
-            <li>Elevated BNP</li>
-        </ul>
-        """
-    else:
-        # Fallback for other queries
-        html_response += f"<p>{answer}</p>"
+    # Add sections dynamically
+    for section, items in sections.items():
+        if items:  # Only include sections that have content
+            html_response += f"<h3>{section}</h3>"
+            html_response += "<ul>"
+            for item in items:
+                # Remove section prefix if present (e.g., "Symptoms:" from the first line)
+                cleaned_item = re.sub(rf"^{section}:", "", item, flags=re.IGNORECASE).strip()
+                html_response += f"<li>{cleaned_item}</li>"
+            html_response += "</ul>"
 
     html_response += "</div>"
     return html_response
@@ -121,7 +139,7 @@ def respond(message, system_message, max_tokens, temperature, top_p):
 demo = gr.Interface(
     fn=respond,
     inputs=[
-        gr.Textbox(label="Your Query", placeholder="Enter your medical question here (e.g., heart failure)..."),
+        gr.Textbox(label="Your Query", placeholder="Enter your medical question here (e.g., diabetes, heart failure)..."),
         gr.Textbox(
             value="You are a medical AI assistant diagnosing patients based on their query, using relevant context from past records of other patients.",
             label="System Message"
@@ -137,8 +155,8 @@ demo = gr.Interface(
         ),
     ],
     outputs=gr.HTML(label="Diagnosis"),
-    title="🏥 Medical Assistant",
-    description="A simple medical assistant that diagnoses patient queries using AI and past records, with styled output."
+    title="🏥 Medical Query Assistant",
+    description="A medical assistant that diagnoses patient queries using AI and past records, with styled output for any condition."
 )
 
 if __name__ == "__main__":