Update app.py

app.py

@@ -109,7 +109,9 @@ def filter_leaderboard(df, model_name, sort_by):
     return filtered_df
 
 
-def create_scatter_plot(df, x_axis, y_axis):
+def create_scatter_plot(df, x_axis, y_axis, model_filter, highlight_filter):
+    selected_color = 'orange'
+    
     fig = px.scatter(
         df,
         x=x_axis,
@@ -120,10 +122,27 @@ def create_scatter_plot(df, x_axis, y_axis):
         trendline='ols',
         trendline_options=dict(log_x=True, log_y=True),
         color='highlighted',
-        color_discrete_map={True: 'orange', False: 'blue'},
+        color_discrete_map={True: selected_color, False: 'blue'},
         title=f'{y_axis} vs {x_axis}'
     )
-    fig.update_layout(showlegend=False)
+
+    # Create legend labels
+    legend_labels = {}
+    if highlight_filter:
+        legend_labels[True] = f'{highlight_filter}'
+        legend_labels[False] = f'{model_filter or "all models"}'
+    else:
+        legend_labels[False] = f'{model_filter or "all models"}'
+    
+    # Update legend
+    for trace in fig.data:
+        if isinstance(trace.marker.color, str):  # This is for the scatter traces
+            trace.name = legend_labels.get(trace.marker.color == selected_color, '')
+    
+    fig.update_layout(
+        showlegend=True,
+        legend_title_text='Model Selection'
+    )
     
     return fig
 
@@ -160,7 +179,7 @@ def update_leaderboard_and_plot(
     else:
         combined_df = filtered_df
            
-    fig = create_scatter_plot(combined_df, x_axis, y_axis)
+    fig = create_scatter_plot(combined_df, x_axis, y_axis, model_name, highlight_name)
     display_df = combined_df.drop(columns=['highlighted'])
     display_df = display_df.style.apply(lambda x: ['background-color: #FFA500' if combined_df.loc[x.name, 'highlighted'] else '' for _ in x], axis=1).format(precision=2)
     return display_df, fig