feat: add quantization to "Model Size vs Performance" plot

src/components/visualizations.py

@@ -7,6 +7,7 @@ import plotly.express as px
 import pandas as pd
 from typing import Optional, Dict, List, Set
 import plotly.graph_objects as go
+from ..core.scoring import get_quantization_tier
 
 
 def clean_device_id(device_id: str) -> str:
@@ -16,6 +17,24 @@ def clean_device_id(device_id: str) -> str:
     return device_id
 
 
+def get_quant_name(factor: float) -> str:
+    """Get human-readable name for quantization factor"""
+    if pd.isna(factor):
+        return "Unknown"
+    if factor >= 1.0:
+        return "No Quantization (F16/F32)"
+    quant_map = {
+        0.8: "[i]Q8_x",
+        0.6: "[i]Q6_x",
+        0.5: "[i]Q5_x",
+        0.4: "[i]Q4_x",
+        0.3: "[i]Q3_x",
+        0.2: "[i]Q2_x",
+        0.1: "[i]Q1_x",
+    }
+    return quant_map.get(factor, f"Q{int(factor*10)}_x")
+
+
 def create_performance_plot(
     df: pd.DataFrame, metric: str, title: str, hover_data: List[str] = None
 ):
@@ -113,7 +132,9 @@ def filter_dataframe(df: pd.DataFrame, filters: Dict) -> pd.DataFrame:
     return filtered_df
 
 
-def create_model_size_performance_plot(df: pd.DataFrame, device_id: str, title: str):
+def create_model_size_performance_plot(
+    df: pd.DataFrame, device_id: str, quant_filter: str, title: str
+):
     """Create a plot showing model size vs performance metrics for a specific device"""
     if df.empty:
         return None
