Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from gradio_client import Client, handle_file
+import seaborn as sns
 import matplotlib.pyplot as plt
 import os
 
@@ -39,12 +40,25 @@ def stream_chat_with_rag(
     
 
 def heatmap(top_n):
-    fig = client.predict(
+    pivot_table = client.predict(
     	top_n= top_n,
-		api_name="/heatmap"
+		api_name="/get_heatmap_pivot_table"
     )
-    print(type(fig))
-    return fig
+    
+    plt.figure(figsize=(10, 8))
+    sns.heatmap(pivot_table,
+                cmap='YlOrRd',
+                cbar_kws={'label': 'Weighted Frequency'},
+                square=True)
+    
+    plt.title(f'Top {top_n} Emotions vs Topics Weighted Frequency')
+    plt.xlabel('Topics')
+    plt.ylabel('Emotions')
+    plt.xticks(rotation=45, ha='right')
+    plt.tight_layout()
+    
+    return plt.gcf()
+
 
 # Create Gradio interface
 with gr.Blocks(title="Reddit Election Analysis") as demo: