app.py

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

from src.assets.text_content import TITLE, INTRODUCTION_TEXT, CITATION_BUTTON_LABEL, CITATION_BUTTON_TEXT
from src.assets.css_html_js import custom_css, get_window_url_params
from src.utils import restart_space, load_dataset_repo, make_clickable_model


LLM_PERF_LEADERBOARD_REPO = "optimum/llm-perf-leaderboard"
LLM_PERF_DATASET_REPO = "optimum/llm-perf-dataset"
OPTIMUM_TOKEN = os.environ.get("OPTIMUM_TOKEN")

OLD_COLUMNS = ["model", "backend.name", "backend.torch_dtype",
               "generate.latency(s)", "generate.throughput(tokens/s)"]

NEW_COLUMNS = ["Model", "Backend 🏭", "Load Datatype",
               "Latency (s) ⬇️", "Throughput (tokens/s) ⬆️"]

COLUMNS_DATATYPES = ["markdown", "str", "str", "number", "number"]

SORTING_COLUMN = ["Throughput (tokens/s) ⬆️"]


llm_perf_dataset_repo = load_dataset_repo(LLM_PERF_DATASET_REPO, OPTIMUM_TOKEN)


def get_benchmark_df(benchmark):
    if llm_perf_dataset_repo:
        llm_perf_dataset_repo.git_pull()

    # load
    df = pd.read_csv(
        f"./llm-perf-dataset/reports/{benchmark}/inference_report.csv")

    # preprocess
    df["model"] = df["model"].apply(make_clickable_model)

    # filter
    df = df[OLD_COLUMNS]

    # rename
    df.rename(columns={
        df_col: rename_col for df_col, rename_col in zip(OLD_COLUMNS, NEW_COLUMNS)
    }, inplace=True)

    # sort
    df.sort_values(by=SORTING_COLUMN, ascending=False, inplace=True)

    return df


# Define demo interface
demo = gr.Blocks(css=custom_css)
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("🖥️ A100-80GB Benchmark 🏋️", elem_id="A100-benchmark", id=0):

            SINGLE_A100_TEXT = """<h4>Specifications:</h4>
            - Single-GPU (1)
            - Singleton Batch (1)
            - Thousand Tokens (1000)"""
            gr.HTML(SINGLE_A100_TEXT)

            single_A100_df = get_benchmark_df(benchmark="1xA100-80GB")
            leaderboard_table_lite = gr.components.Dataframe(
                value=single_A100_df,
                datatype=COLUMNS_DATATYPES,
                headers=NEW_COLUMNS,
                elem_id="1xA100-table",
            )

            MULTI_A100_TEXT = """<h4>Specifications:</h4>
            - Multi-GPU (4)
            - Singleton Batch (1)
            - Thousand Tokens (1000)"""
            gr.HTML(MULTI_A100_TEXT)

            multi_A100_df = get_benchmark_df(benchmark="4xA100-80GB")
            leaderboard_table_full = gr.components.Dataframe(
                value=multi_A100_df,
                datatype=COLUMNS_DATATYPES,
                headers=NEW_COLUMNS,
                elem_id="4xA100-table",
            )

    with gr.Row():
        with gr.Column():
            with gr.Accordion("📙 Citation", open=False):
                citation_button = gr.Textbox(
                    value=CITATION_BUTTON_TEXT,
                    label=CITATION_BUTTON_LABEL,
                    elem_id="citation-button",
                ).style(show_copy_button=True)

# Restart space every hour
scheduler = BackgroundScheduler()
scheduler.add_job(restart_space, "interval", seconds=3600,
                  args=[LLM_PERF_LEADERBOARD_REPO, OPTIMUM_TOKEN])
scheduler.start()

# Launch demo
demo.queue(concurrency_count=40).launch()