app.py

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

from app_configs import AVAILABLE_MODELS, DEFAULT_SELECTIONS, THEME
from components.quizbowl.bonus import BonusInterface
from components.quizbowl.tossup import TossupInterface
from display.custom_css import css_bonus, css_pipeline, css_tossup
from display.guide import GUIDE_MARKDOWN

# Constants
from envs import (
    API,
    EVAL_REQUESTS_PATH,
    EVAL_RESULTS_PATH,
    PLAYGROUND_DATASET_NAMES,
    QUEUE_REPO,
    REPO_ID,
    RESULTS_REPO,
    TOKEN,
)
from workflows import factory


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


# Space initialisation
try:
    print(EVAL_REQUESTS_PATH)
    snapshot_download(
        repo_id=QUEUE_REPO,
        local_dir=EVAL_REQUESTS_PATH,
        repo_type="dataset",
        tqdm_class=None,
        etag_timeout=30,
        token=TOKEN,
    )
except Exception:
    restart_space()
try:
    print(EVAL_RESULTS_PATH)
    snapshot_download(
        repo_id=RESULTS_REPO,
        local_dir=EVAL_RESULTS_PATH,
        repo_type="dataset",
        tqdm_class=None,
        etag_timeout=30,
        token=TOKEN,
    )
except Exception:
    restart_space()

fonts_header = """
<link rel="preconnect" href="https://fonts.googleapis.com">
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
<link href="https://fonts.googleapis.com/css2?family=Shantell+Sans:ital,wght@0,300..800;1,300..800&display=swap" rel="stylesheet">
<link href="https://fonts.googleapis.com/css2?family=Space+Mono:ital,wght@0,400;0,700;1,400;1,700&display=swap" rel="stylesheet">
<link href="https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:ital,wght@0,100;0,200;0,300;0,400;0,500;0,600;0,700;1,100;1,200;1,300;1,400;1,500;1,600;1,700&display=swap" rel="stylesheet">
"""

js_head = """
<script>
    const gradioApp = document.getElementsByTagName('gradio-app')[0];
    console.log("Gradio app:", gradioApp);
    console.log(gradioApp.querySelectorAll('.token'));
    console.log(document.querySelectorAll('.token'));

    // Function to trigger Python callback
    const setHiddenIndex = (index) => {
        console.log("Setting hidden index to:", index);
        const hiddenIndex = gradioApp.querySelector("#hidden-index textarea");
        if (hiddenIndex) {
            hiddenIndex.value = index;
            let event = new Event("input", { bubbles: true});
            Object.defineProperty(event, "target", { value: hiddenIndex});
            hiddenIndex.dispatchEvent(event);
        }
    };

    // Add event listeners to all tokens
    function setupTokenListeners() {
        const tokens = gradioApp.querySelectorAll('.token');
        console.log("Tokens:", tokens);
        tokens.forEach(token => {
            token.addEventListener('mouseover', function() {
                const index = parseInt(this.getAttribute('data-index'));
                console.log("Mouseover token index:", index);

                // Reset all tokens
                gradioApp.querySelectorAll('.token').forEach(el => {
                    el.classList.remove('highlighted');
                });

                // Highlight this token
                this.classList.add('highlighted');

                // Update the hidden index to trigger the Python callback
                setHiddenIndex(index);
            });
        });
    }
    console.log("Preamble complete");

    document.addEventListener("DOMContentLoaded", function() {
        // Setup initial listeners
        console.log("DOM fully loaded and parsed");
        setupTokenListeners();

        // Setup a mutation observer to handle dynamically added tokens
        const observer = new MutationObserver(function(mutations) {
            mutations.forEach(function(mutation) {
                if (mutation.addedNodes.length) {
                    setupTokenListeners();
                }
            });
        });

        // Start observing the token container for changes
        const tokenContainer = gradioApp.querySelector('.token-container');
        console.log("Token container:", tokenContainer);
        if (tokenContainer) {
            observer.observe(tokenContainer.parentNode, { childList: true, subtree: true });
        }
        console.log("Listener setup complete");
    });
</script>
"""


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


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

    full_css = css_pipeline + css_tossup + css_bonus
    tossup_ds = load_dataset("tossup")
    bonus_ds = load_dataset("bonus")
    with gr.Blocks(
        css=full_css,
        head=fonts_header + js_head,
        theme=THEME,
        title="Quizbowl Bot",
    ) as demo:
        with gr.Sidebar(width=400):
            gr.Markdown(GUIDE_MARKDOWN)
        with gr.Row():
            gr.Markdown("## Welcome to Quizbowl Bot! This is a tool for creating and testing quizbowl agents.")
        with gr.Tabs():
            with gr.Tab("Tossup Agents"):
                defaults = DEFAULT_SELECTIONS["tossup"] | {
                    "init_workflow": factory.create_quizbowl_simple_workflow(),
                }
                tossup_interface = TossupInterface(demo, tossup_ds, AVAILABLE_MODELS, defaults)
            with gr.Tab("Bonus Round Agents"):
                defaults = DEFAULT_SELECTIONS["bonus"] | {
                    "init_workflow": factory.create_quizbowl_bonus_simple_workflow(),
                }
                bonus_interface = BonusInterface(demo, bonus_ds, AVAILABLE_MODELS, defaults)

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