app.py

from __future__ import annotations

import html
import logging
import re
from datetime import datetime, timezone
from pathlib import Path
from threading import RLock
from typing import Any

import gradio as gr
import pandas as pd
from apscheduler.schedulers.background import BackgroundScheduler
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from huggingface_hub import snapshot_download

from src.about import (
    CITATION_BUTTON_LABEL,
    CITATION_BUTTON_TEXT,
    EVALUATION_QUEUE_TEXT,
    LLM_BENCHMARKS_TEXT,
    TITLE,
    Tasks,
)
from src.display.utils import (
    BENCHMARK_COLS,
    COLS,
    EVAL_COLS,
    ModelType,
    Precision,
    WeightType,
)
from src.envs import EVAL_REQUESTS_PATH, EVAL_RESULTS_PATH, QUEUE_REPO, RESULTS_REPO, TOKEN
from src.populate import get_evaluation_queue_df, get_leaderboard_df
from src.submission.submit import add_new_eval

logging.basicConfig(level=logging.INFO)
LOGGER = logging.getLogger("leaderboard")

ROOT = Path(__file__).parent
STATIC_DIR = ROOT / "static"
MODEL_LINK_RE = re.compile(r'href=["\']https://huggingface\.co/([^"\']+)["\']')
DATA_LOCK = RLock()
DATA: dict[str, Any] = {}


def _download_dataset(repo_id: str, local_dir: str) -> None:
    Path(local_dir).mkdir(parents=True, exist_ok=True)
    snapshot_download(
        repo_id=repo_id,
        local_dir=local_dir,
        repo_type="dataset",
        tqdm_class=None,
        etag_timeout=30,
        token=TOKEN,
    )


def _json_value(value: Any) -> Any:
    if pd.isna(value):
        return None
    if hasattr(value, "item"):
        value = value.item()
    return value


def _model_id(value: Any) -> str:
    text = str(value or "")
    if match := MODEL_LINK_RE.search(text):
        return html.unescape(match.group(1))
    return html.unescape(re.sub("<[^>]+>", "", text))


def _records(frame: pd.DataFrame) -> list[dict[str, Any]]:
    records = []
    for row in frame.to_dict(orient="records"):
        clean_row = {str(key): _json_value(value) for key, value in row.items()}
        if "Model" in clean_row:
            clean_row["Model"] = _model_id(clean_row["Model"])
        if "model" in clean_row:
            clean_row["model"] = _model_id(clean_row["model"])
        records.append(clean_row)
    return records


def _empty_payload(message: str | None = None) -> dict[str, Any]:
    return {
        "rows": [],
        "queues": {"finished": [], "running": [], "pending": []},
        "benchmarks": [task.value.col_name for task in Tasks],
        "model_types": [item.to_str(" · ") for item in ModelType if item is not ModelType.Unknown],
        "precisions": [item.value.name for item in Precision if item is not Precision.Unknown],
        "weight_types": [item.value.name for item in WeightType],
        "copy": {
            "title": TITLE,
            "about": LLM_BENCHMARKS_TEXT.strip(),
            "submission": EVALUATION_QUEUE_TEXT.strip(),
            "citation_label": CITATION_BUTTON_LABEL,
            "citation": CITATION_BUTTON_TEXT.strip(),
        },
        "updated_at": datetime.now(timezone.utc).isoformat(),
        "error": message,
    }


def refresh_data(*, download: bool = True) -> dict[str, Any]:
    """Refresh leaderboard and queue data without taking the Space offline."""
    errors = []
    if download:
        for repo_id, local_dir in (
            (QUEUE_REPO, EVAL_REQUESTS_PATH),
            (RESULTS_REPO, EVAL_RESULTS_PATH),
        ):
            try:
                _download_dataset(repo_id, local_dir)
            except Exception as exc:  # Keep serving the last successful snapshot.
                LOGGER.exception("Could not refresh %s", repo_id)
                errors.append(f"Could not refresh {repo_id}: {exc}")

    payload = _empty_payload("; ".join(errors) or None)
    try:
        leaderboard = get_leaderboard_df(EVAL_RESULTS_PATH, EVAL_REQUESTS_PATH, COLS, BENCHMARK_COLS)
        payload["rows"] = _records(leaderboard)
    except Exception as exc:
        LOGGER.exception("Could not build leaderboard data")
        payload["error"] = f"Could not build leaderboard data: {exc}"

    try:
        finished, running, pending = get_evaluation_queue_df(EVAL_REQUESTS_PATH, EVAL_COLS)
        payload["queues"] = {
            "finished": _records(finished),
            "running": _records(running),
            "pending": _records(pending),
        }
    except Exception as exc:
        LOGGER.exception("Could not build evaluation queues")
        payload["error"] = f"Could not build evaluation queues: {exc}"

    with DATA_LOCK:
        DATA.clear()
        DATA.update(payload)
        return dict(DATA)


server = gr.Server(
    title="Open LLM Benchmark",
    description="A clean, customizable Gradio 6 leaderboard template.",
    version="2.0.0",
)
server.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


@server.get("/", response_class=HTMLResponse, include_in_schema=False)
async def homepage() -> HTMLResponse:
    return HTMLResponse((STATIC_DIR / "index.html").read_text())


@server.api(name="leaderboard_data", queue=False, concurrency_limit=None)
def leaderboard_data() -> dict[str, Any]:
    """Return the current leaderboard, evaluation queues, filters, and template copy."""
    with DATA_LOCK:
        return dict(DATA) if DATA else _empty_payload("Data is still loading.")


@server.api(name="submit_model", concurrency_limit=4)
def submit_model(
    model: str,
    base_model: str,
    revision: str,
    precision: str,
    weight_type: str,
    model_type: str,
) -> dict[str, str]:
    """Validate a public Hugging Face model and add it to the evaluation queue."""
    result = add_new_eval(model, base_model, revision, precision, weight_type, model_type)
    if result["status"] == "success":
        refresh_data(download=False)
    return result


if __name__ == "__main__":
    refresh_data()
    scheduler = BackgroundScheduler(daemon=True)
    scheduler.add_job(refresh_data, "interval", minutes=10, max_instances=1, coalesce=True)
    scheduler.start()
    server.launch()