"""DNA Foundation Model Benchmark Leaderboard — Gradio Blocks app.""" import gradio as gr import pandas as pd from src.constants import CLASSIFICATION_METRICS from src.data import apply_filters, deduplicate, load_data from src.plots import ( bar_plot_per_category, bar_plot_per_task, bubble_context_length, efficiency_table, heatmap_variants, ranking_bar, scatter_speed, strip_plot, top_n_variants, ) # --------------------------------------------------------------------------- # Startup: load & deduplicate once # --------------------------------------------------------------------------- print("Loading benchmark data from HF Hub …") try: _LOADED_DF = load_data() _RAW_DF = deduplicate(_LOADED_DF) print(f" Loaded {len(_RAW_DF)} rows, {_RAW_DF['model_alias'].nunique()} models.") except Exception as exc: print(f" WARNING: Could not load data — {exc}") _LOADED_DF = pd.DataFrame() _RAW_DF = pd.DataFrame() _METRIC_CHOICES = list(CLASSIFICATION_METRICS.values()) _METRIC_KEYS = list(CLASSIFICATION_METRICS.keys()) _CATEGORY_CHOICES = sorted(_RAW_DF["task_category"].unique().tolist()) if not _RAW_DF.empty else [] _SPEED_COL_META = [ ("test_throughput_seq_s", "Throughput (seq/s)"), ("test_embedding_time_s", "Embedding Time (s)"), ("vram_model_mb", "VRAM (MB)"), ] _SPEED_CHOICES = [lbl for col, lbl in _SPEED_COL_META if not _RAW_DF.empty and col in _RAW_DF.columns] _SPEED_COL_MAP = {lbl: col for col, lbl in _SPEED_COL_META} _DATASET_CHOICES = ( sorted(_RAW_DF["dataset_name"].dropna().unique().tolist()) if not _RAW_DF.empty and "dataset_name" in _RAW_DF.columns else [] ) def _metric_key(label: str) -> str: if not label or label not in _METRIC_CHOICES: return _METRIC_KEYS[0] return _METRIC_KEYS[_METRIC_CHOICES.index(label)] def _numeric_range(df: pd.DataFrame, col: str) -> tuple[float, float]: if col not in df.columns or df.empty: return 0.0, 1.0 lo, hi = float(df[col].min()), float(df[col].max()) if lo == hi: return lo, lo + 1.0 return lo, hi # --------------------------------------------------------------------------- # Render helpers # --------------------------------------------------------------------------- def render_leaderboard(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 ranking_bar(df, key), strip_plot(df, key) def render_per_category(df: pd.DataFrame, metric_label: str, top_n: int = 20, datasets: list[str] | None = None): if datasets and "dataset_name" in df.columns: df = df[df["dataset_name"].isin(datasets)] key = _metric_key(metric_label) df = top_n_variants(df, key, int(top_n)) return bar_plot_per_category(df, key) def render_per_task(df: pd.DataFrame, metric_label: str, category: str, top_n: int = 20, datasets: list[str] | None = None): if datasets and "dataset_name" in df.columns: df = df[df["dataset_name"].isin(datasets)] key = _metric_key(metric_label) df = top_n_variants(df, key, int(top_n)) return bar_plot_per_task(df, key, category) def render_speed( df: pd.DataFrame, metric_label: str, speed_label: str | None, log_scale: bool, top_n: int = 20, ): key = _metric_key(metric_label) df = top_n_variants(df, key, int(top_n)) col = _SPEED_COL_MAP.get(speed_label) if speed_label else None return scatter_speed(df, key, speed_col=col, log_scale=log_scale), efficiency_table(df, key) 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) def render_variant_heatmap( df: pd.DataFrame, metric_label: str, model_alias: list[str] | str, aggregate_by_dataset: bool = False, ): key = _metric_key(metric_label) return heatmap_variants(df, key, model_alias, aggregate_by_dataset) def update_category_choices(df: pd.DataFrame, datasets: list[str] | None = None): if datasets and "dataset_name" in df.columns: df = df[df["dataset_name"].isin(datasets)] cats = sorted(df["task_category"].unique().tolist()) if not df.empty else [] return gr.update(choices=cats, value=cats[0] if cats else None) def apply_and_preview( df: pd.DataFrame, prod_only: bool, exclude_models: list[str], params_min: float, params_max: float, dim_min: float, dim_max: float, ctx_min: float, ctx_max: float, pooling_strategies: list[str] | None = None, ) -> tuple[pd.DataFrame, pd.DataFrame, str]: filters = { "subsample_train": "full" if prod_only else None, "exclude_models": exclude_models or [], "model_params_min": params_min, "model_params_max": params_max, "embedding_dim_min": dim_min, "embedding_dim_max": dim_max, "max_context_size_min": ctx_min, "max_context_size_max": ctx_max, "pooling_strategies": pooling_strategies or [], } result = apply_filters(df, filters) preview = result.head(50) if not result.empty else pd.DataFrame() return result, preview, f"{len(result)} rows active" def _hash_label(row: pd.Series) -> str: """Human-readable label for an embedding config hash dropdown entry.""" def _s(val) -> str: return "" if val is None or (isinstance(val, float) and pd.isna(val)) else str(val) category = _s(row.get("task_category")) task = _s(row.get("task_name")) pool = _s(row.get("pooling_strategy")) strategy = _s(row.get("layer_selection_strategy")) idx = row.get("layer_index") layer = strategy + (f"({int(idx)})" if idx is not None and not pd.isna(idx) else "") h = _s(row.get("embedding_config_hash"))[:8] parts = [f"{category}/{task}" if category else task, pool, layer] return " · ".join(p for p in parts if p) + f" ({h})" def get_hash_choices_for_model(loaded_df: pd.DataFrame, model: str) -> list[tuple[str, str]]: if not model or loaded_df.empty or "embedding_config_hash" not in loaded_df.columns: return [] sub = loaded_df[loaded_df["model_alias"] == model].drop_duplicates("embedding_config_hash") return sorted( [(_hash_label(row), row["embedding_config_hash"]) for _, row in sub.iterrows()], key=lambda x: x[0], ) def get_head_types_for_hash(loaded_df: pd.DataFrame, model: str, key: str) -> list[str]: if loaded_df.empty or "head_type" not in loaded_df.columns: return [] sub = loaded_df if model: sub = sub[sub["model_alias"] == model] if key and "embedding_config_hash" in loaded_df.columns: sub = sub[sub["embedding_config_hash"] == key] return sorted(sub["head_type"].dropna().unique().tolist()) _DEFAULT_INSP_COLS = ["kept", "run_id", "run_at", "subsample_train", "task_name", "mcc_test", "accuracy_test"] def build_inspector_table( loaded_df: pd.DataFrame, raw_df: pd.DataFrame, model: str, key: str, head_type: str, selected_cols: list[str] | None = None, task_cats: list[str] | None = None, tasks: list[str] | None = None, pooling_strategies: list[str] | None = None, ) -> pd.DataFrame: if loaded_df.empty: return pd.DataFrame() masks = [] if model: masks.append(loaded_df["model_alias"] == model) if key and "embedding_config_hash" in loaded_df.columns: masks.append(loaded_df["embedding_config_hash"] == key) if head_type and "head_type" in loaded_df.columns: masks.append(loaded_df["head_type"] == head_type) if task_cats and "task_category" in loaded_df.columns: masks.append(loaded_df["task_category"].isin(task_cats)) if tasks and "task_name" in loaded_df.columns: masks.append(loaded_df["task_name"].isin(tasks)) if pooling_strategies and "pooling_strategy" in loaded_df.columns: masks.append(loaded_df["pooling_strategy"].isin(pooling_strategies)) if not masks: return pd.DataFrame() combined = masks[0] for m in masks[1:]: combined = combined & m sub = loaded_df[combined].copy() kept_ids = set(raw_df["run_id"]) if "run_id" in raw_df.columns else set() sub["kept"] = sub["run_id"].isin(kept_ids).map({True: "✓", False: "✗"}) cols = selected_cols if selected_cols else _DEFAULT_INSP_COLS return sub[[c for c in cols if c in sub.columns or c == "kept"]].reset_index(drop=True) def exclude_run_ids(df: pd.DataFrame, run_ids_text: str) -> tuple[pd.DataFrame, pd.DataFrame, str]: ids = [r.strip() for r in run_ids_text.split(",") if r.strip()] result = apply_filters(df, {"exclude_run_ids": ids}) preview = result.head(50) if not result.empty else pd.DataFrame() return result, preview, f"{len(result)} rows active" def _filter_status_banner(filtered_df: pd.DataFrame, raw_df: pd.DataFrame) -> str: if filtered_df.empty or raw_df.empty: return "" if len(filtered_df) < len(raw_df): return ( f"⚠️ **Active scope filters** — {len(filtered_df):,} / {len(raw_df):,} rows visible. " "Go to 🔬 Scope / Filter to review." ) return "" # --------------------------------------------------------------------------- # Build UI # --------------------------------------------------------------------------- all_models = sorted(_RAW_DF["model_alias"].unique().tolist()) if not _RAW_DF.empty else [] params_min_v, params_max_v = _numeric_range(_RAW_DF, "model_params") dim_min_v, dim_max_v = _numeric_range(_RAW_DF, "embedding_dim") ctx_min_v, ctx_max_v = _numeric_range(_RAW_DF, "max_context_size") _insp_default_model = all_models[0] if all_models else None _insp_default_keys = get_hash_choices_for_model(_LOADED_DF, _insp_default_model) _insp_default_heads = get_head_types_for_hash(_LOADED_DF, _insp_default_model, None) _ALL_INSP_COLS = ["kept"] + _LOADED_DF.columns.tolist() if not _LOADED_DF.empty else _DEFAULT_INSP_COLS _all_task_cats = ( sorted(_LOADED_DF["task_category"].dropna().unique().tolist()) if not _LOADED_DF.empty and "task_category" in _LOADED_DF.columns else [] ) _all_tasks = ( sorted(_LOADED_DF["task_name"].dropna().unique().tolist()) if not _LOADED_DF.empty and "task_name" in _LOADED_DF.columns else [] ) _all_pooling_strategies = ( sorted(_LOADED_DF["pooling_strategy"].dropna().unique().tolist()) if not _LOADED_DF.empty and "pooling_strategy" in _LOADED_DF.columns else [] ) with gr.Blocks(title="DNA Benchmark Leaderboard", theme=gr.themes.Soft()) as demo: gr.Markdown( "# 🧬 DNA Foundation Model Benchmark Leaderboard\n" "Compare DNA language models across genomics classification tasks." ) filter_banner = gr.Markdown("") # ── Global controls ─────────────────────────────────────────────────────── with gr.Row(): metric_dd = gr.Dropdown(choices=_METRIC_CHOICES, value=_METRIC_CHOICES[0], label="Metric", scale=2) top_n_sl = gr.Slider(minimum=5, maximum=50, step=5, value=20, label="Top N variants", scale=1) # ── Shared state ────────────────────────────────────────────────────────── loaded_df_state = gr.State(_LOADED_DF) raw_df_state = gr.State(_RAW_DF) filtered_df_state = gr.State(_RAW_DF) # ── Tabs ────────────────────────────────────────────────────────────────── with gr.Tabs() as tabs: # ── Tab 1: Leaderboard ──────────────────────────────────────────────── with gr.Tab("🏅 Leaderboard"): leaderboard_fig = gr.Plot(label="Variant ranking (best at top)") violin_fig = gr.Plot(label="Per-task scores by variant (rank 1 at top)") gr.Markdown("---\n### Drill into model variants") model_dd = gr.Dropdown( choices=all_models, value=[all_models[0]] if all_models else [], multiselect=True, label="Select model(s) — single model: all variants as rows; multiple: best variant per family", ) agg_dataset_cb = gr.Checkbox( label="Aggregate by dataset (show mean per benchmark dataset instead of per task)", value=False, ) variant_violin_fig = gr.Plot(label="Variant comparison (metric per task)") # ── Tab 2: Per-Category Performance ─────────────────────────────────── with gr.Tab("📊 Per-Category Performance"): category_fig = gr.Plot(label="Mean performance per task category") gr.Markdown("---\n### Drill into individual tasks") category_dd = gr.Dropdown( choices=_CATEGORY_CHOICES, value=_CATEGORY_CHOICES[0] if _CATEGORY_CHOICES else None, label="Select category", ) dataset_dd = gr.Dropdown( choices=_DATASET_CHOICES, value=[], multiselect=True, label="Filter by dataset", ) per_task_fig = gr.Plot(label="Performance per task within selected category") # ── Tab 3: Speed vs Performance ─────────────────────────────────────── with gr.Tab("⚡ Speed vs Performance"): gr.Markdown( "Each point = one variant (mean metric across all tasks vs mean speed across tasks). " "**Pareto frontier** = variants not dominated on both axes (black outline)." ) with gr.Row(): speed_metric_radio = gr.Radio( choices=_SPEED_CHOICES, value=_SPEED_CHOICES[0] if _SPEED_CHOICES else None, label="X-axis metric", scale=3, ) log_scale_cb = gr.Checkbox(value=True, label="Log scale on X-axis", scale=1) speed_fig = gr.Plot(label="Performance vs efficiency tradeoff") gr.Markdown("### Efficiency ranking") efficiency_tbl = gr.Dataframe( label="Variants sorted by performance — speed and derived efficiency metrics", interactive=False, ) # ── Tab 4: Context Length ───────────────────────────────────────────── with gr.Tab("🔬 Context Length"): context_fig = gr.Plot(label="Performance vs task sequence length") # ── Tab 5: Scope / Filter ───────────────────────────────────────────── with gr.Tab("🔬 Scope / Filter"): prod_only_cb = gr.Checkbox( label="Production runs only (subsample_train IS NULL)", value=False, info="Currently returns no data — all runs use subsample_train=0.05.", ) exclude_models_ms = gr.Dropdown(choices=all_models, multiselect=True, label="Exclude models", value=[]) pooling_filter_ms = gr.Dropdown( choices=_all_pooling_strategies, multiselect=True, label="Pooling strategy", value=[], info="Leave empty to include all pooling strategies.", ) with gr.Row(): params_min_sl = gr.Slider( minimum=params_min_v, maximum=params_max_v, value=params_min_v, label="Model params (min)" ) params_max_sl = gr.Slider( minimum=params_min_v, maximum=params_max_v, value=params_max_v, label="Model params (max)" ) with gr.Row(): dim_min_sl = gr.Slider( minimum=dim_min_v, maximum=dim_max_v, value=dim_min_v, label="Embedding dim (min)" ) dim_max_sl = gr.Slider( minimum=dim_min_v, maximum=dim_max_v, value=dim_max_v, label="Embedding dim (max)" ) with gr.Row(): ctx_min_sl = gr.Slider( minimum=ctx_min_v, maximum=ctx_max_v, value=ctx_min_v, label="Max context size (min bp)" ) ctx_max_sl = gr.Slider( minimum=ctx_min_v, maximum=ctx_max_v, value=ctx_max_v, label="Max context size (max bp)" ) apply_btn = gr.Button("Apply Filters", variant="primary") gr.Markdown("---") run_ids_tb = gr.Textbox(label="Exclude run_ids (comma-separated)", placeholder="run_id_1, run_id_2, …") exclude_runs_btn = gr.Button("Exclude run_ids") gr.Markdown("---") settings_count = gr.Label(label="Active rows") settings_preview = gr.Dataframe(label="Preview (first 50 rows)", interactive=False) gr.Markdown( "---\n### Raw Run Inspector\nSelect a model and embedding config to audit which runs were kept or dropped by dedup." ) with gr.Row(): insp_model_dd = gr.Dropdown(choices=all_models, value=_insp_default_model, label="Model", scale=2) insp_cache_dd = gr.Dropdown( choices=_insp_default_keys, value=None, label="Embedding config hash", scale=3 ) insp_head_dd = gr.Dropdown(choices=_insp_default_heads, value=None, label="Head type", scale=2) with gr.Row(): insp_task_cat_ms = gr.Dropdown( choices=_all_task_cats, value=None, multiselect=True, label="Task category", scale=2 ) insp_task_ms = gr.Dropdown(choices=_all_tasks, value=None, multiselect=True, label="Task", scale=3) insp_pooling_ms = gr.Dropdown( choices=_all_pooling_strategies, value=None, multiselect=True, label="Pooling strategy", scale=2 ) insp_cols_dd = gr.Dropdown( choices=_ALL_INSP_COLS, value=_DEFAULT_INSP_COLS, multiselect=True, label="Columns to display", ) insp_table = gr.Dataframe( label="Runs for selection (✓ = kept by dedup, ✗ = dropped)", interactive=False, ) # ── Event wiring ────────────────────────────────────────────────────────── _leaderboard_triggers = [filtered_df_state, metric_dd, top_n_sl] _context_triggers = [filtered_df_state, metric_dd, top_n_sl] _speed_triggers = [filtered_df_state, metric_dd, speed_metric_radio, log_scale_cb, top_n_sl] for trigger in [filtered_df_state, metric_dd, top_n_sl]: trigger.change(render_leaderboard, _leaderboard_triggers, [leaderboard_fig, violin_fig]) trigger.change(render_context, _context_triggers, context_fig) for trigger in [filtered_df_state, metric_dd, speed_metric_radio, log_scale_cb, top_n_sl]: trigger.change(render_speed, _speed_triggers, [speed_fig, efficiency_tbl]) # Category bar chart: reacts to df, metric, top_n, AND dataset filter _cat_triggers = [filtered_df_state, metric_dd, top_n_sl, dataset_dd] for trigger in [filtered_df_state, metric_dd, top_n_sl, dataset_dd]: trigger.change(render_per_category, _cat_triggers, category_fig) # Per-task drill-down: reacts to category/dataset dropdowns OR global controls _per_task_triggers = [filtered_df_state, metric_dd, category_dd, top_n_sl, dataset_dd] for trigger in [filtered_df_state, metric_dd, category_dd, top_n_sl, dataset_dd]: trigger.change(render_per_task, _per_task_triggers, per_task_fig) # Variant heatmap: reacts to model selector, metric, filtered data, or aggregation toggle _variant_triggers = [filtered_df_state, metric_dd, model_dd, agg_dataset_cb] for trigger in [filtered_df_state, metric_dd, model_dd, agg_dataset_cb]: trigger.change(render_variant_heatmap, _variant_triggers, variant_violin_fig) # Update category dropdown choices when global filter OR dataset filter changes for trigger in [filtered_df_state, dataset_dd]: trigger.change(update_category_choices, [filtered_df_state, dataset_dd], category_dd) # Settings — Apply Filters apply_btn.click( apply_and_preview, inputs=[ raw_df_state, prod_only_cb, exclude_models_ms, params_min_sl, params_max_sl, dim_min_sl, dim_max_sl, ctx_min_sl, ctx_max_sl, pooling_filter_ms, ], outputs=[filtered_df_state, settings_preview, settings_count], ).then( _filter_status_banner, inputs=[filtered_df_state, raw_df_state], outputs=filter_banner, ) # Settings — Exclude run_ids exclude_runs_btn.click( exclude_run_ids, inputs=[filtered_df_state, run_ids_tb], outputs=[filtered_df_state, settings_preview, settings_count], ).then( _filter_status_banner, inputs=[filtered_df_state, raw_df_state], outputs=filter_banner, ) # Inspector — model updates hash choices; model+hash each update head_type choices independently insp_model_dd.change( lambda df, m: gr.update(choices=get_hash_choices_for_model(df, m), value=None), [loaded_df_state, insp_model_dd], insp_cache_dd, ) insp_model_dd.change( lambda df, m, k: gr.update(choices=get_head_types_for_hash(df, m, k), value=None), [loaded_df_state, insp_model_dd, insp_cache_dd], insp_head_dd, ) insp_cache_dd.change( lambda df, m, k: gr.update(choices=get_head_types_for_hash(df, m, k), value=None), [loaded_df_state, insp_model_dd, insp_cache_dd], insp_head_dd, ) _insp_inputs = [ loaded_df_state, raw_df_state, insp_model_dd, insp_cache_dd, insp_head_dd, insp_cols_dd, insp_task_cat_ms, insp_task_ms, insp_pooling_ms, ] for trigger in [ insp_model_dd, insp_cache_dd, insp_head_dd, insp_cols_dd, insp_task_cat_ms, insp_task_ms, insp_pooling_ms, ]: trigger.change(build_inspector_table, _insp_inputs, insp_table) # Trigger resize on tab switch so Plotly reflows to correct container width tabs.select(fn=None, js="() => setTimeout(() => window.dispatchEvent(new Event('resize')), 100)") # Initial render on load demo.load(render_leaderboard, inputs=_leaderboard_triggers, outputs=[leaderboard_fig, violin_fig]) demo.load(render_per_category, inputs=_cat_triggers, outputs=category_fig) demo.load(render_per_task, inputs=_per_task_triggers, outputs=per_task_fig) demo.load(render_speed, inputs=_speed_triggers, outputs=[speed_fig, efficiency_tbl]) demo.load(render_context, inputs=_context_triggers, outputs=context_fig) demo.load(render_variant_heatmap, inputs=_variant_triggers, outputs=variant_violin_fig) if __name__ == "__main__": demo.launch(share=False, server_name="0.0.0.0")