made code simpler and less repetitive for the plots

app.py

@@ -8,7 +8,7 @@ from src.display.css_html_js import custom_css
 from src.display.utils import BenchRawColumn, fields
 from src.envs import API, EVAL_RESULTS_PATH, REPO_ID, RESULTS_REPO, TOKEN
 from src.populate import get_leaderboard_df_from_hf_dataset, summarize_model_task_type_performance
-from src.display.plotting import extract_mean_std, prepare_leaderboard_df, make_wrapped_task_plot, make_wrapped_model_plot, plot_task_metric, plot_model_across_tasks
+from src.display.plotting import METRICS_FOR_PLOTS, extract_mean_std, prepare_leaderboard_df, make_plot_wrapper, plot_metric_bar
 
 def restart_space():
     API.restart_space(repo_id=REPO_ID)
@@ -27,8 +27,9 @@ try:
     LEADERBOARD_DF = get_leaderboard_df_from_hf_dataset(EVAL_RESULTS_PATH)
     LEADERBOARD_DF = summarize_model_task_type_performance(LEADERBOARD_DF)
     LEADERBOARD_DF = prepare_leaderboard_df(LEADERBOARD_DF)
-    wrapped_task_plot = make_wrapped_task_plot(LEADERBOARD_DF)
-    wrapped_model_plot = make_wrapped_model_plot(LEADERBOARD_DF)
+    
+    wrapped_task_plot = make_plot_wrapper(leaderboard_df=LEADERBOARD_DF, group_by="Model", filter_col="Task", orientation="h")
+    wrapped_model_plot = make_plot_wrapper(leaderboard_df=LEADERBOARD_DF, group_by="Task", filter_col="Model", orientation="v")
 
 except Exception:
     restart_space()
@@ -80,11 +81,10 @@ with demo:
                 gr.Markdown("Visualize model performance for a specific task and metric")
 
                 task_choices = sorted(LEADERBOARD_DF["Task"].dropna().unique())
-                metric_choices = ["Accuracy", "MCC"]
 
                 with gr.Row():
                     task_dropdown = gr.Dropdown(choices=task_choices, label="Select Task")
-                    metric_dropdown = gr.Dropdown(choices=metric_choices, value="Accuracy", label="Select Metric")
+                    metric_dropdown = gr.Dropdown(choices=METRICS_FOR_PLOTS, value="Accuracy", label="Select Metric")
         
                 performance_plot = gr.Plot()
 
@@ -96,11 +96,10 @@ with demo:
                 gr.Markdown("Visualize model performance across tasks according to a specific metric")
 
                 model_choices = sorted(LEADERBOARD_DF["Model"].dropna().unique())
-                metric_choices = ["Accuracy", "MCC"]
 
                 with gr.Row():
                     model_dropdown = gr.Dropdown(choices=model_choices, label="Select Model")
-                    metric2_dropdown = gr.Dropdown(choices=metric_choices, value="Accuracy", label="Select Metric")
+                    metric2_dropdown = gr.Dropdown(choices=METRICS_FOR_PLOTS, value="Accuracy", label="Select Metric")
 
                 model_plot = gr.Plot()
 

src/display/plotting.py

@@ -1,6 +1,8 @@
 import plotly.express as px
 import pandas as pd
 
+METRICS_FOR_PLOTS = ["Accuracy", "MCC"]  # TO DO: Add F1 later (to leaderboard as a column) and also here to make it available for plots
+
 def extract_mean_std(col):
     if isinstance(col, str) and '±' in col:
         try:
@@ -13,108 +15,92 @@ def extract_mean_std(col):
     return None, None
 
 def prepare_leaderboard_df(df):
-    for metric in ["Accuracy", "MCC"]:
+    for metric in METRICS_FOR_PLOTS:
         means, stds = zip(*df[metric].apply(extract_mean_std))
         df[f"{metric}_mean"] = means
         df[f"{metric}_std"] = stds
     return df
 
-def make_wrapped_task_plot(leaderboard_df):
-    def wrapped_task_plot_fn(task, metric):
-        return plot_task_metric(leaderboard_df, task, metric)
-    return wrapped_task_plot_fn
-
-def make_wrapped_model_plot(leaderboard_df):
-    def wrapped_model_plot_fn(model, metric):
-        return plot_model_across_tasks(leaderboard_df, model, metric)
-    return wrapped_model_plot_fn
-
-# Plot 1:  Model Performance per Task
-def plot_task_metric(df, task, metric):
+def make_plot_wrapper(leaderboard_df, group_by, filter_col, orientation):
+    def plot_fn(filter_val, metric):
+        return plot_metric_bar(
+            df=leaderboard_df,
+            group_by=group_by,
+            filter_col=filter_col,
+            filter_val=filter_val,
+            metric=metric,
+            orientation=orientation
+        )
+    return plot_fn
+
+def plot_metric_bar(df, group_by, filter_col, filter_val, metric, orientation="v"):
     df = df.copy()
