Blame for src/utils.py

import functools
from pathlib import Path

import altair as alt
import fev
import pandas as pd
import pandas.io.formats.style

# Color constants - all colors defined in one place

COLORS = {
    "bar_fill": "#8d5eb7",
    "error_bar": "#222222",
    "point": "#111111",
    "text_white": "white",
    "text_black": "black",
    "text_default": "#111",
    "gold": "#F7D36B",
    "silver": "#E5E7EB",
    "bronze": "#E6B089",
    "leakage_impute": "#3B82A0",
    "failure_impute": "#E07B39",
}
HEATMAP_COLOR_SCHEME = "purplegreen"

# Display metadata for every model on the leaderboard. Generated by save_tables.py from
# benchmarks/fev_bench/models.yaml in the fev repo, where submitters declare it themselves.
MODELS_CSV = Path(__file__).resolve().parent.parent / "tables" / "models.csv"

# Human-readable labels for the model types.
MODEL_TYPE_LABELS = {
    "pretrained": "Pretrained",
    "task-specific": "Task-specific",
    "statistical": "Statistical",
    "system": "System",
    "closed-api": "Closed API",
}
# Tags shown in the leaderboard's Tags column, in display order, with the chip color for each.
# Pale tints drawn from the page palette (the purple bars, the purple-green heatmap and the medal
# colors), so the tags stay readable without competing with the numbers. Non-commercial gets the warmest tint,
# as the one tag readers usually scan for.
ZERO_SHOT_TAG = "Zero-shot"
NON_COMMERCIAL_TAG = "Non-commercial"
TAG_COLORS = {
    "Pretrained": "#E9DDF5",  # pale lavender
    "Task-specific": "#DCE6F2",  # pale blue
    "Statistical": "#ECE6DD",  # pale sand
    "System": "#DCEEE4",  # pale sage
    "Closed API": "#F3DEDE",  # pale rose
    ZERO_SHOT_TAG: "#ECEDEF",  # neutral gray
    NON_COMMERCIAL_TAG: "#F7E2B0",  # pale gold
}
# Types shown unless the reader opts into more. Results for the others either cannot be reproduced
# independently (closed-api) or are not comparable to single models (system).
DEFAULT_VISIBLE_MODEL_TYPES = ("pretrained", "task-specific", "statistical")


@functools.lru_cache(maxsize=1)
def get_model_metadata() -> dict[str, dict]:
    """Display metadata for each model, keyed by the `model_name` used in the results."""
    df = pd.read_csv(MODELS_CSV)
    return {row["model_name"]: row for row in df.to_dict("records")}


def validate_model_metadata(model_names) -> None:
    """Raise if any displayed model lacks metadata, so the app never shows blank attributes."""
    metadata = get_model_metadata()
    missing = sorted({m for m in model_names if m not in metadata})
    if missing:
        raise ValueError(
            f"Missing metadata (organization, link, type) for models: {missing}. "
            "Add an entry to benchmarks/fev_bench/models.yaml in the fev repo, "
            "then regenerate the tables with `uv run python save_tables.py`."
        )


def select_model_names(model_types, include_non_commercial: bool = True) -> list[str]:
    """Names of the models with a type in `model_types`, optionally excluding non-commercial ones."""
    model_types = set(model_types)
    return sorted(
        name
        for name, entry in get_model_metadata().items()
        if entry["model_type"] in model_types and (include_non_commercial or entry["commercial_use"])
    )


ALL_METRICS = {
    "SQL": (
        "SQL: Scaled Quantile Loss",
        "The [Scaled Quantile Loss (SQL)](https://auto.gluon.ai/dev/tutorials/timeseries/forecasting-metrics.html#autogluon.timeseries.metrics.SQL) is a **scale-invariant** metric for evaluating **probabilistic** forecasts.",
    ),
    "MASE": (
        "MASE: Mean Absolute Scaled Error",
        "The [Mean Absolute Scaled Error (MASE)](https://auto.gluon.ai/dev/tutorials/timeseries/forecasting-metrics.html#autogluon.timeseries.metrics.MASE) is a **scale-invariant** metric for evaluating **point** forecasts.",
    ),
    "WQL": (
        "WQL: Weighted Quantile Loss",
        "The [Weighted Quantile Loss (WQL)](https://auto.gluon.ai/dev/tutorials/timeseries/forecasting-metrics.html#autogluon.timeseries.metrics.WQL), is a **scale-dependent** metric for evaluating **probabilistic** forecasts.",
    ),
    "WAPE": (
        "WAPE: Weighted Absolute Percentage Error",
        "The [Weighted Absolute Percentage Error (WAPE)](https://auto.gluon.ai/dev/tutorials/timeseries/forecasting-metrics.html#autogluon.timeseries.metrics.WAPE) is a **scale-dependent** metric for evaluating **point** forecasts.",
    ),
}


def format_metric_name(metric_name: str):
    return ALL_METRICS[metric_name][0]


def get_metric_description(metric_name: str):
    return ALL_METRICS[metric_name][1]


