Update app.py

app.py

@@ -39,16 +39,22 @@ def make_link(mname):
 
 def get_plots(task):
     df = pd.read_csv('data/energy/' + task)
-    df['energy_score'] = df['energy_score'].astype(int)
+    # 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
+    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"}
+    # 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",  # Ensure correct column for x-axis
-        y="Display Model",  # Keep model name for y-axis
-        color="energy_score",  # Ensure correct column for point color
+        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
         custom_data=['energy_score'],
         height=500,
         width=800,
@@ -70,18 +76,19 @@ def get_all_plots():
         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)
+        # Convert energy_score to categorical string
+        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)
     all_df = all_df.drop_duplicates(subset=['model'])
     
-    color_map = {1: "red", 2: "orange", 3: "yellow", 4: "lightgreen", 5: "green"}
+    color_map = {"1": "red", "2": "orange", "3": "yellow", "4": "lightgreen", "5": "green"}
     
     fig = px.scatter(
         all_df,
-        x="total_gpu_energy",  # Ensure correct column for x-axis
+        x="total_gpu_energy",  # x-axis: GPU energy consumption
         y="Display Model",
-        color="energy_score",  # Ensure correct column for point color
+        color="energy_score",  # Discrete color mapping
         custom_data=['energy_score'],
         height=500,
         width=800,
@@ -241,9 +248,10 @@ 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():
-                    plot = gr.Plot(get_all_plots)
+                    # 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")
+                    table = gr.Dataframe(get_all_model_names(), datatype="markdown")
                     
     with gr.Accordion("📙 Citation", open=False):
         citation_button = gr.Textbox(
@@ -257,4 +265,4 @@ Click through the tasks below to see how different models measure up in terms of
         """Last updated: February 2025"""
     )
 
-demo.launch()
+demo.launch()