Update app.py

app.py

@@ -26,12 +26,6 @@ tasks = [
 ]
 
 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 styled with color #3fa45bff
-    and with a font size increased to 2em.
-    The '!important' rules force the styling immediately.
-    """
     try:
         score_int = int(score)
     except Exception:
@@ -39,47 +33,27 @@ def format_stars(score):
     return f'<span style="color: #3fa45bff !important; font-size:2em !important;">{"★" * score_int}</span>'
 
 def make_link(mname):
-    """
-    Create a markdown link for the model.
-    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):
-    """
-    Read the energy CSV for a given task and return a Plotly scatter plot.
-    The x-axis uses the 'total_gpu_energy' column (rounded to 4 decimals) and 
-    the y-axis displays only the model name (extracted from the 'model' column).
-    """
     df = pd.read_csv('data/energy/' + task)
-    # If an extra unnamed index column exists, drop it.
     if df.columns[0].startswith("Unnamed:"):
         df = df.iloc[:, 1:]
     df['energy_score'] = df['energy_score'].astype(int)
-    # Use the correct column: "total_gpu_energy"
-    df['GPU Energy (Wh)'] = df['total_gpu_energy'].round(4)
-    # Create a column that displays only the model name (the part after '/')
     df['Display Model'] = df['model'].apply(lambda m: m.split('/')[-1])
     
-    # Define a 5-level color mapping: 1 = red, 5 = green.
-    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(
         df,
-        x="GPU Energy (Wh)",
-        y="Display Model",
+        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
         custom_data=['energy_score'],
         height=500,
         width=800,
-        color="energy_score",
         color_discrete_map=color_map
     )
     fig.update_traces(
@@ -93,36 +67,26 @@ def get_plots(task):
     return fig
 
 def get_all_plots():
-    """
-    Combine data from all tasks and return a scatter plot.
-    Duplicate models are dropped.
-    """
     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['energy_score'] = df['energy_score'].astype(int)
-        df['GPU Energy (Wh)'] = df['total_gpu_energy'].round(4)
         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="GPU Energy (Wh)",
+        x="total_gpu_energy",  # Ensure correct column for x-axis
         y="Display Model",
+        color="energy_score",  # Ensure correct column for point color
         custom_data=['energy_score'],
         height=500,
         width=800,
-        color="energy_score",
         color_discrete_map=color_map
     )
     fig.update_traces(