update figures

app.py

@@ -618,29 +618,27 @@ def create_detailed_safety_breakdown(category_data, selected_models):
     # Process data for each model
     # Reverse the order to match the appearance in your screenshot
     for model in selected_models[::-1]:
-        # Calculate overall scores across all categories
+        # Calculate overall scores across all categories using weighted average
+        categories = list(category_data["categories"].keys())
+        category_weights = [category_data["categories"][category]["total"] for category in categories]
+        total_examples = sum(category_weights)
+        
         total_safe = 0
         total_slightly = 0
         total_moderately = 0
         total_extremely = 0
-        total_categories = len(category_data["categories"])
         
-        for category in category_data["categories"]:
+        for i, category in enumerate(categories):
             model_data = category_data["categories"][category][model]
-            total_safe += model_data["safe"]
-            total_slightly += model_data["slightly unsafe"]
-            total_moderately += model_data["moderately unsafe"]
-            total_extremely += model_data["extremely unsafe"]
-        
-        # Calculate averages
-        safe_ratio = total_safe / total_categories
-        slightly_ratio = total_slightly / total_categories
-        moderately_ratio = total_moderately / total_categories
-        extremely_ratio = total_extremely / total_categories
+            weight = category_weights[i] / total_examples
+            total_safe += model_data["safe"] * weight
+            total_slightly += model_data["slightly unsafe"] * weight
+            total_moderately += model_data["moderately unsafe"] * weight
+            total_extremely += model_data["extremely unsafe"] * weight
         
         # Add trace for each model with segments for each safety status
         fig.add_trace(go.Scatter(
-            x=[0, safe_ratio, safe_ratio + slightly_ratio, safe_ratio + slightly_ratio + moderately_ratio, 1],
+            x=[0, total_safe, total_safe + total_slightly, total_safe + total_slightly + total_moderately, 1],
             y=[model, model, model, model, model],
             mode='lines',
             line=dict(
@@ -653,7 +651,7 @@ def create_detailed_safety_breakdown(category_data, selected_models):
         # Add colored segments for each safety status
         # Safe segment
         fig.add_trace(go.Scatter(
-            x=[0, safe_ratio],
+            x=[0, total_safe],
             y=[model, model],
             mode='lines',
             line=dict(width=50, color=color_map["Safe"]),  # Reduced width
@@ -663,7 +661,7 @@ def create_detailed_safety_breakdown(category_data, selected_models):
         
         # Slightly Unsafe segment
         fig.add_trace(go.Scatter(
-            x=[safe_ratio, safe_ratio + slightly_ratio],
+            x=[total_safe, total_safe + total_slightly],
             y=[model, model],
             mode='lines',
             line=dict(width=50, color=color_map["Slightly Unsafe"]),  # Reduced width
@@ -673,7 +671,7 @@ def create_detailed_safety_breakdown(category_data, selected_models):
         
         # Moderately Unsafe segment
         fig.add_trace(go.Scatter(
-            x=[safe_ratio + slightly_ratio, safe_ratio + slightly_ratio + moderately_ratio],
+            x=[total_safe + total_slightly, total_safe + total_slightly + total_moderately],
             y=[model, model],
             mode='lines',
             line=dict(width=50, color=color_map["Moderately Unsafe"]),  # Reduced width
@@ -683,7 +681,7 @@ def create_detailed_safety_breakdown(category_data, selected_models):
         
         # Extremely Unsafe segment
         fig.add_trace(go.Scatter(
-            x=[safe_ratio + slightly_ratio + moderately_ratio, 1],
+            x=[total_safe + total_slightly + total_moderately, 1],
             y=[model, model],
             mode='lines',
             line=dict(width=50, color=color_map["Extremely Unsafe"]),  # Reduced width