app.py

import gradio as gr
import asyncio
import time
import threading
from src.retrieve_data import populate_lb_data

leaderboard_data = {}


async def fetch_data():
    """Fetch the leaderboard data asynchronously"""
    global leaderboard_data
    try:
        data = await populate_lb_data()
        leaderboard_data = data
        return True
    except Exception as e:
        print(f"Error fetching data: {e}")
        return False


def background_update():
    """Background thread function to update data every 5 minutes"""
    while True:
        print("Updating leaderboard data...")
        asyncio.run(fetch_data())
        time.sleep(300)  # 5 minutes


def create_table_for_lb(lb_name, gpu_name):
    """Create a formatted table for a specific leaderboard and GPU"""
    if (
        not lb_name
        or not gpu_name
        or lb_name not in leaderboard_data
        or gpu_name not in leaderboard_data[lb_name]
    ):
        return gr.Dataframe(
            headers=["Rank", "Submission Name", "User ID", "Score", "Date"],
            datatype=["number", "str", "str", "str", "str"],
            value=[],
            interactive=False,
        )

    lb_data = leaderboard_data[lb_name][gpu_name]

    headers = ["Rank", "Submission Name", "User ID", "Score", "Date"]

    rows = []
    for i, result in enumerate(lb_data.results, 1):
        rows.append(
            [
                i,
                result.submission_name,
                result.user_id,
                f"{float(result.submission_score):.4f}",
                result.submission_time,
            ]
        )

    return gr.Dataframe(
        headers=headers,
        datatype=["number", "str", "str", "str", "str"],
        value=rows,
        interactive=False,
    )


def refresh_ui():
    """Force refresh the UI with latest data"""
    asyncio.run(fetch_data())
    return "Data refreshed!", get_lb_names(), [], None


def get_lb_names():
    """Get list of available leaderboard names"""
    return list(leaderboard_data.keys())


def get_gpu_names(lb_name):
    """Get list of available GPUs for a specific leaderboard"""
    if not lb_name or lb_name not in leaderboard_data:
        return []
    return list(leaderboard_data[lb_name].keys())


def on_lb_change(lb_name):
    gpu_choices = get_gpu_names(lb_name)

    return (
        gr.update(choices=gpu_choices, value=gpu_choices[0] if gpu_choices else None),
        update_table(lb_name, gpu_choices[0]),
    )


def update_table(lb_name, gpu_name):
    """Update the table based on selected leaderboard and GPU"""
    if not lb_name or not gpu_name:
        return None
    return create_table_for_lb(lb_name, gpu_name)


def build_ui():
    with gr.Blocks(title="ML Leaderboards", theme=gr.themes.Soft()) as app:
        gr.Markdown("# 🍿 KernelBot Leaderboard 🍿")

        asyncio.run(fetch_data())

        with gr.Row():
            with gr.Column(scale=1):
                lb_dropdown = gr.Dropdown(
                    choices=get_lb_names(),
                    label="Select Leaderboard",
                    interactive=True,
                )
                gpu_dropdown = gr.Dropdown(
                    choices=get_gpu_names(lb_dropdown.value),
                    label="Select GPU",
                    interactive=True,
                )

        with gr.Row():
            results_table = gr.Dataframe(
                headers=["Rank", "Submission Name", "User ID", "Score", "Date"],
                datatype=["number", "str", "str", "str", "str"],
                interactive=False,
                label="Results",
            )

        lb_dropdown.change(
            fn=on_lb_change,
            inputs=[lb_dropdown],
            outputs=[gpu_dropdown, results_table],
        )
        
        gpu_dropdown.change(
            fn=update_table, inputs=[lb_dropdown, gpu_dropdown], outputs=results_table
        )

    return app


if __name__ == "__main__":
    update_thread = threading.Thread(target=background_update, daemon=True)
    update_thread.start()
    app = build_ui()
    app.launch()