Update app.py

app.py

@@ -37,16 +37,9 @@ def make_link(mname):
     display_name = parts[1] if len(parts) > 1 else mname
     return f'[{display_name}](https://huggingface.co/{mname})'
 
-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 expected "total_gpu_energy" column is missing, drop the first column.
-    if "total_gpu_energy" not in df.columns:
-        df = df.iloc[:, 1:]
-    return df
-
 def get_plots(task):
-    df = read_csv_drop_extra(task)
+    # Read CSV using the first column as index so that only the useful columns remain.
+    df = pd.read_csv('data/energy/' + task, index_col=0)
     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])
@@ -55,9 +48,9 @@ def get_plots(task):
     
     fig = px.scatter(
         df,
-        x="total_gpu_energy",  # Use the correct energy column
+        x="total_gpu_energy",
         y="Display Model",
-        color="energy_score",  # Map energy score to the color
+        color="energy_score",
         custom_data=['model', 'total_gpu_energy', 'energy_score'],
         height=500,
         width=800,
@@ -80,7 +73,7 @@ def get_plots(task):
 def get_all_plots():
     all_df = pd.DataFrame()
     for task in tasks:
-        df = read_csv_drop_extra(task)
+        df = pd.read_csv('data/energy/' + task, index_col=0)
         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])
@@ -121,7 +114,7 @@ def get_model_names(task):
       - Score (a star rating based on energy_score)
     For text_generation.csv only, also add the "Class" column if present.
     """
-    df = read_csv_drop_extra(task)
+    df = pd.read_csv('data/energy/' + task, index_col=0)
     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}")
@@ -134,7 +127,7 @@ def get_model_names(task):
     else:
         df = df[['Model', 'GPU Energy (Wh)', 'Score']]
         
-    # Sort by the numeric value (not the formatted string)
+    # Now sort by the numeric value in total_gpu_energy.
     df = df.sort_values(by='total_gpu_energy')
     return df
 
@@ -146,7 +139,7 @@ def get_all_model_names():
     """
     all_df = pd.DataFrame()
     for task in tasks:
-        df = read_csv_drop_extra(task)
+        df = pd.read_csv('data/energy/' + task, index_col=0)
         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}")
@@ -184,7 +177,6 @@ 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():
-                # 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):