app.py

import threading

import gradio as gr
import gradio.components as grc
import pandas as pd
import requests
import uvicorn
from apscheduler.schedulers.background import BackgroundScheduler
from huggingface_hub import snapshot_download
from rich import print

from src.about import (
    BENCHMARKS,
    CITATION_BUTTON_LABEL,
    CITATION_BUTTON_TEXT,
    EVALUATION_QUEUE_TEXT,
    INTRODUCTION_TEXT,
    LLM_BENCHMARKS_TEXT,
    TITLE,
)
from src.backend.app import create_app
from src.display.css_html_js import (
    backend_status_indicator_css,
    backend_status_indicator_html,
    backend_status_js,
    custom_css,
)
from src.display.utils import (
    BASE_COLS,
    BENCHMARK_COLS,
    COLS,
    EVAL_COLS,
    EVAL_TYPES,
    AutoEvalColumn,
    ModelType,
    Precision,
    WeightType,
)
from src.envs import API, settings
from src.populate import get_evaluation_queue_df, get_leaderboard_df
from src.submission.submit import add_new_eval


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


print("///// --- Settings --- /////", settings.model_dump())

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


LEADERBOARD_DF = get_leaderboard_df(
    settings.EVAL_RESULTS_PATH,
    settings.EVAL_REQUESTS_PATH,
    COLS,
    BENCHMARK_COLS,
)

(
    finished_eval_queue_df,
    running_eval_queue_df,
    pending_eval_queue_df,
) = get_evaluation_queue_df(settings.EVAL_REQUESTS_PATH, EVAL_COLS)


def filter_dataframe_by_columns(selected_cols: list[str], original_df: pd.DataFrame) -> pd.DataFrame:
    """
    根据选择的列过滤 DataFrame
    """
    # 始终包含基础列 'T' 和 'Model'
    base_cols = ['T', 'Model']
    all_selected_cols = [col for col in base_cols if col in original_df.columns]

    # 添加用户选择的列（排除已存在的基础列）
    for col in selected_cols:
        if col in original_df.columns and col not in all_selected_cols:
            all_selected_cols.append(col)

    # 确保列的顺序：基础列在前，然后是按原始顺序的选中列
    ordered_cols = []
    for col in original_df.columns:
        if col in all_selected_cols:
            ordered_cols.append(col)

    # 确保总是返回 DataFrame，即使是单列也使用 [[]] 来保持 DataFrame 类型
    if ordered_cols:
        filtered_df = original_df.loc[:, ordered_cols]
    else:
        filtered_df = original_df
    return filtered_df


def filter_dataframe_by_precision(selected_precisions: list[str], df: pd.DataFrame) -> pd.DataFrame:
    """
    根据选择的 precision 筛选 DataFrame
    如果没有选择 precision，返回空的 DataFrame
    """
    if not selected_precisions:
        return df.iloc[0:0].copy()  # 返回相同结构但为空的 DataFrame

    precision_col = AutoEvalColumn.precision.name
    if precision_col not in df.columns:
        return df

    # 筛选包含任一选定 precision 的行
    mask = df[precision_col].isin(selected_precisions)
    filtered_df = df.loc[mask, :]
    return filtered_df


def search_models_in_dataframe(search_text: str, df: pd.DataFrame) -> pd.DataFrame:
    """
    在 DataFrame 中搜索包含关键词的 Model 名称
    支持逗号分隔的多个关键词，匹配包含任一关键词的行
    """
    if not search_text or not search_text.strip():
        return df

    # 分割逗号，去除空白并转换为小写用于匹配
    import re

    keywords = [keyword.strip().lower() for keyword in search_text.split(',') if keyword.strip()]
    if not keywords:
        return df

    if 'Model' not in df.columns:
        return df

    # 匹配函数：从 HTML 中提取纯文本并检查是否包含关键词
    def matches_search(model_cell):
        if pd.isna(model_cell):
            return False

        # 从 HTML 链接中提取纯文本（model_name）
        # 格式: <a ...>model_name</a> 或直接是文本
        text = str(model_cell)

        # 提取 HTML 标签内的文本
        # 匹配 <a>...</a> 标签内的内容，或直接使用文本
        match = re.search(r'<a[^>]*>([^<]+)</a>', text, re.IGNORECASE)
        if match:
            model_name = match.group(1).lower()
        else:
            model_name = text.lower()

        # 检查是否包含任一关键词
        return any(keyword in model_name for keyword in keywords)

    # 应用搜索过滤
    mask = df['Model'].apply(matches_search)
    filtered_df = df.loc[mask, :]
    return filtered_df