@@ -123,60 +144,103 @@ def create_model_size_performance_plot(df: pd.DataFrame, device_id: str, title:
     if device_df.empty:
         return None
 
+    # Filter by quantization if specified
+    if quant_filter != "All":
+        device_df = device_df[
+            device_df["Model ID"].apply(
+                lambda x: get_quantization_tier(x) == float(quant_filter)
+            )
+        ]
+        if device_df.empty:
+            return None
+
     # Create a new figure with secondary y-axis
     fig = go.Figure()
 
+    # Define shapes for different quantization levels
+    quant_shapes = {
+        1.0: "circle",  # F16/F32
+        0.8: "square",  # Q8
+        0.6: "diamond",  # Q6
+        0.5: "triangle-up",  # Q5
+        0.4: "triangle-down",  # Q4
+        0.3: "star",  # Q3
+        0.2: "pentagon",  # Q2
+        0.1: "hexagon",  # Q1
+    }
+
     # Add Token Generation data (left y-axis)
-    fig.add_trace(
-        go.Scatter(
-            x=device_df["Model Size"],
-            y=device_df["Token Generation"],
-            name="Token Generation",
-            mode="markers",
-            marker=dict(color="#2ecc71"),
-            yaxis="y",
+    for quant in sorted(device_df["quant_factor"].unique()):
+        quant_df = device_df[device_df["quant_factor"] == quant]
+        if quant_df.empty:
+            continue
+
+        quant_name = get_quant_name(quant)
+        fig.add_trace(
+            go.Scatter(
+                x=quant_df["Model Size"],
+                y=quant_df["Token Generation"],
+                name=f"{quant_name}",
+                mode="markers",
+                marker=dict(
+                    color="#2ecc71",
+                    symbol=quant_shapes.get(quant, "circle"),
+                    size=10,
+                ),
+                yaxis="y",
+                legendgroup="quant",
+                showlegend=True,
+            )
         )
-    )
 
     # Add Prompt Processing data (right y-axis)
-    fig.add_trace(
-        go.Scatter(
-            x=device_df["Model Size"],
-            y=device_df["Prompt Processing"],
-            name="Prompt Processing",
-            mode="markers",
-            marker=dict(color="#e74c3c"),
-            yaxis="y2",
+    for quant in sorted(device_df["quant_factor"].unique()):
+        quant_df = device_df[device_df["quant_factor"] == quant]
+        if quant_df.empty:
+            continue
+
+        fig.add_trace(
+            go.Scatter(
+                x=quant_df["Model Size"],
+                y=quant_df["Prompt Processing"],
+                name=f"{quant_name}",
+                mode="markers",
+                marker=dict(
+                    color="#e74c3c",
+                    symbol=quant_shapes.get(quant, "circle"),
+                    size=10,
+                ),
+                yaxis="y2",
+                legendgroup="quant",
+                showlegend=False,  # Don't show duplicate quantization entries in legend
+            )
         )
-    )
 
     # Add trend lines if enough points
     if len(device_df) > 2:
         # TG trend line
         tg_trend = px.scatter(
             device_df, x="Model Size", y="Token Generation", trendline="lowess"
-        ).data[
-            1
-        ]  # Get the trend line trace
+        ).data[1]
         tg_trend.update(
             line=dict(color="#2ecc71", dash="solid"),
-            name="TG Trend",
-            showlegend=False,
+            name="Token Generation",
+            showlegend=False,  # Hide from legend
             yaxis="y",
+            legendgroup="metric",
         )
         fig.add_trace(tg_trend)
 
         # PP trend line
         pp_trend = px.scatter(
             device_df, x="Model Size", y="Prompt Processing", trendline="lowess"
-        ).data[
-            1
-        ]  # Get the trend line trace
+        ).data[1]
         pp_trend.update(
             line=dict(color="#e74c3c", dash="solid"),
-            name="PP Trend",
-            showlegend=False,
+            name="Prompt Processing",
+            showlegend=False,  # Hide from legend
             yaxis="y2",
+            legendgroup="metric",
         )
         fig.add_trace(pp_trend)
 
@@ -186,10 +250,7 @@ def create_model_size_performance_plot(df: pd.DataFrame, device_id: str, title:
         xaxis=dict(
             title="Model Size (B)",
             gridcolor="lightgrey",
-            range=[
-                0,
-                max(device_df["Model Size"]) * 1.05,
-            ],  # Start from 0, add 5% padding to max
+            range=[0, max(device_df["Model Size"]) * 1.05],
         ),
         yaxis=dict(
             title="Token Generation (t/s)",
@@ -197,10 +258,7 @@ def create_model_size_performance_plot(df: pd.DataFrame, device_id: str, title:
             tickfont=dict(color="#2ecc71"),
             gridcolor="lightgrey",
             side="left",
-            range=[
-                0,
-                max(device_df["Token Generation"]) * 1.05,
-            ],  # Start from 0, add 5% padding to max
+            range=[0, max(device_df["Token Generation"]) * 1.05],
         ),
         yaxis2=dict(
             title="Prompt Processing (t/s)",
@@ -209,22 +267,21 @@ def create_model_size_performance_plot(df: pd.DataFrame, device_id: str, title:
             anchor="x",
             overlaying="y",
             side="right",
-            range=[
-                0,
-                max(device_df["Prompt Processing"]) * 1.05,
-            ],  # Start from 0, add 5% padding to max
+            range=[0, max(device_df["Prompt Processing"]) * 1.05],
         ),
         height=400,
         showlegend=True,
         plot_bgcolor="white",
         legend=dict(
-            yanchor="middle",
-            y=0.8,
+            yanchor="top",
+            y=0.99,
             xanchor="right",
             x=0.99,
-            bgcolor="rgba(255, 255, 255, 0.8)",  # Semi-transparent white background
+            bgcolor="rgba(255, 255, 255, 0.8)",
             bordercolor="lightgrey",
             borderwidth=1,
+            groupclick="togglegroup",  # Toggle all traces in the same group
+            title="Quantization",  # Add legend title
         ),
     )
 
@@ -255,30 +312,48 @@ def render_model_size_performance(df: pd.DataFrame, filters: Dict):
         device_id: clean_device_id(device_id) for device_id in device_ids
     }
 
-    # Device selector for size vs performance plots
-    selected_device_id = st.selectbox(
-        "Select Device",
-        options=device_ids,
-        format_func=lambda x: device_display_names[
-            x
-        ],  # Display clean names in dropdown
-        help="Select a device to view its performance across different model sizes",
-        key="size_perf_device_selector",
-        placeholder="Search for a device...",
-        index=default_index,
-    )
+    # Create columns for device and quantization selectors
+    col1, col2 = st.columns([2, 1])
+
+    with col1:
+        # Device selector
+        selected_device_id = st.selectbox(
+            "Select Device",
+            options=device_ids,
+            format_func=lambda x: device_display_names[x],
+            help="Select a device to view its performance across different model sizes",
+            key="size_perf_device_selector",
+            placeholder="Search for a device...",
+            index=default_index,
+        )
+
+    with col2:
+        # Quantization filter
+        quant_options = ["All"] + [
+            str(q) for q in sorted(size_perf_df["quant_factor"].unique())
+        ]
+        quant_filter = st.selectbox(
+            "Filter by Quantization",
+            options=quant_options,
+            format_func=lambda x: (
+                "All Quantizations" if x == "All" else get_quant_name(float(x))
+            ),
+            help="Filter data points by quantization level",
+            key="size_perf_quant_selector",
+        )
 
     # Create and display the model size vs performance plot
     size_perf_fig = create_model_size_performance_plot(
         size_perf_df,
         selected_device_id,
+        quant_filter,
         f"Model Size vs Performance Metrics for {device_display_names[selected_device_id]}",
     )
 
     if size_perf_fig:
         st.plotly_chart(size_perf_fig, use_container_width=True)
     else:
-        st.warning("No data available for the selected device.")
+        st.warning("No data available for the selected device and quantization level.")
 
 
 def render_performance_plots(df: pd.DataFrame, filters: Dict):
@@ -721,21 +796,6 @@ def render_device_rankings(df: pd.DataFrame):
     with rank_tab3:
         st.subheader("🔍 Rankings by Quantization")
 
-        # Helper function to get quantization name from factor
-        def get_quant_name(factor: float) -> str:
-            if factor >= 1.0:
-                return "No Quantization (F16/F32)"
-            quant_map = {
-                0.8: "[i]Q8_x",
-                0.6: "[i]Q6_x",
-                0.5: "[i]Q5_x",
-                0.4: "[i]Q4_x",
-                0.3: "[i]Q3_x",
-                0.2: "[i]Q2_x",
-                0.1: "[i]Q1_x",
-            }
-            return quant_map.get(factor, f"Q{int(factor*10)}_x")
-
         # Group by device and quantization level
         quant_rankings = df.copy()
         quant_summary = (