add vis

app.py

@@ -52,11 +52,18 @@ def filter_and_display(selected_columns, model_types, datasets, stage):
     display_columns = ["Method", "Model"] + [col for col in selected_columns if col in filtered.columns]
     return filtered[display_columns] if not filtered.empty else pd.DataFrame(columns=display_columns)
 
-def create_visualization(filtered_data):
+def create_prefill_visualization(filtered_data):
     if filtered_data.empty:
         return None
     fig = px.bar(filtered_data, x='Model', y='Quality', color='Method', barmode='group',
-                 title='Quality by Model and Method')
+                 title='Prefill Stage: Quality by Model and Method')
+    return fig
+
+def create_decode_visualization(filtered_data):
+    if filtered_data.empty:
+        return None
+    fig = px.bar(filtered_data, x='Model', y='Throughput (token/s)', color='Method', barmode='group',
+                 title='Decode Stage: Throughput by Model and Method')
     return fig
 
 # Load the data from the /data folder
@@ -109,6 +116,11 @@ This demo leaderboard allows users to explore and compare different KV cache imp
                 
                 prefill_results = gr.Dataframe(value=filter_and_display(["Quality", "TTFT"], list(data["Model"].unique()), list(data["Dataset"].unique()), "prefill"), headers=["Method", "Model", "Quality", "TTFT", "Link"])
 
+                # Prefill-stage visualization
+                with gr.Row():
+                    gr.Markdown("### Prefill-stage Visualization")
+                    prefill_plot = gr.Plot(value=create_prefill_visualization(filter_and_display(["Quality"], list(data["Model"].unique()), list(data["Dataset"].unique()), "prefill")))
+
                 # Decode-stage selection
                 with gr.Row():
                     gr.Markdown("## Decode-stage Selection")
@@ -143,27 +155,30 @@ This demo leaderboard allows users to explore and compare different KV cache imp
                 
                 decode_results = gr.Dataframe(value=filter_and_display(["Throughput (token/s)", "Quality"], list(data["Model"].unique()), list(data["Dataset"].unique()), "decode"), headers=["Method", "Model", "Throughput (token/s)", "Quality", "Link"])
 
+                # Decode-stage visualization
                 with gr.Row():
-                    gr.Markdown("### Visualization")
-                    plot = gr.Plot(value=create_visualization(data))
+                    gr.Markdown("### Decode-stage Visualization")
+                    decode_plot = gr.Plot(value=create_decode_visualization(filter_and_display(["Throughput (token/s)"], list(data["Model"].unique()), list(data["Dataset"].unique()), "decode")))
 
                 def auto_update_prefill(selected_columns, model_types, datasets):
                     if not model_types or not datasets:
-                        return pd.DataFrame(columns=["Method", "Model"] + selected_columns)
-                    return filter_and_display(selected_columns, model_types, datasets, "prefill")
+                        return pd.DataFrame(columns=["Method", "Model"] + selected_columns), None
+                    filtered_data = filter_and_display(selected_columns, model_types, datasets, "prefill")
+                    return filtered_data, create_prefill_visualization(filtered_data)
 
                 def auto_update_decode(selected_columns, model_types, datasets):
                     if not model_types or not datasets:
-                        return pd.DataFrame(columns=["Method", "Model"] + selected_columns)
-                    return filter_and_display(selected_columns, model_types, datasets, "decode")
+                        return pd.DataFrame(columns=["Method", "Model"] + selected_columns), None
+                    filtered_data = filter_and_display(selected_columns, model_types, datasets, "decode")
+                    return filtered_data, create_decode_visualization(filtered_data)
 
-                prefill_columns_to_display.change(auto_update_prefill, inputs=[prefill_columns_to_display, prefill_model_types, prefill_datasets], outputs=[prefill_results])
-                prefill_model_types.change(auto_update_prefill, inputs=[prefill_columns_to_display, prefill_model_types, prefill_datasets], outputs=[prefill_results])
-                prefill_datasets.change(auto_update_prefill, inputs=[prefill_columns_to_display, prefill_model_types, prefill_datasets], outputs=[prefill_results])
+                prefill_columns_to_display.change(auto_update_prefill, inputs=[prefill_columns_to_display, prefill_model_types, prefill_datasets], outputs=[prefill_results, prefill_plot])
+                prefill_model_types.change(auto_update_prefill, inputs=[prefill_columns_to_display, prefill_model_types, prefill_datasets], outputs=[prefill_results, prefill_plot])
+                prefill_datasets.change(auto_update_prefill, inputs=[prefill_columns_to_display, prefill_model_types, prefill_datasets], outputs=[prefill_results, prefill_plot])
 
-                decode_columns_to_display.change(auto_update_decode, inputs=[decode_columns_to_display, decode_model_types, decode_datasets], outputs=[decode_results])
-                decode_model_types.change(auto_update_decode, inputs=[decode_columns_to_display, decode_model_types, decode_datasets], outputs=[decode_results])
-                decode_datasets.change(auto_update_decode, inputs=[decode_columns_to_display, decode_model_types, decode_datasets], outputs=[decode_results])
+                decode_columns_to_display.change(auto_update_decode, inputs=[decode_columns_to_display, decode_model_types, decode_datasets], outputs=[decode_results, decode_plot])
+                decode_model_types.change(auto_update_decode, inputs=[decode_columns_to_display, decode_model_types, decode_datasets], outputs=[decode_results, decode_plot])
+                decode_datasets.change(auto_update_decode, inputs=[decode_columns_to_display, decode_model_types, decode_datasets], outputs=[decode_results, decode_plot])
 
                 # Add a reload button
                 def reload_data():
@@ -171,10 +186,11 @@ This demo leaderboard allows users to explore and compare different KV cache imp
                     data = load_data(data_dir)
                     return (filter_and_display(prefill_columns_to_display.value, prefill_model_types.value, prefill_datasets.value, "prefill"),
                             filter_and_display(decode_columns_to_display.value, decode_model_types.value, decode_datasets.value, "decode"),
-                            create_visualization(data))
+                            create_prefill_visualization(filter_and_display(prefill_columns_to_display.value, prefill_model_types.value, prefill_datasets.value, "prefill")),
+                            create_decode_visualization(filter_and_display(decode_columns_to_display.value, decode_model_types.value, decode_datasets.value, "decode")))
 
                 reload_button = gr.Button("Reload Data")
-                reload_button.click(reload_data, outputs=[prefill_results, decode_results, plot])
+                reload_button.click(reload_data, outputs=[prefill_results, decode_results, prefill_plot, decode_plot])
 
             with gr.TabItem("About"):
                 gr.Markdown(about_markdown)  # Use the imported about page content