Generating stats from arena website.

arena_diag.py

from __future__ import annotations

import argparse
from datetime import datetime
from datetime import timedelta

import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from datasets import load_dataset
from matplotlib.lines import Line2D


DATASET_NAME = "lmarena-ai/leaderboard-dataset"
DATASET_SPLIT = "full"

# Provider metadata with canonical display names and colors.
PROVIDER_STYLE: dict[str, tuple[str, str]] = {
    "01 ai": ("01 AI", "#E54882"),
    "ai21 labs": ("AI21 Labs", "#D23F7A"),
    "alibaba": ("Alibaba", "#E67E22"),
    "allenai/uw": ("Allen Institute for AI", "#E24CA0"),
    "anthropic": ("Anthropic", "#C97B5A"),
    "amazon": ("Amazon", "#F28B18"),
    "baidu": ("Baidu", "#2E44B8"),
    "bytedance": ("ByteDance Seed", "#3C6DD4"),
    "cohere": ("Cohere", "#BA86D8"),
    "cognitive computations": ("Cognitive Computations", "#3D4652"),
    "deepseek": ("DeepSeek", "#2F4BE0"),
    "google": ("Google", "#3CA64C"),
    "huggingface": ("HuggingFace", "#F5A623"),
    "kimi": ("Kimi", "#2D7BE7"),
    "lg ai research": ("LG AI Research", "#D83A8C"),
    "mbzuai institute of foundation models": (
        "MBZUAI Institute of Foundation Models",
        "#2735B5",
    ),
    "meituan": ("LongCat", "#18A05A"),
    "meta": ("Meta", "#1F8CEB"),
    "microsoft azure": ("Microsoft Azure", "#0E5A8A"),
    "minimax": ("MiniMax", "#E63B73"),
    "mistral": ("Mistral", "#F07C12"),
    "moonshot": ("Moonshot", "#2EA3D8"),
    "nexusflow": ("Nexus Research", "#12A57B"),
    "nousresearch": ("Nous Research", "#0EA854"),
    "nvidia": ("NVIDIA", "#76B82A"),
    "openchat": ("OpenChat", "#6B79D8"),
    "openai": ("OpenAI", "#222222"),
    "reka ai": ("Reka AI", "#2E3F66"),
    "stepfun": ("StepFun", "#1F78E0"),
    "tencent": ("Tencent", "#4C89CC"),
    "uc berkeley": ("UC Berkeley", "#4D6A95"),
    "upstage": ("Upstage", "#6B55E0"),
    "upstage ai": ("Upstage", "#6B55E0"),
    "xai": ("xAI", "#6C63FF"),
    "xiaomi": ("Xiaomi", "#F37021"),
    "z ai": ("Z AI", "#1560D8"),
    "zai": ("Z AI", "#1560D8"),
    "zhipu": ("Z AI", "#1560D8"),
    "zhipu ai": ("Z AI", "#1560D8"),
    "lmsys": ("LMSYS", "#939AA1"),
    "uw": ("UW", "#8A929A"),
    "microsoft": ("Microsoft", "#8E949B"),
}

DEFAULT_PROVIDER_COLOR = "#9AA0A6"
DEFAULT_PROVIDER_NAME = "Unknown"

LEGEND_ORDER = [
    "alibaba",
    "anthropic",
    "deepseek",
    "google",
    "kimi",
    "lg ai research",
    "mbzuai institute of foundation models",
    "meta",
    "minimax",
    "mistral",
    "openai",
    "upstage",
    "xai",
    "xiaomi",
    "z ai",
    "lmsys",
    "uw",
    "microsoft",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config",
        default="text",
        help="Dataset configuration (ex: text, text_style_control).",
    )
    parser.add_argument(
        "--category",
        default="overall",
        help="Categorie a filtrer (par defaut: overall).",
    )
    parser.add_argument(
        "--output",
        default="best_model_elo_over_time.png",
        help="Output image path.",
    )
    parser.add_argument(
        "--min-rating",
        type=float,
        default=1050.0,
        help="Hide all scores strictly below this rating.",
    )
    parser.add_argument(
        "--top-rank-limit",
        type=int,
        default=20,
        help=(
            "Include additional competitor points up to this rank "
            "(excluding the top model for each day)."
        ),
    )
    parser.add_argument(
        "--label-layout",
        choices=["ordered", "no_cross"],
        default="ordered",
        help=(
            "Right label layout: 'ordered' keeps timeline ordering; "
            "'no_cross' reorders labels to avoid crossing connector lines."
        ),
    )
    return parser.parse_args()


