Update app.py

app.py

@@ -1,120 +1,177 @@
 import gradio as gr
 import pandas as pd
-from huggingface_hub import list_models
 import plotly.express as px
 
 CITATION_BUTTON_LABEL = "Copy the following snippet to cite these results"
 CITATION_BUTTON_TEXT = r"""@misc{energystarai-leaderboard,
-                  author = {Sasha Luccioni and Boris Gamazaychikov and Emma Strubell and Sara Hooker and Yacine Jernite and Carole-Jean Wu and Margaret Mitchell},
-                  title = {AI Energy Score Leaderboard v.0},
-                  year = {2024},
-                  publisher = {Hugging Face},
-                  howpublished = "\url{https://huggingface.co/spaces/EnergyStarAI/2024_Leaderboard}",
-                    }
-                 """
-tasks = ['asr.csv', 'object_detection.csv', 'text_classification.csv', 'image_captioning.csv',
-        'question_answering.csv', 'text_generation.csv', 'image_classification.csv',
-        'sentence_similarity.csv', 'image_generation.csv', 'summarization.csv']
+    author = {Sasha Luccioni and Boris Gamazaychikov and Emma Strubell and Sara Hooker and Yacine Jernite and Carole-Jean Wu and Margaret Mitchell},
+    title = {AI Energy Score Leaderboard v.0},
+    year = {2024},
+    publisher = {Hugging Face},
+    howpublished = "\url{https://huggingface.co/spaces/EnergyStarAI/2024_Leaderboard}",
+}"""
+
+# List of tasks (CSV filenames)
+tasks = [
+    'asr.csv',
+    'object_detection.csv',
+    'text_classification.csv',
+    'image_captioning.csv',
+    'question_answering.csv',
+    'text_generation.csv',
+    'image_classification.csv',
+    'sentence_similarity.csv',
+    'image_generation.csv',
+    'summarization.csv'
+]
+
+def format_stars(score):
+    """
+    Convert the energy_score (assumed to be an integer from 1 to 5)
+    into that many star characters wrapped in a span with the given color.
+    """
+    try:
+        score_int = int(score)
+    except Exception:
+        score_int = 0
+    return f'<span style="color: #3fa45bff; font-size:1.2em;">{"★" * score_int}</span>'
+
+def make_link(mname):
+    """
+    Create a markdown link from the model identifier.
+    For example, if mname is "org/model", display "model" and link to its HF page.
+    """
+    parts = str(mname).split('/')
+    display_name = parts[1] if len(parts) > 1 else mname
+    return f'[{display_name}](https://huggingface.co/{mname})'
 
 def get_plots(task):
-    #TO DO : hover text with energy efficiency number, parameters
-    task_df= pd.read_csv('data/energy/'+task)
-    params_df = pd.read_csv('data/params/'+task)
-    params_df= params_df.rename(columns={"Link": "model"})
-    all_df = pd.merge(task_df, params_df, on='model')
-    all_df['Total GPU Energy (Wh)'] = all_df['total_gpu_energy']*1000
-    all_df = all_df.sort_values(by=['Total GPU Energy (Wh)'])
-    all_df['parameters'] = all_df['parameters'].apply(format_params)
-    all_df['energy_star'] = pd.cut(all_df['Total GPU Energy (Wh)'], 3, labels=["⭐⭐⭐", "⭐⭐", "⭐"])
-    fig = px.scatter(all_df, x="model", y='Total GPU Energy (Wh)', custom_data=['parameters'], height= 500, width= 800, color = 'energy_star', color_discrete_map={"⭐": 'red', "⭐⭐": "orange", "⭐⭐⭐": "green"})
+    """
+    Read the energy CSV for a given task and return a Plotly scatter plot.
+    The y-axis shows the total GPU energy (Wh) and the color is determined by energy_score.
+    """
+    df = pd.read_csv('data/energy/' + task)
+    # Ensure energy_score is an integer (for discrete color mapping)
+    df['energy_score'] = df['energy_score'].astype(int)
+    # Convert kWh to Wh and round to 4 decimal places.
+    df['Total GPU Energy (Wh)'] = (df['total_gpu_energy'] * 1000).round(4)
+    
+    # Define a 5-level color mapping: 1 = red, 5 = green.
+    color_map = {
+        1: "red",
+        2: "orange",
+        3: "yellow",
+        4: "lightgreen",
+        5: "green"
+    }
+    
+    fig = px.scatter(
+        df,
+        x="model",
+        y="Total GPU Energy (Wh)",
+        custom_data=['energy_score'],
+        height=500,
+        width=800,
+        color="energy_score",
+        color_discrete_map=color_map
+    )
     fig.update_traces(
-    hovertemplate="<br>".join([
-        "Total Energy: %{y}",
-        "Parameters: %{customdata[0]}"])
+        hovertemplate="<br>".join([
+            "Model: %{x}",
+            "Total Energy (Wh): %{y}",
+            "Energy Score: %{customdata[0]}"
+        ])
     )
+    fig.update_layout(xaxis_title="Model", yaxis_title="Total GPU Energy (Wh)")
     return fig
 
 def get_all_plots():
-    all_df = pd.DataFrame(columns= ['model', 'parameters', 'total_gpu_energy'])
+    """
+    Combine data from all tasks and return a scatter plot.
+    Duplicate models (if any) are dropped.
+    """
+    all_df = pd.DataFrame()
     for task in tasks:
-        task_df= pd.read_csv('data/energy/'+task)
-        params_df = pd.read_csv('data/params/'+task)
-        params_df= params_df.rename(columns={"Link": "model"})
-        tasks_df = pd.merge(task_df, params_df, on='model')
-        tasks_df= tasks_df[['model', 'parameters', 'total_gpu_energy']]
-        tasks_df['Total GPU Energy (Wh)'] = tasks_df['total_gpu_energy']*1000
-        tasks_df['energy_star'] = pd.cut(tasks_df['Total GPU Energy (Wh)'], 3, labels=["⭐⭐⭐", "⭐⭐", "⭐"])
-        all_df = pd.concat([all_df, tasks_df])
-    all_df = all_df.sort_values(by=['Total GPU Energy (Wh)'])
-    all_df['parameters'] = all_df['parameters'].apply(format_params)
-    fig = px.scatter(all_df, x="model", y='Total GPU Energy (Wh)', custom_data=['parameters'], height= 500, width= 800, color = 'energy_star', color_discrete_map={"⭐": 'red', "⭐⭐": "orange", "⭐⭐⭐": "green"})
+        df = pd.read_csv('data/energy/' + task)
+        df['energy_score'] = df['energy_score'].astype(int)
+        df['Total GPU Energy (Wh)'] = (df['total_gpu_energy'] * 1000).round(4)
+        all_df = pd.concat([all_df, df], ignore_index=True)
+    all_df = all_df.drop_duplicates(subset=['model'])
+    
+    color_map = {
+        1: "red",
+        2: "orange",
+        3: "yellow",
+        4: "lightgreen",
+        5: "green"
+    }
+    fig = px.scatter(
+        all_df,
+        x="model",
+        y="Total GPU Energy (Wh)",
+        custom_data=['energy_score'],
+        height=500,
+        width=800,
+        color="energy_score",
+        color_discrete_map=color_map
+    )
     fig.update_traces(
-    hovertemplate="<br>".join([
-        "Total Energy: %{y}",
-        "Parameters: %{customdata[0]}"])
+        hovertemplate="<br>".join([
+            "Model: %{x}",
+            "Total Energy (Wh): %{y}",
+            "Energy Score: %{customdata[0]}"
+        ])
     )
+    fig.update_layout(xaxis_title="Model", yaxis_title="Total GPU Energy (Wh)")
     return fig
 
-def make_link(mname):
-    link = "["+ str(mname).split('/')[1] +'](https://huggingface.co/'+str(mname)+")"
-    return link
-
 def get_model_names(task):
-    task_df= pd.read_csv('data/params/'+task)
-    energy_df= pd.read_csv('data/energy/'+task)
-    task_df= task_df.rename(columns={"Link": "model"})
-    all_df = pd.merge(task_df, energy_df, on='model')
-    all_df=all_df.drop_duplicates(subset=['model'])
-    all_df['Parameters'] = all_df['parameters'].apply(format_params)
-    all_df['Model'] = all_df['model'].apply(make_link)
-    all_df['Total GPU Energy (Wh)'] = all_df['total_gpu_energy']*1000
-    all_df['Total GPU Energy (Wh)'] = all_df['Total GPU Energy (Wh)'].round(2)
-    all_df['Rating'] = pd.cut(all_df['Total GPU Energy (Wh)'], 3, labels=["⭐⭐⭐", "⭐⭐", "⭐"])
-    all_df= all_df.sort_values('Total GPU Energy (Wh)')
-    model_names = all_df[['Model','Rating','Total GPU Energy (Wh)', 'Parameters']]
+    """
+    For a given task, load the energy CSV and return a dataframe with three columns:
+    - Model (a markdown link),
+    - Rating (the star rating based on energy_score),
+    - Total GPU Energy (Wh)
+    """
+    df = pd.read_csv('data/energy/' + task)
+    df['energy_score'] = df['energy_score'].astype(int)
+    df['Total GPU Energy (Wh)'] = (df['total_gpu_energy'] * 1000).round(4)
+    df['Model'] = df['model'].apply(make_link)
+    df['Rating'] = df['energy_score'].apply(format_stars)
+    df = df.sort_values(by='Total GPU Energy (Wh)')
+    model_names = df[['Model', 'Rating', 'Total GPU Energy (Wh)']]
     return model_names
 
 def get_all_model_names():
-    #TODO: add link to results in model card of each model
-    all_df = pd.DataFrame(columns = ['model', 'parameters', 'total_gpu_energy'])
+    """
+    Combine data from all tasks and return a table of models.
+    Duplicate models are dropped.
+    """
+    all_df = pd.DataFrame()
     for task in tasks:
-        task_df= pd.read_csv('data/params/'+task)
-        energy_df= pd.read_csv('data/energy/'+task)
-        task_df= task_df.rename(columns={"Link": "model"})
-        tasks_df = pd.merge(task_df, energy_df, on='model')
-        tasks_df= tasks_df[['model', 'parameters', 'total_gpu_energy']]
-        tasks_df['Total GPU Energy (Wh)'] = tasks_df['total_gpu_energy']*1000
-        tasks_df['Total GPU Energy (Wh)'] = tasks_df['Total GPU Energy (Wh)'].round(2)
-        tasks_df['Rating'] = pd.cut(tasks_df['Total GPU Energy (Wh)'], 3, labels=["⭐⭐⭐", "⭐⭐", "⭐"])
-        all_df = pd.concat([all_df, tasks_df])
-    all_df=all_df.drop_duplicates(subset=['model'])
-    all_df['Parameters'] = all_df['parameters'].apply(format_params)
-    all_df['Model'] = all_df['model'].apply(make_link)
-    all_df= all_df.sort_values('Total GPU Energy (Wh)')
-    model_names = all_df[['Model','Rating','Total GPU Energy (Wh)', 'Parameters']]
+        df = pd.read_csv('data/energy/' + task)
+        df['energy_score'] = df['energy_score'].astype(int)
+        df['Total GPU Energy (Wh)'] = (df['total_gpu_energy'] * 1000).round(4)
+        df['Model'] = df['model'].apply(make_link)
+        df['Rating'] = df['energy_score'].apply(format_stars)
+        all_df = pd.concat([all_df, df], ignore_index=True)
+    all_df = all_df.drop_duplicates(subset=['model'])
+    all_df = all_df.sort_values(by='Total GPU Energy (Wh)')
+    model_names = all_df[['Model', 'Rating', 'Total GPU Energy (Wh)']]
     return model_names
 
-
-def format_params(num):
-    if num > 1000000000:
-        if not num % 1000000000:
-            return f'{num // 1000000000}B'
-        return f'{round(num / 1000000000, 1)}B'
-    return f'{num // 1000000}M'
-
-
-
+# Build the Gradio interface.
 demo = gr.Blocks()
 
 with demo:
     gr.Markdown(
         """# AI Energy Score Leaderboard - v.0 (2024) 🌎 💻 🌟
