import gradio as gr
import pandas as pd

LEADERBOARD_FILE = "results.csv"

def load_leaderboard():
    return pd.read_csv(LEADERBOARD_FILE)

def display_data(data_type="Overall"):
    df = load_leaderboard()
    for col in ["Safe Completion Rate", "Harmful Completion Rate", "Refusal Rate", "Normalized Safety Score", "Bias Completion Rate", "Cybercrime Completion Rate", "Harassment Completion Rate","Misinformation Completion Rate", "Illegal Activity Completion Rate"]:
        df[col] = df[col].apply(lambda x: f"{x:.1f}")
    if data_type == "Overall":
        new_df = df[["Model", "Safe Completion Rate", "Harmful Completion Rate", "Refusal Rate", "Normalized Safety Score", "License"]].copy()
    else:
        category_col = f"{data_type} Completion Rate"
        new_df = df[["Model", category_col, "Safe Completion Rate", "License"]].copy()
    sort_col = "Model"
    return new_df.sort_values(by=sort_col)

with gr.Blocks() as demo:
    gr.Markdown("# SafeArena Leaderboard")
    
    gr.Markdown("""
    # SafeArena Leaderboard
    
    | [**🤗Dataset**](https://huggingface.co/datasets/McGill-NLP/safearena) | [**📄Paper**](https://arxiv.org/abs/2503.04957) | [**🌐Website**](https://safearena.github.io) | [**💾Code**](https://github.com/McGill-NLP/safearena) |
    | :--: | :--: | :--: | :--: |
    """)
    
    data_type_dropdown = gr.Dropdown(
        label="Data Type",
        choices=["Overall", "Bias", "Cybercrime", "Harassment", "Misinformation", "Illegal Activity"],
        value="Overall"
    )

    table = gr.Dataframe(value=display_data("Overall"), interactive=False)
    
    data_type_dropdown.change(
        fn=display_data,
        inputs=data_type_dropdown,
        outputs=table
    )
    
    demo.launch(share=True)