def init_leaderboard_tabs(dataframe: pd.DataFrame, cols: list[str]):
    # 存储原始 DataFrame 以便后续过滤使用（使用闭包保存）
    original_df = dataframe.copy()

    available_precisions = sorted(original_df["Precision"].dropna().unique().tolist())
    default_precision = (
        ['bfloat16']
        if 'bfloat16' in available_precisions
        else (available_precisions[:1] if available_precisions else [])
    )

    # 初始化显示的列（包含基础列和默认选中的列）
    default_selected = [col for col in dataframe.columns if col in cols] + ['Average ⬆️']

    # 先按 precision 筛选 original_df
    precision_filtered_df = filter_dataframe_by_precision(default_precision, original_df)
    # 根据默认选择再筛选一次 DataFrame
    initial_filtered_df = filter_dataframe_by_columns(default_selected, precision_filtered_df)

    with gr.Row():
        with gr.Column(scale=1):
            search = gr.Textbox(label="Search", placeholder="Separate multiple queries with commas")
            show_columns = gr.CheckboxGroup(
                choices=[col for col in dataframe.columns if col not in ['T', 'Model']],
                label="Select Columns to Display",
                value=default_selected,
                interactive=True,
            )
        with gr.Column(scale=1):
            _model_type = gr.CheckboxGroup(
                [],
                label="Model Type",
                value=[],
            )
            precision = gr.CheckboxGroup(
                choices=available_precisions,
                label="Precision",
                value=default_precision,
                interactive=True,
            )
            _hide_models = gr.CheckboxGroup(
                ['Deleted/incomplete'],
                label="Hide Models",
                value=['Deleted/incomplete'],
                interactive=True,
            )

    with gr.Row():
        with gr.Column(scale=3):
            leaderboard = gr.Dataframe(
                value=initial_filtered_df,  # 使用初始筛选后的 DataFrame
                interactive=False,
                wrap=False,
                datatype='markdown',
                elem_id="auto-width-dataframe",
            )

    # 统一的更新函数：同时处理 precision、列筛选和搜索
    def update_dataframe(search_text: str, selected_cols: list[str], selected_precisions: list[str]):
        # 先按 precision 筛选 original_df
        precision_filtered_df = filter_dataframe_by_precision(selected_precisions, original_df)
        # 再按列筛选
        column_filtered_df = filter_dataframe_by_columns(selected_cols, precision_filtered_df)
        # 最后按搜索关键词筛选
        final_df = search_models_in_dataframe(search_text, column_filtered_df)
        return final_df

    # 绑定搜索、列选择和 precision 的变化事件，动态更新 DataFrame
    search.change(
        fn=update_dataframe,
        inputs=[search, show_columns, precision],
        outputs=leaderboard,
    )

    show_columns.change(
        fn=update_dataframe,
        inputs=[search, show_columns, precision],
        outputs=leaderboard,
    )

    precision.change(
        fn=update_dataframe,
        inputs=[search, show_columns, precision],
        outputs=leaderboard,
    )

    return leaderboard