-    ### Welcome to the leaderboard for the [AI Energy Score Project!](https://huggingface.co/EnergyStarAI)
-    Click through the tasks below to see how different models measure up in terms of energy efficiency"""
+### Welcome to the leaderboard for the [AI Energy Score Project!](https://huggingface.co/EnergyStarAI)
+Click through the tasks below to see how different models measure up in terms of energy efficiency."""
     )
     gr.Markdown(
         """Test your own models via the [submission portal](https://huggingface.co/spaces/AIEnergyScore/submission_portal)!"""
-        )
+    )
+    
     with gr.Tabs():
         with gr.TabItem("Text Generation 💬"):
             with gr.Row():
@@ -122,91 +179,87 @@ with demo:
                     plot = gr.Plot(get_plots('text_generation.csv'))
                 with gr.Column(scale=1):
                     table = gr.Dataframe(get_model_names('text_generation.csv'), datatype="markdown")
-
+                    
         with gr.TabItem("Image Generation 📷"):
             with gr.Row():
                 with gr.Column():
                     plot = gr.Plot(get_plots('image_generation.csv'))
                 with gr.Column():
                     table = gr.Dataframe(get_model_names('image_generation.csv'), datatype="markdown")
-
+                    
         with gr.TabItem("Text Classification 🎭"):
             with gr.Row():
                 with gr.Column():
                     plot = gr.Plot(get_plots('text_classification.csv'))
                 with gr.Column():
                     table = gr.Dataframe(get_model_names('text_classification.csv'), datatype="markdown")
