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_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)


def get_backend_status_undefined_html() -> str:
    """
    返回未定义状态（首次检查前）的 HTML
    """
    return """
    <div id="backend-status-indicator">
        <span class="backend-status-light undefined"></span>
        <span>Backend Status: Checking...</span>
    </div>
    """


def check_backend_health() -> tuple[bool, str]:
    """
    查询后端健康状态
    返回: (is_healthy, status_html)
    """
    try:
        response = requests.get("http://localhost:8000/api/v1/health/", timeout=2)
        if response.status_code == 200:
            data = response.json()
            if data.get("code") == 0:
                return (
                    True,
                    """
                <div id="backend-status-indicator">
                    <span class="backend-status-light healthy"></span>
                    <span>Backend Status: Healthy</span>
                </div>
                """,
                )
        return (
            False,
            """
        <div id="backend-status-indicator">
            <span class="backend-status-light unhealthy"></span>
            <span>Backend Status: Unhealthy</span>
        </div>
        """,
        )
    except Exception:
        return (
            False,
            """
        <div id="backend-status-indicator">
            <span class="backend-status-light unhealthy"></span>
            <span>Backend Status: Unavailable</span>
        </div>
        """,
        )


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


demo = gr.Blocks(css_paths=[custom_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):
            # Backend status indicator - 初始状态为 undefined
            backend_status = gr.HTML(value=get_backend_status_undefined_html(), elem_id="backend-status-container")

            # 定时更新后端状态
            def update_backend_status():
                return check_backend_health()[1]

            # 创建一个隐藏的按钮用于定时触发更新
            status_trigger = gr.Button(visible=False, elem_id="backend-status-trigger-btn")

            # 绑定按钮点击事件来更新状态
            status_trigger.click(
                fn=update_backend_status,
                inputs=None,
                outputs=backend_status,
            )

            # 加载外部 JavaScript 文件
            js_content = backend_status_js.read_text(encoding="utf-8")
            status_trigger_js_html = f'<script>{js_content}</script>'
            gr.HTML(status_trigger_js_html)

            # 页面加载时立即检查一次
            demo.load(
                fn=update_backend_status,
                inputs=None,
                outputs=backend_status,
            )

            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,
            )


if __name__ == "__main__":
    # Backend server - 在单独的线程中运行
    app = create_app()

    def run_fastapi(host: str = "127.0.0.1", port: int = 8000):
        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()