Update app.py

app.py

@@ -39,22 +39,21 @@ def make_link(mname):
 
 def get_plots(task):
     df = pd.read_csv('data/energy/' + task)
-    # Remove extra unnamed column if present
     if df.columns[0].startswith("Unnamed:"):
         df = df.iloc[:, 1:]
-    
-    # Convert energy_score to int and then to str so it's treated as categorical
+    # Ensure total_gpu_energy is a float so that values are not misinterpreted
+    df['total_gpu_energy'] = df['total_gpu_energy'].astype(float)
+    # Convert energy_score to a categorical string for discrete coloring
     df['energy_score'] = df['energy_score'].astype(int).astype(str)
     df['Display Model'] = df['model'].apply(lambda m: m.split('/')[-1])
     
-    # Update color_map keys to be strings
     color_map = {"1": "red", "2": "orange", "3": "yellow", "4": "lightgreen", "5": "green"}
     
     fig = px.scatter(
         df,
-        x="total_gpu_energy",  # x-axis: GPU energy consumption
-        y="Display Model",     # y-axis: Model name for display
-        color="energy_score",  # Discrete color based on energy score
+        x="total_gpu_energy",
+        y="Display Model",
+        color="energy_score",
         custom_data=['energy_score'],
         height=500,
         width=800,
@@ -76,7 +75,7 @@ def get_all_plots():
         df = pd.read_csv('data/energy/' + task)
         if df.columns[0].startswith("Unnamed:"):
             df = df.iloc[:, 1:]
-        # Convert energy_score to categorical string
+        df['total_gpu_energy'] = df['total_gpu_energy'].astype(float)
         df['energy_score'] = df['energy_score'].astype(int).astype(str)
         df['Display Model'] = df['model'].apply(lambda m: m.split('/')[-1])
         all_df = pd.concat([all_df, df], ignore_index=True)
@@ -86,9 +85,9 @@ def get_all_plots():
     
     fig = px.scatter(
         all_df,
-        x="total_gpu_energy",  # x-axis: GPU energy consumption
+        x="total_gpu_energy",
         y="Display Model",
-        color="energy_score",  # Discrete color mapping
+        color="energy_score",
         custom_data=['energy_score'],
         height=500,
         width=800,
@@ -105,20 +104,11 @@ def get_all_plots():
     return fig
 
 def get_model_names(task):
-    """
-    For a given task, load the energy CSV and return a dataframe with the following columns:
-      - Model (a markdown link)
-      - GPU Energy (Wh) formatted as a string with 4 decimal places
-      - Score (a star rating based on energy_score)
-    For text_generation.csv only, also add the "Class" column from the CSV.
-    The final column order is: Model, GPU Energy (Wh), Score, [Class].
-    """
     df = pd.read_csv('data/energy/' + task)
     if df.columns[0].startswith("Unnamed:"):
         df = df.iloc[:, 1:]
     df['energy_score'] = df['energy_score'].astype(int)
-    # Format the energy as a string with 4 decimals
-    df['GPU Energy (Wh)'] = df['total_gpu_energy'].apply(lambda x: f"{x:.4f}")
+    df['GPU Energy (Wh)'] = df['total_gpu_energy'].astype(float).apply(lambda x: f"{x:.4f}")
     df['Model'] = df['model'].apply(make_link)
     df['Score'] = df['energy_score'].apply(format_stars)
     
@@ -132,16 +122,11 @@ def get_model_names(task):
     return df
 
 def get_all_model_names():
-    """
-    Combine data from all tasks and return a leaderboard table with:
-      - Model, GPU Energy (Wh), Score
-    Duplicate models are dropped.
-    """
     all_df = pd.DataFrame()
     for task in tasks:
         df = pd.read_csv('data/energy/' + task)
         df['energy_score'] = df['energy_score'].astype(int)
-        df['GPU Energy (Wh)'] = df['total_gpu_energy'].apply(lambda x: f"{x:.4f}")
+        df['GPU Energy (Wh)'] = df['total_gpu_energy'].astype(float).apply(lambda x: f"{x:.4f}")
         df['Model'] = df['model'].apply(make_link)
         df['Score'] = df['energy_score'].apply(format_stars)
         all_df = pd.concat([all_df, df], ignore_index=True)
@@ -149,8 +134,64 @@ def get_all_model_names():
     all_df = all_df.sort_values(by='GPU Energy (Wh)')
     return all_df[['Model', 'GPU Energy (Wh)', 'Score']]
 
-# Build the Gradio interface.
-# The css argument below makes all tables (e.g. leaderboard) use a fixed layout with narrower columns.
+# === New functions for Text Generation filtering by model class ===
+def get_text_generation_plots(model_class):
+    df = pd.read_csv('data/energy/text_generation.csv')
+    if df.columns[0].startswith("Unnamed:"):
+        df = df.iloc[:, 1:]
+    # Filter to the selected model class (if a "class" column exists)
+    if 'class' in df.columns:
+        df = df[df['class'] == model_class]
+    df['total_gpu_energy'] = df['total_gpu_energy'].astype(float)
+    df['energy_score'] = df['energy_score'].astype(int).astype(str)
+    df['Display Model'] = df['model'].apply(lambda m: m.split('/')[-1])
+    
+    color_map = {"1": "red", "2": "orange", "3": "yellow", "4": "lightgreen", "5": "green"}
+    
+    fig = px.scatter(
+        df,
+        x="total_gpu_energy",
+        y="Display Model",
+        color="energy_score",
+        custom_data=['energy_score'],
+        height=500,
+        width=800,
+        color_discrete_map=color_map
+    )
+    fig.update_traces(
+        hovertemplate="<br>".join([
+            "Model: %{y}",
+            "GPU Energy (Wh): %{x}",
+            "Energy Score: %{customdata[0]}"
+        ])
+    )
+    fig.update_layout(xaxis_title="GPU Energy (Wh)", yaxis_title="Model")
+    return fig
+
+def get_text_generation_model_names(model_class):
+    df = pd.read_csv('data/energy/text_generation.csv')
+    if df.columns[0].startswith("Unnamed:"):
+        df = df.iloc[:, 1:]
+    if 'class' in df.columns:
+        df = df[df['class'] == model_class]
+    df['energy_score'] = df['energy_score'].astype(int)
+    df['GPU Energy (Wh)'] = df['total_gpu_energy'].astype(float).apply(lambda x: f"{x:.4f}")
+    df['Model'] = df['model'].apply(make_link)
+    df['Score'] = df['energy_score'].apply(format_stars)
+    if 'class' in df.columns:
+        df['Class'] = df['class']
+        df = df[['Model', 'GPU Energy (Wh)', 'Score', 'Class']]
+    else:
+        df = df[['Model', 'GPU Energy (Wh)', 'Score']]
+    df = df.sort_values(by='GPU Energy (Wh)')
+    return df
+
+def update_text_generation(model_class):
+    plot = get_text_generation_plots(model_class)
+    table = get_text_generation_model_names(model_class)
+    return plot, table
+
+# === Build the Gradio Interface ===
 demo = gr.Blocks(css="""
 .gr-dataframe table {
     table-layout: fixed;
@@ -175,13 +216,20 @@ Click through the tasks below to see how different models measure up in terms of
     )
     
     with gr.Tabs():
+        # --- Text Generation Tab with Dropdown for Model Class ---
         with gr.TabItem("Text Generation 💬"):
             with gr.Row():
                 with gr.Column(scale=1.3):
-                    plot = gr.Plot(get_plots('text_generation.csv'))
+                    tg_plot = gr.Plot(get_text_generation_plots("A"))
                 with gr.Column(scale=1):
-                    table = gr.Dataframe(get_model_names('text_generation.csv'), datatype="markdown")
-                    
+                    tg_table = gr.Dataframe(get_text_generation_model_names("A"), datatype="markdown")
+            model_class_dropdown = gr.Dropdown(choices=["A", "B", "C"],
+                                                 label="Select Model Class",
+                                                 value="A")
+            model_class_dropdown.change(fn=update_text_generation,
+                                        inputs=model_class_dropdown,
+                                        outputs=[tg_plot, tg_table])
+        
         with gr.TabItem("Image Generation 📷"):
             with gr.Row():
                 with gr.Column():
@@ -248,7 +296,6 @@ Click through the tasks below to see how different models measure up in terms of
         with gr.TabItem("All Tasks 💡"):
             with gr.Row():
                 with gr.Column():
-                    # Call the functions to generate the plot and table
                     plot = gr.Plot(get_all_plots())
                 with gr.Column():
                     table = gr.Dataframe(get_all_model_names(), datatype="markdown")