Update app.py

app.py

@@ -17,7 +17,9 @@ embedding_arr = model.encode(df['symptoms']).astype(np.float32)
 # ------ Clustering Setup ------
 kmeans = MiniBatchKMeans(n_clusters=10, random_state=42)
 cluster_labels = kmeans.fit_predict(embedding_arr)
-cluster_prognosis_map = df.groupby(cluster_labels)['prognosis'].agg(lambda x: x.mode().tolist())
+cluster_prognosis_map = df.groupby(cluster_labels)['prognosis'].agg(
+    lambda x: x.value_counts(normalize=True).head(3).to_dict()
+)
 
 # ------ PCA Initialization ------
 pca = PCA(n_components=2).fit(embedding_arr)
@@ -31,7 +33,7 @@ def cached_encode(text):
 client = AsyncInferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
 # ------ Streaming Response Function ------
-async def respond(message, history, system_message, max_tokens, temperature, top_p):
+async def respond(message, history, system_message):
     try:
         # Encoding and clustering
         query_embedding = cached_encode(message)
@@ -46,10 +48,10 @@ async def respond(message, history, system_message, max_tokens, temperature, top
                 "role": "user",
                 "content": message
             }],
-            max_tokens=max_tokens,
+            max_tokens=512,
             stream=True,
-            temperature=temperature,
-            top_p=top_p
+            temperature=0.7,
+            top_p=0.95
         )
         
         full_response = ""
@@ -57,36 +59,25 @@ async def respond(message, history, system_message, max_tokens, temperature, top
             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])}"
+        # Format prognosis likelihoods
+        prognosis_info = cluster_prognosis_map[query_cluster]
+        formatted_prognoses = ", ".join([f"{k}: {v*100:.1f}%" for k, v in prognosis_info.items()])
+        
+        # Append formatted prognosis to response
+        yield f"{full_response}\n\nCommon prognoses: {formatted_prognoses}"
         
     except Exception as e:
         yield f"Error: {str(e)}"
 
-# ------ Gradio Interface with Example Query Button ------
-with gr.Blocks() as demo:
-    with gr.Row():
-        chatbot = 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")
-            ]
-        )
-    with gr.Row():
-        example_button = gr.Button("Use Example Query")
-        example_query = gr.Textbox(
-            value="I have a headache and fever",
-            label="Example Query",
-            interactive=False
-        )
-    
-    # Trigger example query manually when button is clicked
-    example_button.click(fn=lambda: "I have a headache and fever", outputs=[chatbot])
+# ------ Gradio Interface ------
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Textbox(value="Medical diagnosis assistant", label="System Role")
+    ],
+    examples=[["I have a headache and fever"]]  # Example query for guidance
+)
 
 if __name__ == "__main__":
     demo.launch()