-
+                    
         with gr.TabItem("Image Classification 🖼️"):
             with gr.Row():
                 with gr.Column():
                     plot = gr.Plot(get_plots('image_classification.csv'))
                 with gr.Column():
                     table = gr.Dataframe(get_model_names('image_classification.csv'), datatype="markdown")
-
+                    
         with gr.TabItem("Image Captioning 📝"):
             with gr.Row():
                 with gr.Column():
                     plot = gr.Plot(get_plots('image_captioning.csv'))
                 with gr.Column():
                     table = gr.Dataframe(get_model_names('image_captioning.csv'), datatype="markdown")
+                    
         with gr.TabItem("Summarization 📃"):
             with gr.Row():
                 with gr.Column():
                     plot = gr.Plot(get_plots('summarization.csv'))
                 with gr.Column():
                     table = gr.Dataframe(get_model_names('summarization.csv'), datatype="markdown")
-
-        with gr.TabItem("Automatic Speech Recognition 💬 "):
+                    
+        with gr.TabItem("Automatic Speech Recognition 💬"):
             with gr.Row():
                 with gr.Column():
                     plot = gr.Plot(get_plots('asr.csv'))
                 with gr.Column():
                     table = gr.Dataframe(get_model_names('asr.csv'), datatype="markdown")
