Update app.py

app.py

@@ -7,7 +7,7 @@ import os
 HF_TOKEN = os.getenv("HF_TOKEN")  # Replace with your actual token if not using an environment variable
 
 # Initialize the Gradio Client for the specified API
-client = Client("mangoesai/Elections_Comparison_Agent_V3", hf_token=HF_TOKEN)
+client = Client("mangoesai/Elections_Comparison_Agent_V4", hf_token=HF_TOKEN)
 
 client_name = ['2016 Election','2024 Election', 'Comparison two years']
 
@@ -15,7 +15,6 @@ client_name = ['2016 Election','2024 Election', 'Comparison two years']
 
 def stream_chat_with_rag(
     message: str,
-    # history: list,
     client_name: str
 ):
     # print(f"Message: {message}")
@@ -33,11 +32,36 @@ def stream_chat_with_rag(
 
     return answer, fig
 
+def heatmap(top_n):
+    map = client.predict(
+        top_n = top_n,
+        api_name = "/heatmap"
+    )
+
+    return map
+
 
 
 # Create Gradio interface
 with gr.Blocks(title="Reddit Election Analysis") as demo:
     gr.Markdown("# Reddit Public sentiment & Social topic distribution ")
+    with gr.Row():
+        with gr.Column():
+            with gr.Row():
+                top_n = gr.Dropdown(choices= [1,2,3,4,5,6,7,8,9,10])
+            with gr.Row():
+                fresh_btn = gr.Button("Refresh Heatmap")
+        with gr.Column():
+            output_heatmap = gr.Plot(
+                label="Top Public sentiment & Social topic Heatmap",
+                container=True,  # Ensures the plot is contained within its area
+                elem_classes="heatmap-plot"  # Add a custom class for styling
+            )
+    fresh_btn.click(
+        fn=heatmap,
+        inputs=top_n,
+        outputs=output_heatmap
+    )
 
 
 
@@ -54,7 +78,8 @@ with gr.Blocks(title="Reddit Election Analysis") as demo:
 
             query_input = gr.Textbox(
                 label="Your Question",
-                placeholder="Ask about election comments or posts..."
+                placeholder="Ask about election comments or posts, like, Is there any comments don't like the election results?",
+                value = "Is there any comments don't like the election results"
             )
 
             submit_btn = gr.Button("Submit")
@@ -86,18 +111,22 @@ with gr.Blocks(title="Reddit Election Analysis") as demo:
                 width: 100%;
                 margin: auto;
             }
+            .heatmap-plot {
+                min-height: 400px;
+                width: 100%;
+                margin: auto;
+            }
         </style>
     """)
 
     # Update both outputs when submit is clicked
-    submit_btn.click(
-        fn=stream_chat_with_rag,
+    submit_btn.stream_chat_with_rag(
+        fn=process_query,
         inputs=[query_input, year_selector],
         outputs=[output_text, output_plot]
     )
 
 
-
 if __name__ == "__main__":
     demo.launch(share=True)