main.py

from pathlib import Path

from apscheduler.schedulers.background import BackgroundScheduler
import pandas as pd
import gradio as gr
from gradio_leaderboard import Leaderboard, ColumnFilter

from constants import MethodTypes, model_type_emoji


TITLE = """<h1 align="center" id="space-title">TabArena: Public leaderboard for Tabular methods</h1>"""

INTRODUCTION_TEXT = ("TabArena Leaderboard measures the performance of tabular models on a collection of tabular "
                     "datasets manually curated. The datasets are collected to make sure they are tabular, with "
                     "permissive license without ethical issues and so on, we refer to the paper XXX for a "
                     "description of our approach.")

ABOUT_TEXT = f"""
## How It Works.

To evaluate the leaderboard, follow install instructions in 
`https://github.com/autogluon/tabrepo/tree/tabarena` and run 
`https://github.com/autogluon/tabrepo/blob/tabarena/examples/tabarena/run_tabarena_eval.py`.


This will generate a leaderboard. You can add your own method and contact the authors if you want it to be added
to the leaderboard. We require method to have public code available to be considered in the leaderboard. 
"""

CITATION_BUTTON_LABEL = "If you use this leaderboard in your research please cite the following:"
CITATION_BUTTON_TEXT = r"""
@article{
TODO update when arxiv version is ready,
}
"""


def get_model_family(model_name: str) -> str:
    prefixes_mapping = {
        MethodTypes.automl: ["AutoGluon"],
        MethodTypes.finetuned: ["REALMLP", "TabM", "FASTAI", "MNCA", "NN_TORCH"],
        MethodTypes.tree: ["GBM", "CAT", "EBM", "XGB"],
        MethodTypes.foundational: ["TABDPT", "TABICL", "TABPFN"],
        MethodTypes.baseline: ["KNN", "LR"]
    }
    for method_type, prefixes in prefixes_mapping.items():
        for prefix in prefixes:
            if prefix.lower() in model_name.lower():
                return method_type
    return MethodTypes.other


def load_data(filename: str):
    df_leaderboard = pd.read_csv(Path(__file__).parent / "data" / f"{filename}.csv.zip")
    print(f"Loaded dataframe with {len(df_leaderboard)} rows and columns {df_leaderboard.columns}")
    df_leaderboard["family"] = df_leaderboard.loc[:, "method"].apply(get_model_family)
    df_leaderboard["family"] = df_leaderboard.loc[:, "family"].apply(lambda s: s + " " + model_type_emoji[s])
    df_leaderboard = df_leaderboard.loc[:, ["method", "family", "time_train_s", "time_infer_s", "rank", "elo"]]
    df_leaderboard = df_leaderboard.round(1)
    df_leaderboard.rename(columns={
        "time_train_s": "training time (s)",
        "time_infer_s": "inference time (s)",
    }, inplace=True)
    return df_leaderboard


def make_leaderboard(df_leaderboard: pd.DataFrame) -> Leaderboard:
    return Leaderboard(
        value=df_leaderboard,
        search_columns=["method"],
        filter_columns=[
            # "method",
            ColumnFilter("family", type="dropdown", label="Filter by family"),
        ]
    )


def main():

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

        with gr.Tabs(elem_classes="tab-buttons") as tabs:
            with gr.TabItem('🏅 Overall', elem_id="llm-benchmark-tab-table", id=2):
                df_leaderboard = load_data("leaderboard-all")
                leaderboard = make_leaderboard(df_leaderboard)

            with gr.TabItem('🏅 Regression', elem_id="llm-benchmark-tab-table", id=0):
                df_leaderboard = load_data("leaderboard-regression")
                leaderboard = make_leaderboard(df_leaderboard)

            with gr.TabItem('🏅 Classification', elem_id="llm-benchmark-tab-table", id=1):
                df_leaderboard = load_data("leaderboard-classification")
                leaderboard = make_leaderboard(df_leaderboard)

            with gr.TabItem("📝 About", elem_id="llm-benchmark-tab-table", id=4):
                gr.Markdown(ABOUT_TEXT, elem_classes="markdown-text")
        with gr.Row():
            with gr.Accordion("📙 Citation", open=False):
                citation_button = gr.Textbox(
                    value=CITATION_BUTTON_TEXT,
                    label=CITATION_BUTTON_LABEL,
                    lines=20,
                    elem_id="citation-button",
                    show_copy_button=True,
                )

    scheduler = BackgroundScheduler()
    # scheduler.add_job(restart_space, "interval", seconds=1800)
    scheduler.start()
    demo.queue(default_concurrency_limit=40).launch()
    demo.launch()


if __name__ == "__main__":
    main()