-
+                    
         with gr.TabItem("Object Detection 🚘"):
             with gr.Row():
                 with gr.Column():
                     plot = gr.Plot(get_plots('object_detection.csv'))
                 with gr.Column():
                     table = gr.Dataframe(get_model_names('object_detection.csv'), datatype="markdown")
-
+                    
         with gr.TabItem("Sentence Similarity 📚"):
             with gr.Row():
                 with gr.Column():
                     plot = gr.Plot(get_plots('sentence_similarity.csv'))
                 with gr.Column():
                     table = gr.Dataframe(get_model_names('sentence_similarity.csv'), datatype="markdown")
-
+                    
         with gr.TabItem("Extractive QA ❔"):
             with gr.Row():
                 with gr.Column():
                     plot = gr.Plot(get_plots('question_answering.csv'))
                 with gr.Column():
                     table = gr.Dataframe(get_model_names('question_answering.csv'), datatype="markdown")
-
+                    
         with gr.TabItem("All Tasks 💡"):
             with gr.Row():
                 with gr.Column():
                     plot = gr.Plot(get_all_plots)
                 with gr.Column():
                     table = gr.Dataframe(get_all_model_names, datatype="markdown")
-
-    with gr.Accordion("Methodology", open = False):
-        gr.Markdown(
-        """For each of the ten tasks above, we created a custom dataset with 1,000 entries (see all of the datasets on our [org Hub page](https://huggingface.co/EnergyStarAI)).
-        We then tested each of the models from the leaderboard on the appropriate task on Nvidia H100 GPUs, measuring the energy consumed using [Code Carbon](https://mlco2.github.io/codecarbon/), an open-source Python package for tracking the environmental impacts of code.
-        We developed and used a [Docker container](https://github.com/huggingface/EnergyStarAI/) to maximize the reproducibility of results, and to enable members of the community to benchmark internal models.
-        Reach out to us if you want to collaborate!
-        """)
+                    
     with gr.Accordion("📙 Citation", open=False):
-            citation_button = gr.Textbox(
-                value=CITATION_BUTTON_TEXT,
-                label=CITATION_BUTTON_LABEL,
-                elem_id="citation-button",
-                lines=10,
-                show_copy_button=True,
-            )
+        citation_button = gr.Textbox(
+            value=CITATION_BUTTON_TEXT,
+            label=CITATION_BUTTON_LABEL,
+            elem_id="citation-button",
+            lines=10,
+            show_copy_button=True,
+        )
     gr.Markdown(
-    """Last updated: October 1st, 2024 by [Sasha Luccioni](https://huggingface.co/sasha)""")
+        """Last updated: February 2025"""
+    )
+
 demo.launch()