def get_model_link(model_name):
    return get_model_metadata()[model_name]["url"]


def get_model_organization(model_name):
    return get_model_metadata()[model_name]["organization"]


def get_model_type(model_name):
    return get_model_metadata()[model_name]["model_type"]


def get_model_type_label(model_name):
    return MODEL_TYPE_LABELS[get_model_type(model_name)]


def get_model_display_name(model_name):
    """Display name declared in models.yaml, falling back to the name used in the results."""
    entry = get_model_metadata().get(model_name)
    return entry["display_name"] if entry else model_name


def get_model_name_link(model_name):
    """Model URL with the display name as fragment, rendered as a named link by the page's LinkColumn."""
    return f"{get_model_link(model_name)}#{get_model_display_name(model_name)}"


def get_model_tags(model_name) -> list[str]:
    """Tags describing a model: its type, whether it is zero-shot, and whether its license is non-commercial."""
    entry = get_model_metadata()[model_name]
    tags = [MODEL_TYPE_LABELS[entry["model_type"]]]
    if entry["zero_shot"]:
        tags.append(ZERO_SHOT_TAG)
    if not entry["commercial_use"]:
        tags.append(NON_COMMERCIAL_TAG)
    return tags


def format_leaderboard(df: pd.DataFrame):
    df = df.copy()
    df["skill_score"] = df["skill_score"].round(1)
    df["win_rate"] = df["win_rate"].round(1)
    # Leakage only applies to zero-shot models; report 0 for models trained on each task
    is_zero_shot = df["model_name"].map(lambda m: bool(get_model_metadata()[m]["zero_shot"]))
    df["training_corpus_overlap"] = (
        (df["training_corpus_overlap"] * 100).round().astype(int).where(is_zero_shot, 0)
    )
    # Model type, zero-shot and licensing are folded into a single column of tags to keep the table
    # narrow. Non-commercial is only shown as the exception.
    df["tags"] = df["model_name"].map(get_model_tags)
    df["org"] = df["model_name"].apply(get_model_organization)
    df["model_name"] = df["model_name"].apply(get_model_name_link)
    df = df[
        [
            "model_name",
            "win_rate",
            "skill_score",
            "median_inference_time_s_per100",
            "org",
            "tags",
            "training_corpus_overlap",
            "num_failures",
        ]
    ]
    return df.style.apply(
        lambda x: [
            "background-color: #f8f9fa" if i % 2 == 1 else "" for i in range(len(x))
        ],
        axis=0,
    )


def construct_bar_chart(df: pd.DataFrame, col: str, metric_name: str):
    label = "Skill Score" if col == "skill_score" else "Win Rate"

    tooltip = [
        alt.Tooltip("model_name:N"),
        alt.Tooltip(f"{col}:Q", format=".2f"),
        alt.Tooltip(f"{col}_lower:Q", title="95% CI Lower", format=".2f"),
        alt.Tooltip(f"{col}_upper:Q", title="95% CI Upper", format=".2f"),
    ]

    base_encode = {
        "y": alt.Y("model_name:N", title="Forecasting Model", sort=None),
        "tooltip": tooltip,
    }

    bars = (
        alt.Chart(df)
        .mark_bar(color=COLORS["bar_fill"], cornerRadius=4)
        .encode(
            x=alt.X(f"{col}:Q", title=f"{label} (%)", scale=alt.Scale(zero=False)),
            **base_encode,
        )
    )

    error_bars = (
        alt.Chart(df)
        .mark_errorbar(ticks={"height": 5}, color=COLORS["error_bar"])
        .encode(
            y=alt.Y("model_name:N", title=None, sort=None),
            x=alt.X(f"{col}_lower:Q", title=f"{label} (%)"),
            x2=alt.X2(f"{col}_upper:Q"),
            tooltip=tooltip,
        )
    )

    points = (
        alt.Chart(df)
        .mark_point(filled=True, color=COLORS["point"])
        .encode(x=alt.X(f"{col}:Q", title=f"{label} (%)"), **base_encode)
    )

    return (
        (bars + error_bars + points)
        .properties(height=500, title=f"{label} ({metric_name}) with 95% CIs")
        .configure_title(fontSize=16)
    )


