Rebuild Context Length tab as per-variant, model-context-size chart

Switch X-axis from task sequence length to model max context size so the
chart actually answers whether larger context windows improve performance;
split into an overall + by-task-category panel. Task sequence length is
kept via bubble size instead of being dropped. Plots every variant
combination (not aggregated per model) to match production deployment
granularity.

app.py

@@ -8,7 +8,7 @@ from src.data import apply_filters, deduplicate, load_data
 from src.plots import (
     bar_plot_per_category,
     bar_plot_per_task,
-    bubble_context_length,
+    context_length_overview,
     efficiency_table,
     heatmap_variants,
     ranking_bar,
@@ -109,7 +109,7 @@ def render_speed(
 def render_context(df: pd.DataFrame, metric_label: str, top_n: int = 20):
     key = _metric_key(metric_label)
     df = top_n_variants(df, key, int(top_n))
-    return bubble_context_length(df, key)
+    return context_length_overview(df, key)
 
 
 def render_variant_heatmap(
@@ -369,7 +369,7 @@ with gr.Blocks(title="DNA Benchmark Leaderboard", theme=gr.themes.Soft()) as dem
 
         # ── Tab 4: Context Length ─────────────────────────────────────────────
         with gr.Tab("🔬 Context Length"):
-            context_fig = gr.Plot(label="Performance vs task sequence length")
+            context_fig = gr.Plot(label="Performance vs model context size (bubble size = task sequence length)")
 
         # ── Tab 5: Scope / Filter ─────────────────────────────────────────────
         with gr.Tab("🔬 Scope / Filter"):

src/plots.py

@@ -5,6 +5,7 @@ import json
 import numpy as np
 import pandas as pd
 import plotly.graph_objects as go
+from plotly.subplots import make_subplots
 
 from src.constants import (
     CATEGORY_COLORS,
@@ -772,13 +773,30 @@ def efficiency_table(df: pd.DataFrame, metric: str) -> pd.DataFrame:
 
 
 # ---------------------------------------------------------------------------
-# 7. Context length bubble chart
+# 7. Context length overview (overall + by-category, X = model max context size)
 # ---------------------------------------------------------------------------
 
-def bubble_context_length(df: pd.DataFrame, metric: str) -> go.Figure:
-    """Bubble chart: X=task_seq_length, Y=metric, color=model family, size=max_context_size.
+# Marker symbol per task category, used in the by-category panel legend.
+_CATEGORY_SYMBOLS: dict[str, str] = {
+    "histone_marks":  "circle",
+    "promoter":       "square",
+    "enhancer":       "triangle-up",
+    "splice_site":    "diamond",
+    "variant_effect": "triangle-down",
+    "rna_expression": "star",
+    "other":          "cross",
+}
+
+
+def context_length_overview(df: pd.DataFrame, metric: str) -> go.Figure:
+    """Two-panel figure: model max context size (X) vs performance (Y).
 
-    One bubble per (variant, task). Legend grouped by model family.
+    One point per variant combination (model + pooling + layer + head) — not
+    collapsed to a single "best" variant per model, since that's the granularity
+    actually deployed. Left panel: mean metric across that variant's tasks.
+    Right panel: same variant, broken out by task category (color=model,
+    symbol=category). Bubble size in both panels encodes mean task sequence
+    length, so that dimension stays visible even though it's no longer on the axis.
     """
     if not TASK_SEQ_LENGTHS:
         fig = go.Figure()
@@ -792,73 +810,118 @@ def bubble_context_length(df: pd.DataFrame, metric: str) -> go.Figure:
         fig.update_layout(height=480, template="plotly_white")
         return fig
 
-    if df.empty or metric not in df.columns or "task_seq_length" not in df.columns:
+    required = ["model_alias", "task_name", "task_category", "task_seq_length", "max_context_size", metric]
+    if df.empty or any(c not in df.columns for c in required):
         return go.Figure()
 
     _CONFIG_COLS = ["model_hf_id", "pooling_strategy", "layer_selection_params", "head_type"]
     extra_cols = [c for c in _CONFIG_COLS if c in df.columns]
-    base_cols = ["model_alias", "task_name", "task_seq_length", metric, "max_context_size"]
-    plot_df = df[base_cols + extra_cols].dropna(subset=["task_seq_length", metric]).copy()
+    d = df.dropna(subset=[metric, "task_seq_length", "max_context_size"])[required + extra_cols].copy()
+    if d.empty:
+        return go.Figure()
+    d["_vkey"] = _variant_key_col(d)
 
-    plot_df["_vkey"] = _variant_key_col(plot_df)
-    if "layer_selection_params" in plot_df.columns:
-        plot_df["_layer_info"] = plot_df["layer_selection_params"].apply(_fmt_layer)
+    metric_label = CLASSIFICATION_METRICS.get(metric, metric)
+    models = sorted(d["model_alias"].unique())
+    model_color_map = _model_color_map(models)
+    n_tasks = d.groupby("model_alias")["task_name"].nunique().max()
 
-    def _config_str(row: pd.Series) -> str:
-        parts = []
-        for col, lbl in [("model_hf_id", "HF model"), ("pooling_strategy", "Pooling"),
-                         ("_layer_info", "Layer"), ("head_type", "Head")]:
-            if col not in row.index:
-                continue
-            v = row[col]
-            if v is not None and not (isinstance(v, float) and pd.isna(v)):
-                parts.append(f"{lbl}: {v}")
-        return "<br>".join(parts)
+    seq_all = d["task_seq_length"]
+    seq_lo, seq_hi = seq_all.min(), seq_all.max()
 
-    plot_df["_config_str"] = plot_df.apply(_config_str, axis=1)
+    def _size(seq_len: float) -> float:
+        if seq_hi == seq_lo:
+            return 22.0
+        return 14 + (seq_len - seq_lo) / (seq_hi - seq_lo) * 30
 
-    models = sorted(plot_df["model_alias"].unique()) if "model_alias" in plot_df.columns else []
-    model_color_map = _model_color_map(models)
-    max_ctx = plot_df["max_context_size"].max() or 1
-    metric_label = CLASSIFICATION_METRICS.get(metric, metric)
-    fig = go.Figure()
+    vkey_label = {vkey: _variant_label(g.iloc[0]) for vkey, g in d.groupby("_vkey")}
 
-    legend_shown: set[str] = set()
+    fig = make_subplots(
+        rows=1, cols=2,
+        subplot_titles=(
+            "Context length vs. overall performance",
+            "Context length vs. performance, by task category",
+        ),
+        horizontal_spacing=0.1,
+    )
 
-    for vkey in plot_df["_vkey"].unique():
-        vdf = plot_df[plot_df["_vkey"] == vkey]
-        first = vdf.iloc[0]
-        model = _norm(first.get("model_alias", ""))
-        label = _variant_label(first)
+    # ---- Left panel: one point per variant combination ---------------------
+    overall = (
+        d.groupby(["model_alias", "_vkey"])
+        .agg(y=(metric, "mean"), n=(metric, "count"), max_ctx=("max_context_size", "first"),
+             seq_mean=("task_seq_length", "mean"))
+        .reset_index()
+    )
+    legend_shown: set[str] = set()
+    for _, row in overall.iterrows():
+        model = row["model_alias"]
+        vlabel = vkey_label.get(row["_vkey"], row["_vkey"])
         color = model_color_map.get(model, _MODEL_PALETTE[0])
         show_legend = model not in legend_shown
         legend_shown.add(model)
+        fig.add_trace(
+            go.Scatter(
+                x=[row["max_ctx"]], y=[row["y"]], mode="markers",
+                name=model, legendgroup=model, showlegend=show_legend,
+                marker=dict(size=_size(row["seq_mean"]), color=color, opacity=0.75, line=dict(width=1, color="white")),
+                hovertemplate=(
+                    f"<b>{model}</b><br>{vlabel}<br>Max context: %{{x}} bp<br>"
+                    f"{metric_label}: %{{y:.3f}} ({row['n']} tasks)<br>"
+                    "Mean task seq length: %{customdata:.0f} bp<extra></extra>"
+                ),
+                customdata=[row["seq_mean"]],
+            ),
+            row=1, col=1,
+        )
 
-        sizes = (vdf["max_context_size"].fillna(max_ctx) / max_ctx * 40 + 8).clip(8, 50)
-        cd = vdf[["max_context_size", "_config_str"]].to_numpy(dtype=object)
+    # ---- Right panel: per (variant, task category) --------------------------
+    by_cat = (
+        d.groupby(["model_alias", "_vkey", "task_category"])
+        .agg(y=(metric, "mean"), max_ctx=("max_context_size", "first"), seq_mean=("task_seq_length", "mean"))
+        .reset_index()
+    )
+    for _, row in by_cat.iterrows():
+        cat = row["task_category"]
+        vlabel = vkey_label.get(row["_vkey"], row["_vkey"])
+        color = model_color_map.get(row["model_alias"], _MODEL_PALETTE[0])
+        symbol = _CATEGORY_SYMBOLS.get(cat, "cross")
+        cat_label = CATEGORY_DISPLAY_NAMES.get(cat, cat)
         fig.add_trace(
             go.Scatter(
-                x=vdf["task_seq_length"], y=vdf[metric], mode="markers",
-                name=model, text=vdf["task_name"],
-                marker=dict(size=sizes, color=color, opacity=0.7, line=dict(width=1, color="white")),
-                customdata=cd,
-                legendgroup=model, showlegend=show_legend,
+                x=[row["max_ctx"]], y=[row["y"]], mode="markers",
+                name=row["model_alias"], legendgroup=row["model_alias"], showlegend=False,
+                marker=dict(size=_size(row["seq_mean"]), color=color, symbol=symbol, opacity=0.75,
+                            line=dict(width=1, color="white")),
                 hovertemplate=(
-                    f"<b>%{{text}}</b> — {label}<br>"
-                    "Seq length: %{x} bp<br>"
-                    f"{metric_label}: %{{y:.3f}}<br>"
-                    "Max context: %{customdata[0]} bp<br>"
-                    "%{customdata[1]}<extra></extra>"
+                    f"<b>{row['model_alias']}</b> — {cat_label}<br>{vlabel}<br>Max context: %{{x}} bp<br>"
+                    f"{metric_label}: %{{y:.3f}}<br>Mean task seq length: %{{customdata:.0f}} bp<extra></extra>"
                 ),
-            )
+                customdata=[row["seq_mean"]],
+            ),
+            row=1, col=2,
         )
 
+    # Legend-only dummy traces so task-category shapes have a key.
+    for cat in sorted(by_cat["task_category"].unique()):
+        fig.add_trace(
+            go.Scatter(
+                x=[None], y=[None], mode="markers",
+                marker=dict(size=10, color="#888", symbol=_CATEGORY_SYMBOLS.get(cat, "cross")),
+                name=CATEGORY_DISPLAY_NAMES.get(cat, cat),
+                legendgroup="category", legendgrouptitle_text="Task category",
+                showlegend=True,
+            ),
+            row=1, col=2,
+        )
+
+    fig.update_xaxes(title_text="Max context size (bp)", type="log", row=1, col=1)
+    fig.update_xaxes(title_text="Max context size (bp)", type="log", row=1, col=2)
+    fig.update_yaxes(title_text=f"Mean {metric_label} per variant ({n_tasks} tasks)", row=1, col=1)
+    fig.update_yaxes(title_text=f"Mean {metric_label}", row=1, col=2)
     fig.update_layout(
-        title=f"{metric_label} vs task sequence length (bubble size = model max context)",
-        xaxis_title="Task sequence length (bp)",
-        yaxis_title=metric_label,
-        height=520,
+        height=560,
         template="plotly_white",
-        legend=_LEGEND_H,
+        legend=dict(orientation="v", yanchor="top", y=1, xanchor="left", x=1.02),
+        margin=dict(r=160),
     )
     return fig