-    df = df[df["Task"] == task]
+    df = df[df[filter_col] == filter_val]
 
     y_col = f"{metric}_mean"
     std_col = f"{metric}_std"
 
     if df.empty or y_col not in df.columns:
-        return px.bar(title="No models found for selected task.")
+        return px.bar(title="No data found.")
 
     df = df.sort_values(by=y_col, ascending=False)
-    df["Model_wrapped"] = df["Model"]
-
-    fig = px.bar(
-    df,
-    y="Model_wrapped",
-    x=y_col,      
-    orientation='h', 
-    title=f"{metric} per Model on {task}",
-    labels={y_col: metric, "Model_wrapped": "Model"},
-    text_auto=".2f",
-    height=20 * len(df) + 200
-)   
-    fig.update_traces(
-        marker_color="#F97316",
-        hovertemplate=f"<b>%{{y}}</b><br>{metric}: %{{x:.2f}}<br>Std Dev: %{{customdata[0]:.3f}}",
-        customdata=df[[std_col]].values
-    )
-
-    fig.update_layout(
-    title={
-        'text': f"{metric} per Model on {task} tasks",
-        'x': 0.5,
-        'xanchor': 'center',
-        'yanchor': 'top',
-        'font': dict(size=20),
-        'pad': dict(t=0, b=0),
-    },
-    margin=dict(t=10, b=10, l=150, r=10),
-    yaxis_tickfont_size=10,
-    xaxis=dict(range=[0, 1])
-    )
-
-    return fig
-
-# Plot 2: Overview of the performance of a model
-import plotly.express as px
-
-def plot_model_across_tasks(df, model_name, metric):
-    df = df.copy()
-    y_col = f"{metric}_mean"
-    std_col = f"{metric}_std"
-
-    model_df = df[df["Model"] == model_name]
-    if model_df.empty or y_col not in model_df.columns:
-        return px.bar(title="No data found for selected model.")
 
-    model_df = model_df.sort_values(by=y_col, ascending=False)
+    x_axis, y_axis = (y_col, group_by) if orientation == "h" else (group_by, y_col)
+    height = 20 * len(df) + 200 if orientation == "h" else None
 
     fig = px.bar(
-        model_df,
-        x="Task",
-        y=y_col,
-        title=f"{metric} for {model_name} across Tasks",
-        labels={y_col: metric, "Task": "Task"},
+        df,
+        x=x_axis,
+        y=y_axis,
+        orientation=orientation,
         text_auto=".2f",
+        labels={y_col: metric, group_by: group_by},
+        title=f"{metric} for {filter_val}",
+        height=height,
     )
 
+    if orientation == "h":
+        hovertemplate = (
+            "<b>%{y}</b><br>"
+            f"{metric}: %{{x:.2f}}<br>"
+            "Std Dev: %{customdata[0]:.3f}"
+        )
+    else:
+        hovertemplate = (
+            "<b>%{x}</b><br>"
+            f"{metric}: %{{y:.2f}}<br>"
+            "Std Dev: %{customdata[0]:.3f}"
+        )
+
+
     fig.update_traces(
         marker_color="#F97316",
-        hovertemplate=f"<b>%{{x}}</b><br>{metric}: %{{y:.2f}}<br>Std Dev: %{{customdata[0]:.3f}}",
-        customdata=model_df[[std_col]].values
+        hovertemplate=hovertemplate,
+        customdata=df[[std_col]].values
     )
 
-    fig.update_layout(
+    layout_args = dict(
         title={
-        'text': f"{metric} for {model_name} across Tasks",
-        'x': 0.5,
-        'xanchor': 'center',
-        'yanchor': 'top',
-        'font': dict(size=20),
-        'pad': dict(t=0, b=0),
-    },
-        yaxis=dict(range=[0, 1]),
-        margin=dict(t=50, b=100)
+            'text': f"{metric} for {filter_val}",
+            'x': 0.5,
+            'xanchor': 'center',
+            'yanchor': 'top',
+            'font': dict(size=20),
+            'pad': dict(t=0, b=0),
+        },
+        margin=dict(t=10, b=100, l=150, r=10),
+        hoverlabel=dict(
+            font=dict(color="white"),
+            bgcolor="#F97316",
+            bordercolor="#c5580d"
+        )
     )
 
+    if orientation == "h":
+        layout_args["xaxis"] = dict(range=[0, 1])
+        layout_args["yaxis_tickfont_size"] = 10
+    else:
+        layout_args["yaxis"] = dict(range=[0, 1])
+
+    fig.update_layout(**layout_args)
     return fig