Update app.py

app.py

@@ -30,40 +30,34 @@ def format_stars(score):
         score_int = int(score)
     except Exception:
         score_int = 0
-    # Display a star rating (★) based on the energy score.
     return f'<span style="color: #3fa45bff; font-size:2em;">{"★" * score_int}</span>'
 
 def make_link(mname):
-    # Make a Markdown link from the model name.
     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):
-    # Read the CSV for the given task.
+def read_csv_drop_extra(task):
+    """Helper to load CSV and drop the first column if necessary."""
     df = pd.read_csv('data/energy/' + task)
-    # If the first column is unnamed (the extra blank column), drop it.
-    if df.columns[0].startswith("Unnamed:"):
+    # If the expected "total_gpu_energy" column is missing, drop the first column.
+    if "total_gpu_energy" not in df.columns:
         df = df.iloc[:, 1:]
-    # Convert the numeric columns
+    return df
+
+def get_plots(task):
+    df = read_csv_drop_extra(task)
     df['total_gpu_energy'] = pd.to_numeric(df['total_gpu_energy'], errors='coerce')
     df['energy_score'] = pd.to_numeric(df['energy_score'], errors='coerce').astype(int)
-    # Create a short version of the model name for display on the y-axis.
     df['Display Model'] = df['model'].apply(lambda m: m.split('/')[-1])
     
-    # Define a discrete color mapping for energy scores.
     color_map = {1: "red", 2: "orange", 3: "yellow", 4: "lightgreen", 5: "green"}
     
-    # Build a scatter plot:
-    #  - x-axis: total_gpu_energy
-    #  - y-axis: Display Model (short model name)
-    #  - Color: energy_score
-    #  - Custom tooltip will include the full model name, energy value and energy score.
     fig = px.scatter(
         df,
-        x="total_gpu_energy",
+        x="total_gpu_energy",  # Use the correct energy column
         y="Display Model",
-        color="energy_score",
+        color="energy_score",  # Map energy score to the color
         custom_data=['model', 'total_gpu_energy', 'energy_score'],
         height=500,
         width=800,
@@ -84,12 +78,9 @@ def get_plots(task):
     return fig
 
 def get_all_plots():
-    # Combine data from all tasks.
     all_df = pd.DataFrame()
     for task in tasks:
-        df = pd.read_csv('data/energy/' + task)
-        if df.columns[0].startswith("Unnamed:"):
-            df = df.iloc[:, 1:]
+        df = read_csv_drop_extra(task)
         df['total_gpu_energy'] = pd.to_numeric(df['total_gpu_energy'], errors='coerce')
         df['energy_score'] = pd.to_numeric(df['energy_score'], errors='coerce').astype(int)
         df['Display Model'] = df['model'].apply(lambda m: m.split('/')[-1])
@@ -126,13 +117,11 @@ 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 to 4 decimal places)
+      - GPU Energy (Wh) formatted to 4 decimal places
       - Score (a star rating based on energy_score)
-    For text_generation.csv only, also include the "Class" column if it exists.
+    For text_generation.csv only, also add the "Class" column if present.
     """
-    df = pd.read_csv('data/energy/' + task)
-    if df.columns[0].startswith("Unnamed:"):
-        df = df.iloc[:, 1:]
+    df = read_csv_drop_extra(task)
     df['total_gpu_energy'] = pd.to_numeric(df['total_gpu_energy'], errors='coerce')
     df['energy_score'] = pd.to_numeric(df['energy_score'], errors='coerce').astype(int)
     df['GPU Energy (Wh)'] = df['total_gpu_energy'].apply(lambda x: f"{x:.4f}")
@@ -145,7 +134,7 @@ def get_model_names(task):
     else:
         df = df[['Model', 'GPU Energy (Wh)', 'Score']]
         
-    # Sort by the numeric energy value.
+    # Sort by the numeric value (not the formatted string)
     df = df.sort_values(by='total_gpu_energy')
     return df
 
@@ -157,9 +146,7 @@ def get_all_model_names():
     """
     all_df = pd.DataFrame()
     for task in tasks:
-        df = pd.read_csv('data/energy/' + task)
-        if df.columns[0].startswith("Unnamed:"):
-            df = df.iloc[:, 1:]
+        df = read_csv_drop_extra(task)
         df['total_gpu_energy'] = pd.to_numeric(df['total_gpu_energy'], errors='coerce')
         df['energy_score'] = pd.to_numeric(df['energy_score'], errors='coerce').astype(int)
         df['GPU Energy (Wh)'] = df['total_gpu_energy'].apply(lambda x: f"{x:.4f}")
@@ -171,7 +158,6 @@ def get_all_model_names():
     return all_df[['Model', 'GPU Energy (Wh)', 'Score']]
 
 # Build the Gradio interface.
-# The CSS below sets fixed layouts for tables.
 demo = gr.Blocks(css="""
 .gr-dataframe table {
     table-layout: fixed;
@@ -198,7 +184,8 @@ Click through the tasks below to see how different models measure up in terms of
     with gr.Tabs():
         with gr.TabItem("Text Generation 💬"):
             with gr.Row():
-                with gr.Column(scale=1.3):
+                # Changed scale to an integer (2 vs 1) to avoid warnings.
+                with gr.Column(scale=2):
                     plot = gr.Plot(get_plots('text_generation.csv'))
                 with gr.Column(scale=1):
                     table = gr.Dataframe(get_model_names('text_generation.csv'), datatype="markdown")