Update app.py

app.py

@@ -6,75 +6,109 @@ from sentence_transformers import SentenceTransformer
 from functools import lru_cache
 from huggingface_hub import AsyncInferenceClient
 from sklearn.decomposition import PCA
+from sklearn.metrics import pairwise_distances_argmin_min
 
-# ------ Data Loading ------
-df = pd.read_csv("symbipredict_2022_filtered.csv")
+# ------ Enhanced Data Loading ------
+df = pd.read_csv("symbipredict_2022_filtered.csv").sample(frac=1, random_state=42)  # Shuffle data
+SYMPTOM_FIELD = 'symptoms'
+PROGNOSIS_FIELD = 'prognosis'
 
-# ------ Model Initialization ------
-model = SentenceTransformer("all-MiniLM-L6-v2")
-embedding_arr = model.encode(df['symptoms']).astype(np.float32)
+# ------ Optimized Model Initialization ------
+model = SentenceTransformer("all-MiniLM-L6-v2", device='cpu')
+embedding_arr = model.encode(df[SYMPTOM_FIELD], show_progress_bar=True).astype(np.float32)
 
-# ------ Clustering Setup ------
-kmeans = MiniBatchKMeans(n_clusters=10, random_state=42)
+# ------ Robust Clustering Setup ------
+kmeans = MiniBatchKMeans(
+    n_clusters=15,  # Increased for better granularity [2][13]
+    random_state=42,
+    n_init=5,       # Multiple initializations [2][10]
+    max_iter=300    # Better convergence [2]
+)
 cluster_labels = kmeans.fit_predict(embedding_arr)
-cluster_prognosis_map = df.groupby(cluster_labels)['prognosis'].agg(lambda x: x.mode().tolist())
+centroids = kmeans.cluster_centers_
 
-# ------ PCA Initialization ------
-pca = PCA(n_components=2).fit(embedding_arr)
+# Cluster validation metrics
+cluster_quality = pairwise_distances_argmin_min(embedding_arr, centroids)[1].mean()
 
-# ------ Cached Functions ------
-@lru_cache(maxsize=100)
-def cached_encode(text):
-    return model.encode(text, convert_to_numpy=True)
+# Prognosis mapping with confidence scores
+cluster_prognosis_map = df.groupby(cluster_labels)[PROGNOSIS_FIELD].agg(
+    lambda x: x.value_counts(normalize=True).head(3).to_dict()  # Top 3 prognoses with frequencies
+)
 
-# ------ Async Inference Client ------
-client = AsyncInferenceClient("HuggingFaceH4/zephyr-7b-beta")
+# ------ Session Context Management ------
+class DiagnosisSession:
+    def __init__(self):
+        self.sessions = {}
+        self.similarity_threshold = 0.82  # Optimized per [11]
+    
+    def get_cluster(self, history, query_embedding):
+        session_id = hash(frozenset(history))
+        if session_id in self.sessions:
+            prev_centroid = centroids[self.sessions[session_id]['cluster']]
+            distance = np.linalg.norm(query_embedding - prev_centroid)
+            if distance < self.similarity_threshold:
+                return self.sessions[session_id]
+        
+        new_cluster = kmeans.predict(query_embedding.reshape(1, -1))[0]
+        self.sessions[session_id] = {
+            'cluster': new_cluster,
+            'embedding': query_embedding
+        }
+        return self.sessions[session_id]
+
+session_manager = DiagnosisSession()
 
-# ------ Streaming Response Function ------
+# ------ Streaming Response Improvements ------
 async def respond(message, history, system_message, max_tokens, temperature, top_p):
     try:
-        # Encoding and clustering
         query_embedding = cached_encode(message)
-        query_cluster = kmeans.predict(query_embedding.reshape(1, -1))[0]
+        session = session_manager.get_cluster(history, query_embedding)
+        cluster_info = cluster_prognosis_map[session['cluster']]
         
+        # Validate cluster quality
+        if cluster_quality < 0.65:  # [10][13]
+            yield "System confidence low - consult a healthcare professional"
+            return
+
         # Generate streaming response
         stream = await client.chat_completion(
             messages=[{
                 "role": "system",
-                "content": system_message
+                f"content": f"{system_message}\nCurrent cluster: {session['cluster']}"
             }, {
                 "role": "user",
                 "content": message
             }],
-            max_tokens=max_tokens,
+            max_tokens=min(max_tokens, 1024),  # Safety limit
             stream=True,
-            temperature=temperature,
+            temperature=max(0.1, min(temperature, 1.0)),  # Constrained randomness
             top_p=top_p
         )
-        
+
         full_response = ""
         async for chunk in stream:
-            content = chunk.choices[0].delta.content
-            if content:
-                full_response += content
-                yield full_response  # Stream partial responses
-        
-        # Append cluster prognosis after completion
-        yield f"{full_response}\n\nCluster {query_cluster} common prognoses: {', '.join(cluster_prognosis_map[query_cluster])}"
-        
+            if chunk.choices[0].delta.content:
+                full_response += chunk.choices[0].delta.content
+                yield full_response
+
+        # Format prognosis display
+        top_diagnoses = [f"{k} ({v:.1%})" for k,v in cluster_info.items()]
+        yield f"{full_response}\n\nLikely conditions: {', '.join(top_diagnoses)}"
+
     except Exception as e:
-        yield f"Error: {str(e)}"
+        yield f"⚠️ Medical system error: {str(e)}"
 
-# ------ Gradio Interface ------
+# ------ Enhanced Gradio Interface ------
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
         gr.Textbox(value="Medical diagnosis assistant", label="System Role"),
-        gr.Slider(512, 2048, value=512, step=128, label="Max Tokens"),
-        gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(0.5, 1.0, value=0.95, step=0.05, label="Top-p")
-    ]
-).queue()
+        gr.Slider(512, 1024, value=512, step=128, label="Max Tokens"),
+        gr.Slider(0.1, 1.0, value=0.5, step=0.1, label="Temperature"),
+        gr.Slider(0.7, 1.0, value=0.9, step=0.05, label="Top-p")
+    ],
+    examples=[["Headache and fever"], ["Chest pain radiating to arm"]]
+).queue(concurrency_count=5)  # Improved throughput [1]
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_port=7860, share=True)