app.py

import json

import datasets
import gradio as gr
from apscheduler.schedulers.background import BackgroundScheduler
from huggingface_hub import snapshot_download
from loguru import logger

import populate
from about import LEADERBOARD_INTRODUCTION_TEXT, LEADERBOARD_TITLE
from app_configs import DEFAULT_SELECTIONS, THEME
from components.quizbowl.bonus import BonusInterface
from components.quizbowl.tossup import TossupInterface
from components.typed_dicts import PipelineInterfaceDefaults, TossupInterfaceDefaults
from display.css_html_js import fonts_header, js_head, leaderboard_css
from display.custom_css import css_bonus, css_pipeline, css_tossup
from display.guide import BUILDING_MARKDOWN, GUIDE_MARKDOWN, QUICKSTART_MARKDOWN
from display.utils import AutoEvalColumn, fields

# Constants
from envs import (
    API,
    EVAL_REQUESTS_PATH,
    EVAL_RESULTS_PATH,
    LEADERBOARD_REFRESH_INTERVAL,
    PLAYGROUND_DATASET_NAMES,
    QUEUE_REPO,
    REPO_ID,
    RESULTS_REPO,
    SERVER_REFRESH_INTERVAL,
)
from workflows import factory
from workflows.configs import AVAILABLE_MODELS


def restart_space():
    API.restart_space(repo_id=REPO_ID)


def download_dataset_snapshot(repo_id, local_dir):
    try:
        logger.info(f"Downloading dataset snapshot from {repo_id} to {local_dir}")
        snapshot_download(
            repo_id=repo_id,
            local_dir=local_dir,
            repo_type="dataset",
            tqdm_class=None,
        )
    except Exception as e:
        logger.error(f"Error downloading dataset snapshot from {repo_id} to {local_dir}: {e}. Restarting space.")
        restart_space()


download_dataset_snapshot(QUEUE_REPO, EVAL_REQUESTS_PATH)


def fetch_leaderboard_df():
    logger.info("Leaderboard fetched...")
    download_dataset_snapshot(RESULTS_REPO, EVAL_RESULTS_PATH)
    return populate.get_leaderboard_df(EVAL_RESULTS_PATH)


def load_dataset(mode: str):
    if mode == "tossup":
        ds = datasets.load_dataset(PLAYGROUND_DATASET_NAMES["tossup"], split="eval")
        ds = ds.filter(lambda x: x["qid"].split("-")[2] == "1" and int(x["qid"].split("-")[3]) <= 10)
    elif mode == "bonus":
        ds = datasets.load_dataset(PLAYGROUND_DATASET_NAMES["bonus"], split="eval")
        ds = ds.filter(lambda x: x["qid"].split("-")[2] == "1" and int(x["qid"].split("-")[3]) <= 10)
    else:
        raise ValueError(f"Invalid mode: {mode}")

    return ds


def get_default_tab_id(request: gr.Request):
    logger.info(f"Request: {request}")
    tab_key_value = request.query_params.get("tab", "tossup")
    return gr.update(selected=tab_key_value)


def presave_pipeline_state(
    login_btn,
    browser_state: dict,
    tossup_pipeline_state: dict,
    tossup_output_state: dict,
    bonus_pipeline_state: dict,
    bonus_output_state: dict,
):
    browser_state.setdefault("tossup", {})
    browser_state["tossup"]["pipeline_state"] = tossup_pipeline_state
    browser_state["tossup"]["output_state"] = tossup_output_state
    browser_state.setdefault("bonus", {})
    browser_state["bonus"]["pipeline_state"] = bonus_pipeline_state
    browser_state["bonus"]["output_state"] = bonus_output_state
    logger.debug(
        f"Pipeline state before login. Login button: {login_btn}, browser state: {json.dumps(browser_state, indent=4)}"
    )
    return login_btn, browser_state


if __name__ == "__main__":
    scheduler = BackgroundScheduler()
    scheduler.add_job(restart_space, "interval", seconds=SERVER_REFRESH_INTERVAL)
    scheduler.start()

    css = css_pipeline + css_tossup + css_bonus + leaderboard_css
    head = fonts_header + js_head
    tossup_ds = load_dataset("tossup")
    bonus_ds = load_dataset("bonus")
    with gr.Blocks(
        css=css,
        head=head,
        theme=THEME,
        title="Quizbowl Bot",
    ) as demo:
        browser_state = gr.BrowserState(
            {
                "tossup": {"pipeline_state": None, "output_state": None},
                "bonus": {"pipeline_state": None, "output_state": None},
            }
        )
        with gr.Row():
            with gr.Column(scale=5):
                gr.Markdown(
                    "## Welcome to Quizbowl Arena! \n### Create, play around, and submit your quizbowl agents.",
                    elem_classes="welcome-text",
                )
            login_btn = gr.LoginButton(scale=1)
        gr.Markdown(
            "**First time here?** Check out the [❓ Help](#help) tab for a quick introduction and the "
            "[walkthrough documentation](https://github.com/stanford-crfm/quizbowl-arena/blob/main/docs/walkthrough.md) "
            "for detailed examples and tutorials on how to create and compete with your own QuizBowl agents.",
            elem_classes="help-text",
        )
        with gr.Tabs() as gtab:
            with gr.Tab("🛎️ Tossup Agents", id="tossup"):
                defaults = TossupInterfaceDefaults(
                    **DEFAULT_SELECTIONS["tossup"], init_workflow=factory.create_simple_qb_tossup_workflow()
                )
                tossup_interface = TossupInterface(demo, browser_state, tossup_ds, AVAILABLE_MODELS, defaults)
            with gr.Tab("🙋🏻‍♂️ Bonus Round Agents", id="bonus"):
                defaults = PipelineInterfaceDefaults(
                    **DEFAULT_SELECTIONS["bonus"], init_workflow=factory.create_simple_qb_bonus_workflow()
                )
                bonus_interface = BonusInterface(demo, browser_state, bonus_ds, AVAILABLE_MODELS, defaults)
            with gr.Tab("🏅 Leaderboard", elem_id="llm-benchmark-tab-table", id="leaderboard"):
                leaderboard_timer = gr.Timer(LEADERBOARD_REFRESH_INTERVAL)
                gr.Markdown("<a id='leaderboard' href='#leaderboard'>QANTA Leaderboard</a>")
                gr.Markdown(LEADERBOARD_INTRODUCTION_TEXT)
                refresh_btn = gr.Button("🔄 Refresh")
                leaderboard_table = gr.Dataframe(
                    value=fetch_leaderboard_df,
                    every=leaderboard_timer,
                    headers=[c.name for c in fields(AutoEvalColumn)],
                    datatype=[c.type for c in fields(AutoEvalColumn)],
                    elem_id="leaderboard-table",
                    interactive=False,
                    visible=True,
                )
                refresh_btn.click(fn=fetch_leaderboard_df, inputs=[], outputs=leaderboard_table)
            with gr.Tab("❓ Help", id="help"):
                with gr.Row():
                    with gr.Column():
                        gr.Markdown(QUICKSTART_MARKDOWN)
                    with gr.Column():
                        gr.Markdown(BUILDING_MARKDOWN)

        # Event Listeners

        login_btn.click(
            fn=presave_pipeline_state,
            inputs=[
                login_btn,
                browser_state,
                tossup_interface.pipeline_state,
                tossup_interface.output_state,
                bonus_interface.pipeline_state,
                bonus_interface.output_state,
            ],
            outputs=[login_btn, browser_state],
        )

        demo.queue(default_concurrency_limit=40).launch()