def to_date(value: str) -> datetime.date:
    # Supporte les dates au format YYYY-MM-DD et YYYY-MM-DD HH:MM:SS.
    raw = value.strip()
    if " " in raw:
        raw = raw.split(" ", maxsplit=1)[0]
    return datetime.strptime(raw, "%Y-%m-%d").date()


def normalize_provider(name: str | None) -> str:
    if not name:
        return "unknown"
    normalized = " ".join(name.strip().lower().replace("-", " ").split())
    aliases = {
        "x ai": "xai",
        "mbzuai": "mbzuai institute of foundation models",
        "lg": "lg ai research",
        "lmsys org": "lmsys",
        "zai": "z ai",
        "zhipu": "z ai",
        "zhipu ai": "z ai",
        "upstage ai": "upstage",
        "allenai uw": "allenai/uw",
    }
    return aliases.get(normalized, normalized)


def provider_color(name: str | None) -> str:
    key = normalize_provider(name)
    return PROVIDER_STYLE.get(key, (DEFAULT_PROVIDER_NAME, DEFAULT_PROVIDER_COLOR))[1]


def provider_display(name: str | None) -> str:
    key = normalize_provider(name)
    return PROVIDER_STYLE.get(key, (name or DEFAULT_PROVIDER_NAME, DEFAULT_PROVIDER_COLOR))[0]


def short_model_name(name: str, max_len: int = 20) -> str:
    if len(name) <= max_len:
        return name
    return f"{name[: max_len - 1]}..."


def count_segment_crossings(
    source_xs: list[float],
    source_ys: list[float],
    target_x: float,
    target_ys: list[float],
) -> int:
    """Count pairwise crossings for straight segments from sources to a common target x."""
    n = len(source_xs)
    crossings = 0
    eps = 1e-9

    for i in range(n):
        for j in range(i + 1, n):
            x1, y1, ty1 = source_xs[i], source_ys[i], target_ys[i]
            x2, y2, ty2 = source_xs[j], source_ys[j], target_ys[j]

            overlap_start = max(x1, x2)
            if overlap_start >= target_x - eps:
                continue

            denom1 = target_x - x1
            denom2 = target_x - x2
            if abs(denom1) < eps or abs(denom2) < eps:
                continue

            def y_at(x: float, sx: float, sy: float, ty: float) -> float:
                t = (x - sx) / (target_x - sx)
                return sy + (ty - sy) * t

            d_start = y_at(overlap_start, x1, y1, ty1) - y_at(overlap_start, x2, y2, ty2)
            d_end = ty1 - ty2

            if abs(d_start) < eps or abs(d_end) < eps or d_start * d_end < 0:
                crossings += 1

    return crossings


def optimize_no_cross_order(
    source_xs: list[float],
    source_ys: list[float],
    target_x: float,
    lane_ys: list[float],
) -> tuple[list[int], int]:
    """Greedy swap optimization of label assignment to reduce connector crossings."""
    n = len(source_xs)
    order = sorted(range(n), key=lambda i: source_ys[i], reverse=True)

    def score(current_order: list[int]) -> int:
        target_ys = [0.0] * n
        for lane_idx, label_idx in enumerate(current_order):
            target_ys[label_idx] = lane_ys[lane_idx]
        return count_segment_crossings(source_xs, source_ys, target_x, target_ys)

    best_score = score(order)
    improved = True
    while improved and best_score > 0:
        improved = False
        best_candidate = order
        best_candidate_score = best_score

        for i in range(n - 1):
            for j in range(i + 1, n):
                candidate = order.copy()
                candidate[i], candidate[j] = candidate[j], candidate[i]
                s = score(candidate)
                if s < best_candidate_score:
                    best_candidate = candidate
                    best_candidate_score = s

        if best_candidate_score < best_score:
            order = best_candidate
            best_score = best_candidate_score
            improved = True

    return order, best_score


