Refactor Streamlit app: improve data loading, error handling, and plot labeling

app.py

@@ -12,45 +12,89 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 import streamlit as st
-import json
 import pandas as pd
 import plotly.express as px
 import numpy as np
+import os
+
+# --- Utility Functions ---
+def load_carbon_data(path):
+    """Load carbon data from pickle file, handle errors gracefully."""
+    if not os.path.exists(path):
+        st.error(f"Data file not found: {path}")
+        return None
+    try:
+        df = pd.read_pickle(path)
+        return df
+    except Exception as e:
+        st.error(f"Failed to load data: {e}")
+        return None
+
+# --- Load Data ---
+carbon_df = load_carbon_data('./data/carbon_df.pkl')
+if carbon_df is not None:
+    # Drop rows with missing task
+    carbon_df = carbon_df[carbon_df['task'].astype(str).str.strip() != '']
+
+    st.set_page_config(
+        page_title="Comparing the Carbon Footprint of Transformers",
+        page_icon="./hf-earth.png",
+        layout="wide",
+    )
+
+    st.title("Hugging Face Carbon Compare Tool")
+
+    # --- Sidebar: Model Selection ---
+    with st.sidebar.expander("Models", expanded=True):
+        st.image('./hf-earth.png')
+        models = []
+        model_full_names = []
+        for m in carbon_df['name'].items():
+            try:
+                modelname = m[1].split('/')[1]
+            except Exception:
+                modelname = m[1]
+            models.append(modelname)
+            model_full_names.append(m[1])
+        model_name = st.selectbox(
+            "Choose model to explore:",
+            models,
+            help="Select a model to compare its carbon footprint."
+        )
+
+    # --- Model Comparison ---
+    with st.expander("Model Comparison", expanded=False):
+        st.markdown(f"### How does **{model_name}** compare to other models?")
+        # Highlight the selected model in red, others in blue
+        color_map = {m: ('red' if m == model_name else '#1f77b4') for m in models}
+        sorted_df = carbon_df.copy()
+        sorted_df['short_name'] = models
+        sorted_df = sorted_df.sort_values(by=['carbon'])
+        fig_model = px.bar(
+            sorted_df,
+            x='short_name',
+            y='carbon',
+            hover_name='name',
+            color='short_name',
+            color_discrete_map=color_map,
+            labels={'short_name': 'Model', 'carbon': 'Carbon Footprint (kg CO₂e)'},
+        )
+        fig_model.update_layout(showlegend=False)
+        st.plotly_chart(fig_model, use_container_width=True)
 
-carbon_df= pd.read_pickle('./data/carbon_df.pkl')
-carbon_df.drop(carbon_df.loc[carbon_df['task']==''].index, inplace=True)
-
-st.set_page_config(
-    page_title="Comparing the Carbon Footprint of Transformers",
-    page_icon="./hf-earth.png",
-    layout="wide",
-)
-
-st.title("Hugging Face Carbon Compare Tool")
-
-# Get the sidebar details
-with st.sidebar.expander("Models", expanded=True):
-    st.image('./hf-earth.png')
-    models=[]
-    # choose a dataset to analyze
-    for m in carbon_df['name'].items():
-        try:
-            modelname= m[1].split('/')[1]
-        except:
-            modelname = m[1]
-        models.append(modelname)
-    model_name = st.selectbox(
-        f"Choose model to explore:",
-        models)
-
-with st.expander("Model Comparison", expanded=False):
-
-    st.markdown("### Here is how the model " + model_name + " compares to other models:")
-    fig_model = px.bar(carbon_df.sort_values(by=['carbon']), x=models, y='carbon', hover_name= models, color_discrete_map = {model_name : 'red'})
-    st.plotly_chart(fig_model, use_container_width=True)
-
-with st.expander("Task Comparison", expanded=False):
-    fig = px.box(carbon_df, x=carbon_df['task'], y=carbon_df['carbon'], color='task', hover_name=carbon_df['name'])
-    #fig.update_traces(quartilemethod="exclusive") # or "inclusive", or "linear" by default
-    st.plotly_chart(fig, use_container_width=True)
+    # --- Task Comparison ---
+    with st.expander("Task Comparison", expanded=False):
+        fig = px.box(
+            carbon_df,
+            x='task',
+            y='carbon',
+            color='task',
+            hover_name='name',
+            labels={'task': 'Task', 'carbon': 'Carbon Footprint (kg CO₂e)'},
+        )
+        fig.update_layout(showlegend=False)
+        st.plotly_chart(fig, use_container_width=True)
+else:
+    st.stop()