def main():
    demo = gr.Blocks(css_paths=[custom_css, backend_status_indicator_css])
    with demo:
        gr.HTML(TITLE)
        gr.Markdown(INTRODUCTION_TEXT, elem_classes="markdown-text")

        with gr.Tabs(elem_classes="tab-buttons") as _tabs:
            for i, benchmark in enumerate[str](sorted(BENCHMARKS)):
                with gr.TabItem(f"🏅 {benchmark}", elem_id="llm-benchmark-tab-table", id=i):
                    benchmark_cols = [
                        BENCHMARK_COL for BENCHMARK_COL in BENCHMARK_COLS if BENCHMARK_COL.startswith(benchmark)
                    ]
                    cols = BASE_COLS + benchmark_cols
                    BENCHMARK_DF = get_leaderboard_df(
                        settings.EVAL_RESULTS_PATH,
                        settings.EVAL_REQUESTS_PATH,
                        cols,
                        benchmark_cols,
                    )
                    _leaderboard = init_leaderboard_tabs(BENCHMARK_DF, benchmark_cols)

            with gr.TabItem("📝 About", elem_id="about-tab", id=len(BENCHMARKS)):
                gr.Markdown(LLM_BENCHMARKS_TEXT, elem_classes="markdown-text")

            with gr.TabItem("🚀 Submit here! ", elem_id="submit-tab", id=len(BENCHMARKS) + 1):
                with gr.Column():
                    with gr.Row():
                        gr.Markdown(EVALUATION_QUEUE_TEXT, elem_classes="markdown-text")

                    with gr.Column():
                        with gr.Accordion(
                            f"✅ Finished Evaluations ({len(finished_eval_queue_df)})",
                            open=False,
                        ):
                            with gr.Row():
                                _finished_eval_table = grc.Dataframe(
                                    value=finished_eval_queue_df,
                                    headers=EVAL_COLS,
                                    datatype=EVAL_TYPES,
                                    row_count=5,
                                )
                        with gr.Accordion(
                            f"🔄 Running Evaluation Queue ({len(running_eval_queue_df)})",
                            open=False,
                        ):
                            with gr.Row():
                                _running_eval_table = grc.Dataframe(
                                    value=running_eval_queue_df,
                                    headers=EVAL_COLS,
                                    datatype=EVAL_TYPES,
                                    row_count=5,
                                )

                        with gr.Accordion(
                            f"⏳ Pending Evaluation Queue ({len(pending_eval_queue_df)})",
                            open=False,
                        ):
                            with gr.Row():
                                _pending_eval_table = grc.Dataframe(
                                    value=pending_eval_queue_df,
                                    headers=EVAL_COLS,
                                    datatype=EVAL_TYPES,
                                    row_count=5,
                                )
                with gr.Row():
                    gr.Markdown("# ✉️✨ Submit your model here!", elem_classes="markdown-text")

                with gr.Row():
                    search_name = gr.Textbox(label="search model name", placeholder="user/model_name")

                with gr.Row():
                    table = gr.Dataframe(
                        headers=["Model Name", "Pipeline", "Downloads", "Likes"],
                        datatype=["str", "str", "number", "number"],
                        interactive=False,
                        wrap=True,
                        label="click model name to select",
                    )

                with gr.Row():
                    with gr.Column():
                        model_name_textbox = gr.Textbox(label="Model name", placeholder="user/model_name")
                        revision_name_textbox = gr.Textbox(label="Revision commit", placeholder="main")
                        model_type = gr.Dropdown(
                            choices=[t.to_str(" : ") for t in ModelType if t != ModelType.Unknown],
                            label="Model type",
                            multiselect=False,
                            value=None,
                            interactive=True,
                        )

                        def search_models(query):
                            if not query.strip():
                                return []
                            models = API.list_models(search=query, limit=10)
                            results = []
                            for m in models:
                                results.append([m.id, m.pipeline_tag or "N/A", m.downloads or 0, m.likes or 0])
                            return results

                        def on_select(evt: gr.SelectData, data):
                            row_idx = evt.index[0]  # 获取点击行号
                            if row_idx < len(data):
                                return data.iloc[row_idx, 0]  # 返回模型名
                            return ""

                        search_name.change(fn=search_models, inputs=search_name, outputs=table)
                        table.select(fn=on_select, inputs=table, outputs=model_name_textbox)

                    with gr.Column():
                        precision = gr.Dropdown(
                            choices=[i.value.name for i in Precision if i != Precision.Unknown],
                            label="Precision",
                            multiselect=False,
                            value="float16",
                            interactive=True,
                        )
                        weight_type = gr.Dropdown(
                            choices=[i.value.name for i in WeightType],
                            label="Weights type",
                            multiselect=False,
                            value="Original",
                            interactive=True,
                        )
                        base_model_name_textbox = gr.Textbox(label="Base model (for delta or adapter weights)")

                submit_button = gr.Button("Submit Eval")
                submission_result = gr.Markdown()
                submit_button.click(
                    add_new_eval,
                    [
                        model_name_textbox,
                        base_model_name_textbox,
                        revision_name_textbox,
                        precision,
                        weight_type,
                        model_type,
                    ],
                    submission_result,
                )

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

        # Backend status indicator
        backend_status = gr.HTML(
            value=get_backend_status_undefined_html(),
            elem_id="backend-status-container",
        )
        # trigger button to bind the click event
        status_trigger = gr.Button(elem_id="backend-status-trigger-btn", visible=False)
        status_trigger.click(
            fn=lambda: check_backend_health()[1],
            inputs=None,
            outputs=backend_status,
        )
        # load external JavaScript file
        js_content = backend_status_js()
        status_trigger_js_html = f'<script>{js_content}</script>'
        gr.HTML(status_trigger_js_html, visible=False)
        demo.load(
            fn=lambda: check_backend_health()[1],
            inputs=None,
            outputs=backend_status,
        )
    return demo


def get_backend_status_undefined_html() -> str:
    """
    返回未定义状态（首次检查前）的 HTML
    """
    return backend_status_indicator_html("undefined")


def check_backend_health() -> tuple[bool, str]:
    """
    查询后端健康状态
    返回: (is_healthy, status_html)
    """
    try:
        response = requests.get(f"http://localhost:{settings.BACKEND_PORT}/api/v1/health/", timeout=2)
        if response.status_code == 200:
            data = response.json()
            if data.get("code") == 0:
                return (
                    True,
                    backend_status_indicator_html("healthy"),
                )
        return (
            False,
            backend_status_indicator_html("unhealthy"),
        )
    except Exception:
        return (
            False,
            backend_status_indicator_html("unhealthy"),
        )


if __name__ == "__main__":
    demo = main()

    # Backend server - 在单独的线程中运行
    app = create_app()

    def run_fastapi():
        host = settings.BACKEND_HOST
        port = settings.BACKEND_PORT
        print(f"Starting FastAPI server on http://{host}:{port}")
        uvicorn.run(
            app,
            host=host,
            port=port,
            log_level="debug",
            access_log=True,
        )

    fastapi_thread = threading.Thread(target=run_fastapi, daemon=True)
    fastapi_thread.start()

    # Gradio server - 在主线程中运行（阻塞）
    scheduler = BackgroundScheduler()
    scheduler.add_job(restart_space, "interval", seconds=1800)
    scheduler.start()
    demo.queue(default_concurrency_limit=40).launch()