def main() -> None:
    args = parse_args()

    ds = load_dataset(
        DATASET_NAME,
        args.config,
        split=DATASET_SPLIT,
        filters=[("category", "==", args.category)],
        columns=[
            "model_name",
            "organization",
            "rating",
            "rank",
            "leaderboard_publish_date",
        ],
    )

    rows: list[tuple[datetime.date, str, str | None, float, float]] = []
    best_by_day: dict[datetime.date, tuple[str, str | None, float]] = {}
    for row in ds:
        day = to_date(row["leaderboard_publish_date"])
        rating = float(row["rating"])
        model_name = row["model_name"]
        organization = row.get("organization")
        rank = float(row["rank"])

        rows.append((day, model_name, organization, rating, rank))

        current = best_by_day.get(day)
        if current is None or rating > current[2]:
            best_by_day[day] = (model_name, organization, rating)

    if not best_by_day:
        raise RuntimeError("Aucune donnee trouvee avec les filtres fournis.")

    dates = sorted(best_by_day)
    best_models = [best_by_day[d][0] for d in dates]
    best_orgs = [best_by_day[d][1] for d in dates]
    best_ratings = [best_by_day[d][2] for d in dates]

    filtered_triplets = [
        (d, m, o, r)
        for d, m, o, r in zip(dates, best_models, best_orgs, best_ratings)
        if r >= args.min_rating
    ]
    if not filtered_triplets:
        raise RuntimeError(
            f"No scores >= {args.min_rating:.0f} for the selected dataset/config/category."
        )

    dates = [t[0] for t in filtered_triplets]
    best_models = [t[1] for t in filtered_triplets]
    best_orgs = [t[2] for t in filtered_triplets]
    best_ratings = [t[3] for t in filtered_triplets]

    print("Date        | Best model                    | Provider             | Rank")
    print("-" * 74)
    for d, m, o, r in zip(dates, best_models, best_orgs, best_ratings):
        provider = provider_display(o)[:20]
        print(f"{d} | {m[:28]:28s} | {provider:20s} | {r:7.2f}")

    fig, ax = plt.subplots(figsize=(16, 9), dpi=220)
    fig.patch.set_facecolor("#F7F6F3")
    ax.set_facecolor("#FCFCFB")

    for i in range(1, len(dates)):
        segment_color = provider_color(best_orgs[i])
        ax.plot(
            dates[i - 1 : i + 1],
            best_ratings[i - 1 : i + 1],
            color=segment_color,
            linewidth=2.8,
            alpha=0.95,
        )

    # Add top-N competitors (excluding daily leader) as points only.
    other_dates: list[datetime.date] = []
    other_ratings: list[float] = []
    other_colors: list[str] = []
    other_orgs: list[str | None] = []
    for day, model_name, organization, rating, rank in rows:
        if day not in best_by_day:
            continue
        if (
            rank <= args.top_rank_limit
            and rating >= args.min_rating
            and model_name != best_by_day[day][0]
        ):
            other_dates.append(day)
            other_ratings.append(rating)
            other_colors.append(provider_color(organization))
            other_orgs.append(organization)

    ax.scatter(
        other_dates,
        other_ratings,
        c=other_colors,
        s=12,
        alpha=0.35,
        linewidths=0,
        zorder=2,
    )

    # A point is shown only when the top model changes compared to previous date.
    change_indices = [0]
    for i in range(1, len(dates)):
        if best_models[i] != best_models[i - 1]:
            change_indices.append(i)

    change_dates = [dates[i] for i in change_indices]
    change_ratings = [best_ratings[i] for i in change_indices]
    change_models = [best_models[i] for i in change_indices]
    change_colors = [provider_color(best_orgs[i]) for i in change_indices]

    ax.scatter(
        change_dates,
        change_ratings,
        c=change_colors,
        s=40,
        edgecolors="#FFFFFF",
        linewidths=0.8,
        zorder=4,
    )

    # Annotate all change points in a right-side label lane to avoid overlap.
    if change_indices:
        lane_x = max(dates) + timedelta(days=90)
        lane_end = lane_x + timedelta(days=35)
        ax.set_xlim(min(dates) - timedelta(days=20), lane_end)

        min_r = min(change_ratings)
        max_r = max(change_ratings)
        pad_base = max(10.0, (max_r - min_r) * 0.08)

        if args.label_layout == "ordered":
            # Reverse chronological order (latest first), i.e. current visual ordering.
            sorted_for_labels = sorted(change_indices, key=lambda idx: dates[idx], reverse=True)
            pad = pad_base
            if len(sorted_for_labels) == 1:
                lane_ys = [best_ratings[sorted_for_labels[0]]]
            else:
                step = (max_r - min_r + 2 * pad) / (len(sorted_for_labels) - 1)
                lane_ys = [max_r + pad - i * step for i in range(len(sorted_for_labels))]
        else:
            # Reorder labels to minimize connector crossings regardless of semantic order.
            date_nums = mdates.date2num(change_dates)
            lane_x_num = mdates.date2num(lane_x)
            n_labels = len(change_indices)

            if n_labels == 1:
                sorted_for_labels = [change_indices[0]]
                lane_ys = [best_ratings[change_indices[0]]]
            else:
                best_order: list[int] | None = None
                best_lane_ys: list[float] | None = None
                best_crossings = float("inf")

                # Expand vertical spread if needed to reach a crossing-free assignment.
                for spread in (1.0, 1.25, 1.5, 2.0, 2.5, 3.0):
                    pad = pad_base * spread
                    step = (max_r - min_r + 2 * pad) / (n_labels - 1)
                    candidate_lane_ys = [max_r + pad - i * step for i in range(n_labels)]

                    order_in_change, crossings = optimize_no_cross_order(
                        source_xs=list(date_nums),
                        source_ys=list(change_ratings),
                        target_x=lane_x_num,
                        lane_ys=candidate_lane_ys,
                    )

                    if crossings < best_crossings:
                        best_crossings = crossings
                        best_order = order_in_change
                        best_lane_ys = candidate_lane_ys

                    if crossings == 0:
                        break

                assert best_order is not None
                assert best_lane_ys is not None
                sorted_for_labels = [change_indices[i] for i in best_order]
                lane_ys = best_lane_ys

        for idx, lane_y in zip(sorted_for_labels, lane_ys):
            d = dates[idx]
            r = best_ratings[idx]
            model = short_model_name(best_models[idx], max_len=24)

            ax.annotate(
                f"{model}",
                xy=(d, r),
                xytext=(lane_x, lane_y),
                textcoords="data",
                fontsize=8,
                color="#111827",
                ha="left",
                va="center",
                bbox={"boxstyle": "round,pad=0.22", "fc": "white", "ec": "none", "alpha": 0.98},
                arrowprops={
                    "arrowstyle": "-",
                    "color": provider_color(best_orgs[idx]),
                    "lw": 0.8,
                    "alpha": 0.85,
                    "connectionstyle": "arc3,rad=0",
                },
            )

    fig.suptitle(
        f"Top Rank Over Time ({args.config}, {args.category})",
        x=0.07,
        y=0.96,
        ha="left",
        fontsize=28,
        fontweight="bold",
        color="#1F2937",
    )
    ax.set_xlabel("Leaderboard publish date", fontsize=12)
    ax.set_ylabel("Top model rank", fontsize=12)
    ax.grid(alpha=0.22, color="#94A3B8", linewidth=0.8)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

    locator = mdates.AutoDateLocator(minticks=6, maxticks=10)
    formatter = mdates.ConciseDateFormatter(locator)
    ax.xaxis.set_major_locator(locator)
    ax.xaxis.set_major_formatter(formatter)

    # Legend includes all providers that have at least one visible point.
    shown: dict[str, tuple[str, str]] = {}
    for org in [*best_orgs, *other_orgs]:
        key = normalize_provider(org)
        if key not in shown:
            shown[key] = (provider_display(org), provider_color(org))

    ordered_keys = [k for k in LEGEND_ORDER if k in shown] + [
        k for k in shown if k not in LEGEND_ORDER
    ]
    legend_handles = [
        Line2D([0], [0], color=shown[key][1], lw=3, label=shown[key][0])
        for key in ordered_keys
    ]
    if legend_handles:
        # Put legend below the chart to avoid covering title or plot area.
        ax.legend(
            handles=legend_handles,
            loc="upper center",
            bbox_to_anchor=(0.5, -0.16),
            ncol=8,
            fontsize=9,
            frameon=False,
        )

    fig.subplots_adjust(top=0.84, bottom=0.29, left=0.07, right=0.90)
    
    # save to image
    plt.savefig(args.output, dpi=360)


if __name__ == "__main__":
    main()