app.py

import gradio as gr
import pandas as pd
from pathlib import Path

from src.json_leaderboard import create_leaderboard_df
from src.about import (
    CITATION_BUTTON_TEXT,
    INTRODUCTION_TEXT,
    LINKS_AND_INFO,
    TITLE,
)
from src.display.css_html_js import custom_css

# 固定列，永远在前面
FIXED_COLUMNS = ["Model Name (clickable)", "Release Date", "HF Model", "Open Source"]  

def get_json_df():
    """Load the leaderboard DataFrame"""
    json_path = Path(__file__).parent / "leaderboard_data.json"
    df = create_leaderboard_df(str(json_path))
    return df

# 提取大类及其子类
def extract_categories_and_subs(df):
    """
    返回 {大类: {"overall": 大类列, "subs": [子类列]}}
    大类列以 '-Overall' 结尾，紧跟其后的列为子类
    """
    category_dict = {}
    all_cols = list(df.columns)
    skip_cols = set(FIXED_COLUMNS + ["Overall"])

    i = 0
    while i < len(all_cols):
        col = all_cols[i]
        if col.endswith("-Overall") and col not in skip_cols:
            cat_name = col.replace("-Overall", "")
            subs = []
            j = i + 1
            while j < len(all_cols):
                next_col = all_cols[j]
                if next_col.endswith("-Overall") or next_col in skip_cols:
                    break
                subs.append(next_col)
                j += 1
            category_dict[cat_name] = {"overall": col, "subs": subs}
        i += 1
    return category_dict

def build_overall_tab_cols(df, categories):
    # Overall 页面希望把 Real/Game 的 '-Overall' 放在子项之后
    real = categories.get("Real-World", {"overall": None, "subs": []})
    game = categories.get("Game-World", {"overall": None, "subs": []})

    base = FIXED_COLUMNS.copy()
    rest = [c for c in df.columns if c not in base]

    ordered = base + ["Overall"]
    # real: 子项 -> overall
    ordered += [c for c in real.get("subs", []) if c in df.columns]
    if real.get("overall") in df.columns:
        ordered += [real["overall"]]
    # game: 子项 -> overall
    ordered += [c for c in game.get("subs", []) if c in df.columns]
    if game.get("overall") in df.columns:
        ordered += [game["overall"]]
    # 其它剩余列
    ordered += [c for c in rest if c not in ordered]
    return ordered

# 初始化
df = get_json_df()
ALL_COLUMNS_ORDERED = list(df.columns)
categories = extract_categories_and_subs(df)

# 可选列 = 全部列 - 固定列
optional_columns = [col for col in df.columns if col not in FIXED_COLUMNS]

# Gradio interface
demo = gr.Blocks(css=custom_css, title="UniREditBench Leaderboard")

with demo:
    gr.HTML(TITLE)
    gr.HTML(LINKS_AND_INFO)
    gr.Markdown(INTRODUCTION_TEXT, elem_classes="markdown-text")

    with gr.Tabs(elem_classes="tab-buttons") as tabs:
        # Overall leaderboard
        with gr.TabItem("🏅 Overall Leaderboard", elem_id="tab-overall"):
            selected_columns_overall = gr.CheckboxGroup(
                choices=optional_columns,
                label="Select additional columns to display",
                value=optional_columns
            )
            overall_cols = build_overall_tab_cols(df, categories)
            leaderboard_table = gr.Dataframe(
                value=df[overall_cols],
                headers=list(df[overall_cols].columns),
                datatype=["html" if col in ["Model Name (clickable)","HF Model"] else "str" for col in df[overall_cols].columns],
                interactive=False,
                wrap=False
            )
            selected_columns_overall.change(
                fn=lambda selected, d: d[[*FIXED_COLUMNS, *[c for c in d.columns if c in selected and c not in FIXED_COLUMNS]]],
                inputs=[selected_columns_overall, gr.State(value=df[overall_cols])],
                outputs=leaderboard_table
            )

        # 每个大类 leaderboard（保持 overall 在前，便于分组）
        for cat_name, info in categories.items():
            with gr.TabItem(f"🏆 {cat_name}", elem_id=f"tab-{cat_name}"):
                cat_cols = [info["overall"]] + info["subs"]
                cat_df = df[FIXED_COLUMNS + cat_cols]

                optional_columns_cat = [col for col in cat_cols if col not in FIXED_COLUMNS]
                selected_columns_cat = gr.CheckboxGroup(
                    choices=optional_columns_cat,
                    label=f"Select additional columns for {cat_name}",
                    value=optional_columns_cat
                )
                leaderboard_table_cat = gr.Dataframe(
                    value=cat_df,
                    headers=list(cat_df.columns),
                    datatype=["html" if col in ["Model Name (clickable)","HF Model"] else "str" for col in cat_df.columns],
                    interactive=False,
                    wrap=False
                )
                selected_columns_cat.change(
                    fn=lambda selected, d: d[[*FIXED_COLUMNS, *[c for c in d.columns if c in selected and c not in FIXED_COLUMNS]]],
                    inputs=[selected_columns_cat, gr.State(value=cat_df)],
                    outputs=leaderboard_table_cat
                )

    # Citation
    with gr.Row():
        with gr.Column():
            gr.Markdown("## 📙 Citation")
            gr.Markdown("If you use UniREditBench in your research, please cite our work:")
            citation_textbox = gr.Textbox(
                value=CITATION_BUTTON_TEXT,
                elem_id="citation-textbox",
                show_label=False,
                interactive=False,
                lines=8,
                show_copy_button=True
            )

if __name__ == "__main__":
    demo.launch()