app.py

import os

import gradio as gr
import pandas as pd
from apscheduler.schedulers.background import BackgroundScheduler
from huggingface_hub import HfApi

from src.about import CITATION_TEXT, INTRODUCTION_TEXT, LLM_BENCHMARKS_TEXT, TITLE
from src.populate import (
    MU_MATH_COLUMNS_DICT,
    U_MATH_COLUMNS_DICT,
    Field,
    get_mu_math_leaderboard_df,
    get_u_math_leaderboard_df,
)


def restart_space():
    TOKEN = os.environ.get("HF_TOKEN")  # A read/write token for your org
    API = HfApi(token=TOKEN)
    REPO_ID = "toloka/u-math-leaderboard"
    API.restart_space(repo_id=REPO_ID)


LEADERBOARD_U_MATH_DF = get_u_math_leaderboard_df()
LEADERBOARD_MU_MATH_DF = get_mu_math_leaderboard_df()


def init_leaderboard(dataframe: pd.DataFrame, columns_dict: dict[str, Field]) -> gr.components.Component:
    if dataframe is None or dataframe.empty:
        raise ValueError("Leaderboard DataFrame is empty or None.")

    def filter_dataframe_by_selected_columns(full_df: pd.DataFrame, columns: list[str]) -> pd.DataFrame:
        always_here_cols = [c.pretty_name for c in columns_dict.values() if c.never_hidden]
        selected_columns = [c for c in columns if c in full_df.columns and c not in always_here_cols]
        # keep the order of the columns
        filtered_df = full_df[[c for c in full_df.columns if c in (always_here_cols + selected_columns)]]
        return filtered_df

    def filter_dataframe_by_selected_tag_columns(
        full_df: pd.DataFrame, current_tag: str
    ) -> tuple[pd.DataFrame, list[str], str, str]:
        always_here_cols = [c.pretty_name for c in columns_dict.values() if c.never_hidden]
        selected_columns = [
            c.pretty_name for c in columns_dict.values() if current_tag in c.tags and c not in always_here_cols
        ]
        # keep the order of the columns
        filtered_df = full_df[[c for c in full_df.columns if c in (always_here_cols + selected_columns)]]
        _columns_to_select_visibility = [
            c.pretty_name for c in columns_dict.values() if not c.fully_hidden and not c.never_hidden
        ]
        return filtered_df, [c for c in _columns_to_select_visibility if c in filtered_df.columns], "All", "All"

    def filter_dataframe_by_search(full_df: pd.DataFrame, current_df: pd.DataFrame, search: str) -> pd.DataFrame:
        filtered_df = full_df[
            (full_df[columns_dict["model_name"].pretty_name].str.contains(search, case=False, na=False))
        ]
        return filtered_df[current_df.columns]

    def filter_dataframe_by_model_type(
        full_df: pd.DataFrame, current_df: pd.DataFrame, filter_name: str
    ) -> pd.DataFrame:
        if filter_name == "All":
            return full_df[current_df.columns]
        else:
            # actually filter by emoji
            query_symbol = filter_name[0]
            filtered_df = full_df[full_df[columns_dict["model_type_symbol"].pretty_name] == query_symbol]
            return filtered_df[current_df.columns]

    def filter_dataframe_by_model_size(
        full_df: pd.DataFrame, current_df: pd.DataFrame, filter_name: str
    ) -> pd.DataFrame:
        if filter_name == "All":
            return full_df[current_df.columns]
        else:
            # actually filter by emoji
            query_symbol = filter_name[0]
            filtered_df = full_df[full_df[columns_dict["model_size_symbol"].pretty_name] == query_symbol]
            return filtered_df[current_df.columns]

    with gr.Column() as col:
        # Add the controls
        with gr.Accordion("➡️ See All Columns", open=False):
            columns_to_select_visibility = [
                c.pretty_name for c in columns_dict.values() if not c.fully_hidden and not c.never_hidden
            ]
            all_columns_selector = gr.CheckboxGroup(
                choices=columns_to_select_visibility,
                value=[
                    c.pretty_name
                    for c in columns_dict.values()
                    if c.pretty_name in columns_to_select_visibility and c.displayed_by_default
                ],
                label="Select Columns to Display:",
                interactive=True,
                container=False,
            )
        with gr.Row():
            with gr.Column():
                search_bar = gr.Textbox(
                    placeholder="🔍 Search for your model and press ENTER...",
                    show_label=False,
                    elem_id="search-bar",
                )

            # collect all column tags and create buttons for them
            all_tags = {}
            with gr.Column(variant="panel"):
                gr.Markdown("Select Columns:")
                for c in columns_dict.values():
                    for tag in c.tags:
                        if tag not in all_tags:
                            all_tags[tag] = gr.Button(tag, interactive=True, size="sm")

            model_type_filter_selector = gr.Radio(
                label="Filter model types:",
                choices=["All", "💙 Open-Weights", "🟥 Proprietary"],
                value="All",
                elem_id="model-type-filter",
                interactive=True,
            )

            model_size_filter_selector = gr.Radio(
                label="Filter model sizes:",
                choices=["All", "🛴 Tiny (<5B)", "🚗 Small (5-50B)", "🚚 Medium (50-100B)", "🚀 Large (>100B)"],
                value="All",
                elem_id="model-size-filter",
                interactive=True,
            )

        # create the hidden and visible dataframes to display
        hidden_leaderboard_df = gr.components.Dataframe(
            value=dataframe,
            datatype=[c.column_type for c in columns_dict.values()],
            visible=False,
            interactive=False,
        )
        leaderboard_df = gr.components.Dataframe(
            value=dataframe[[c.pretty_name for c in columns_dict.values() if c.displayed_by_default]],
            datatype=[c.column_type for c in columns_dict.values()],
            elem_id="leaderboard-df",
            interactive=False,
        )
        # add the callbacks
        all_columns_selector.change(
            fn=filter_dataframe_by_selected_columns,
            inputs=[hidden_leaderboard_df, all_columns_selector],
            outputs=[leaderboard_df],
        )
        search_bar.submit(
            fn=filter_dataframe_by_search,
            inputs=[hidden_leaderboard_df, leaderboard_df, search_bar],
            outputs=[leaderboard_df],
        )
        model_type_filter_selector.change(
            fn=filter_dataframe_by_model_type,
            inputs=[hidden_leaderboard_df, leaderboard_df, model_type_filter_selector],
            outputs=[leaderboard_df],
        )
        model_size_filter_selector.change(
            fn=filter_dataframe_by_model_size,
            inputs=[hidden_leaderboard_df, leaderboard_df, model_size_filter_selector],
            outputs=[leaderboard_df],
        )
        for tag, button in all_tags.items():
            button.click(
                fn=filter_dataframe_by_selected_tag_columns,
                inputs=[hidden_leaderboard_df, button],
                outputs=[leaderboard_df, all_columns_selector, model_type_filter_selector, model_size_filter_selector],
            )

        # reload the leaderboard on the first load
        filter_dataframe_by_selected_columns(dataframe, all_columns_selector.value)
        return col


