import gradio as gr
import pandas as pd
import plotly.graph_objects as go

df = pd.read_csv("code_eval_board.csv")
df = df.sort_values(by=["Average score"], ascending=False)
headers = df.columns.to_list()

def plot_throughput(bs=1):
    throughput_column = 'Throughput (tokens/s)' if bs==1 else 'Throughput (tokens/s) bs=50'

    df['symbol'] = 2  # Triangle
    df['color'] = ''
    df.loc[df['Models'].str.contains('StarCoder|SantaCoder'), 'color'] = 'orange'
    df.loc[df['Models'].str.contains('CodeGen'), 'color'] = 'pink'
    df.loc[df['Models'].str.contains('Replit'), 'color'] = 'purple'
    df.loc[df['Models'].str.contains('Wizard'), 'color']  = '#00b3b3'
    df.loc[df['Models'].str.contains('CodeGeeX'), 'color'] = '#00cc00'

    fig = go.Figure()

    for i in df.index:
        fig.add_trace(go.Scatter(
            x=[df.loc[i, throughput_column]], 
            y=[df.loc[i, 'Average score']],
            mode='markers',
            marker=dict(
                size=[df.loc[i, 'Size (B)'] + 10],
                color=df.loc[i, 'color'],
                symbol=df.loc[i, 'symbol']
            ),
            name=df.loc[i, 'Models'],
            hovertemplate =
            '<b>%{text}</b><br><br>' +
            f'{throughput_column}: %{{x}}<br>'+
            'Average Score: %{y}<br>' +
            'Peak Memory (MB): ' + str(df.loc[i, 'Peak Memory (MB)']) + '<br>' +
            'Human Eval (Python): ' + str(df.loc[i, 'humaneval-python']),
            text=[df.loc[i, 'Models']],
            showlegend=True
        ))

    fig.update_layout(
        autosize=False,
        width=700,
        height=600,
        title=f'Average Score Vs Throughput (A100-80GB, Float16, Batch Size <b>{bs}</b>)',
        xaxis_title=f'{throughput_column}',
        yaxis_title='Average Code Score',
    )
    return fig


demo = gr.Blocks()
with demo:
    with gr.Row():
        gr.Markdown(
            """<div style="text-align: center;"><h1> ⭐ Multilingual <span style='color: #e6b800;'>Code</span> Models <span style='color: #e6b800;'>Evaluation</span></h1></div>\
            <br>\
            <p>We compare performance of base multilingual code generation models on <a href="https://huggingface.co/datasets/openai_humaneval">HumanEval</a> benchmark and <a href="https://huggingface.co/datasets/nuprl/MultiPL-E">MultiPL-E</a>. Following the <a href="https://huggingface.co/spaces/optimum/llm-perf-leaderboard">🤗 Open LLM-Perf Leaderboard 🏋️</a>, we also measure throughput and provide\
            information about the models. We only compare pre-trained multilingual code models, that people can start from as base models for their trainings.</p>"""
        )

    with gr.Column():
        with gr.Tabs(elem_classes="A100-tabs") as A100_tabs:
            with gr.TabItem("🔍 Evaluation table", id=0):
                leaderboard_df = gr.components.Dataframe(
                    value=df, headers=headers, datatype=["str" for _ in range(len(headers))]
                )
            """
            with gr.TabItem("📊 Performance Plot", id=1):
                with gr.Row():
                    bs_1_plot = gr.components.Plot(
                        value=plot_throughput(bs=1),
                        elem_id="bs1-plot",
                        show_label=False,
                    )
                    bs_50_plt = gr.components.Plot(
                        value=plot_throughput(bs=50),
                        elem_id="bs50-plot",
                        show_label=False,
                    )"""
    with gr.Row():
        gr.Markdown(
            """Notes:
            <ul>
            <li> Throughputs and peak memory usage are measured using <a href="https://github.com/huggingface/optimum-benchmark/tree/main">Optimum-Benchmark</a> which powers <a href="https://huggingface.co/spaces/optimum/llm-perf-leaderboard">Open LLM-Perf Leaderboard</a>. (0 throughput corresponds to OOM).</li>
            <li> All models were evaluated with the <a href="https://github.com/bigcode-project/bigcode-evaluation-harness/tree/main">bigcode-evaluation-harness</a> with top-p=0.95, temperature=0.2 and n_samples=50.</li>
            <li> HumanEval-Python, reports the pass@1 on HumanEval, the rest is from MultiPL-E benchmark.</li>
            <li> Average score is the average pass@1 over all languages. For Win Rate, we compute model rank for each language as <pre><code>num_models - (rank -1)</code></pre> and average their rankings.</li>
            <li> #Languages column represents the number of programming languages included during the pretraining.
            </ul>"""
        )
demo.launch()