def construct_pairwise_chart(df: pd.DataFrame, col: str, metric_name: str):
    # Color ranges kept in sync with the fev repo's generate_fev_bench_figures.py
    config = {
        "win_rate": ("Win Rate", [0, 100], 50, f"abs(datum.{col} - 50) > 30"),
        "skill_score": ("Skill Score", [-15, 15], 0, f"abs(datum.{col}) > 10"),
    }
    cbar_label, domain, domain_mid, text_condition = config[col]

    df = df.copy()
    for c in [col, f"{col}_lower", f"{col}_upper"]:
        df[c] *= 100

    df["model_1"] = df["model_1"].apply(get_model_display_name)
    df["model_2"] = df["model_2"].apply(get_model_display_name)

    model_order = (
        df.groupby("model_1")[col].mean().sort_values(ascending=False).index.tolist()
    )

    tooltip = [
        alt.Tooltip("model_1:N", title="Model 1"),
        alt.Tooltip("model_2:N", title="Model 2"),
        alt.Tooltip(f"{col}:Q", title=cbar_label.split(" ")[0], format=".1f"),
        alt.Tooltip(f"{col}_lower:Q", title="95% CI Lower", format=".1f"),
        alt.Tooltip(f"{col}_upper:Q", title="95% CI Upper", format=".1f"),
    ]

    base = alt.Chart(df).encode(
        x=alt.X(
            "model_2:N",
            sort=model_order,
            title="Model 2",
            axis=alt.Axis(orient="top", labelAngle=-90),
        ),
        y=alt.Y("model_1:N", sort=model_order, title="Model 1"),
    )

    heatmap = base.mark_rect().encode(
        color=alt.Color(
            f"{col}:Q",
            legend=None,
            scale=alt.Scale(
                scheme=HEATMAP_COLOR_SCHEME,
                domain=domain,
                domainMid=domain_mid,
                clamp=True,
            ),
        ),
        tooltip=tooltip,
    )

    text_main = base.mark_text(dy=-8, fontSize=8, baseline="top", yOffset=5).encode(
        text=alt.Text(f"{col}:Q", format=".1f"),
        color=alt.condition(
            text_condition,
            alt.value(COLORS["text_white"]),
            alt.value(COLORS["text_black"]),
        ),
        tooltip=tooltip,
    )

    return (
        (heatmap + text_main)
        .properties(
            height=550,
            title={
                "text": f"Pairwise {cbar_label} ({metric_name}) with 95% CIs",
                "fontSize": 16,
            },
        )
        .configure_axis(labelFontSize=11, titleFontSize=13, titleFontWeight="bold")
        .resolve_scale(color="independent")
    )


def construct_pivot_table_from_df(
    errors: pd.DataFrame, metric_name: str
) -> pd.io.formats.style.Styler:
    """Construct styled pivot table from precomputed DataFrame."""

    def highlight_by_position(styler):
        rank_colors = {1: COLORS["gold"], 2: COLORS["silver"], 3: COLORS["bronze"]}

        for row_idx in errors.index:
            row_ranks = errors.loc[row_idx].rank(method="min")
            for col_idx in errors.columns:
                rank = row_ranks[col_idx]
                style_parts = []

                # Rank background colors
                if rank <= 3:
                    style_parts.append(f"background-color: {rank_colors[rank]}")
                else:
                    style_parts.append(f"color: {COLORS['text_default']}")

                if style_parts:
                    styler = styler.map(
                        lambda x, s="; ".join(style_parts): s,
                        subset=pd.IndexSlice[row_idx:row_idx, col_idx:col_idx],
                    )
        return styler

    return highlight_by_position(errors.style).format(precision=3)


def construct_pivot_table(
    summaries: pd.DataFrame,
    metric_name: str,
    baseline_model: str,
    leakage_imputation_model: str,
) -> pd.io.formats.style.Styler:
    errors = fev.pivot_table(
        summaries=summaries, metric_column=metric_name, task_columns=["task_name"]
    )
    train_overlap = (
        fev.pivot_table(
            summaries=summaries,
            metric_column="trained_on_this_dataset",
            task_columns=["task_name"],
        )
        .fillna(False)
        .astype(bool)
    )

    is_imputed_baseline = errors.isna()
    is_leakage_imputed = train_overlap

    # Handle imputations
    errors = errors.mask(train_overlap, errors[leakage_imputation_model], axis=0)
    for col in errors.columns:
        if col != baseline_model:
            errors[col] = errors[col].fillna(errors[baseline_model])

    errors = errors[errors.rank(axis=1).mean().sort_values().index]
    errors.index.rename("Task name", inplace=True)

    def highlight_by_position(styler):
        rank_colors = {1: COLORS["gold"], 2: COLORS["silver"], 3: COLORS["bronze"]}

        for row_idx in errors.index:
            row_ranks = errors.loc[row_idx].rank(method="min")
            for col_idx in errors.columns:
                rank = row_ranks[col_idx]
                style_parts = []

                # Rank background colors
                if rank <= 3:
                    style_parts.append(f"background-color: {rank_colors[rank]}")

                # Imputation text colors
                if is_leakage_imputed.loc[row_idx, col_idx]:
                    style_parts.append(f"color: {COLORS['leakage_impute']}")
                elif is_imputed_baseline.loc[row_idx, col_idx]:
                    style_parts.append(f"color: {COLORS['failure_impute']}")
                elif not style_parts or (
                    len(style_parts) == 1 and "font-weight" in style_parts[0]
                ):
                    style_parts.append(f"color: {COLORS['text_default']}")

                if style_parts:
                    styler = styler.map(
                        lambda x, s="; ".join(style_parts): s,
                        subset=pd.IndexSlice[row_idx:row_idx, col_idx:col_idx],
                    )
        return styler

    return highlight_by_position(errors.style).format(precision=3)