demo = gr.Blocks()
with demo:
    gr.HTML(TITLE)
    gr.Markdown(INTRODUCTION_TEXT, elem_classes="markdown-text")

    print(LEADERBOARD_U_MATH_DF)
    print(LEADERBOARD_MU_MATH_DF)

    with gr.Tabs(elem_classes="tab-buttons") as tabs:
        with gr.TabItem("🏆 U-MATH", elem_id="u-math-benchmark-tab-table", id=0):
            leaderboard_umath = init_leaderboard(LEADERBOARD_U_MATH_DF, U_MATH_COLUMNS_DICT)

        with gr.TabItem("🏅 μ-MATH (Meta-Benchmark)", elem_id="mu-math-benchmark-tab-table", id=1):
            leaderboard_mumath = init_leaderboard(LEADERBOARD_MU_MATH_DF, MU_MATH_COLUMNS_DICT)

        with gr.TabItem("📝 About", elem_id="about-tab-table", id=2):
            gr.Markdown(LLM_BENCHMARKS_TEXT, elem_classes="markdown-text")

        citation_button = gr.Textbox(
            value=CITATION_TEXT,
            label="📙 Citation",
            lines=9,
            elem_id="citation-button",
            show_copy_button=True,
            container=True,
        )

scheduler = BackgroundScheduler()
scheduler.add_job(restart_space, "interval", seconds=60 * 60)
scheduler.start()
demo.queue(default_concurrency_limit=40).launch(ssr_mode=False)