Update the size vs. score graph

app.py

@@ -295,33 +295,45 @@ def toggle_all_categories(action: str) -> list[gr.CheckboxGroup]:
     return results
 
 
+TASK_AVG_NAME_MAP = {
+    c.name: c.task_type.name for c in fields(AutoEvalColumn) if c.average and c.task_type != TaskType.AVG
+}
+
+
 def plot_size_vs_score(df_filtered: pd.DataFrame) -> go.Figure:
     df = ORIGINAL_DF[ORIGINAL_DF[AutoEvalColumn.row_id.name].isin(df_filtered[AutoEvalColumn.row_id.name])]
     df = df[df["#Params (B)"] > 0]
-    df = df[["model_name_for_query", "#Params (B)", "AVG", "Few-shot"]]
-    df["AVG"] = df["AVG"].astype(float)
+    AVG_COLUMNS = ["AVG"] + list(TASK_AVG_NAME_MAP.keys())
+    df = df[["model_name_for_query", "#Params (B)", "Few-shot"] + AVG_COLUMNS]
+    df[AVG_COLUMNS] = df[AVG_COLUMNS].astype(float)
     df = df.rename(columns={"model_name_for_query": "Model", "Few-shot": "n-shot"})
     df["model_name_without_org_name"] = df["Model"].str.split("/").str[-1] + " (" + df["n-shot"] + "-shot)"
+    df = pd.melt(
+        df,
+        id_vars=["Model", "model_name_without_org_name", "#Params (B)", "n-shot"],
+        value_vars=AVG_COLUMNS,
+        var_name="Task_Category",
+        value_name="Score",
+    )
     fig = px.scatter(
         df,
         x="#Params (B)",
-        y="AVG",
+        y="Score",
         text="model_name_without_org_name",
-        hover_data=["Model", "n-shot"],
+        color="Task_Category",
+        hover_data=["Model", "n-shot", "Task_Category"],
     )
     fig.update_traces(
-        hovertemplate="<b>%{customdata[0]}</b><br>#Params: %{x:.2f}B<br>n-shot: %{customdata[1]}<br>AVG: %{y:.4f}<extra></extra>",
+        hovertemplate="<b>%{customdata[0]}</b><br>#Params: %{x:.2f}B<br>n-shot: %{customdata[1]}<br>%{customdata[2]}: %{y:.4f}<extra></extra>",
         textposition="top right",
     )
+    for trace in fig.data:
+        if trace.name != "AVG":
+            trace.visible = "legendonly"
     fig.update_layout(xaxis_range=[0, MAX_MODEL_SIZE * 1.2], yaxis_range=[0, 1])
     return fig
 
 
-TASK_AVG_NAME_MAP = {
-    c.name: c.task_type.name for c in fields(AutoEvalColumn) if c.average and c.task_type != TaskType.AVG
-}
-
-
 def plot_average_scores(df_filtered: pd.DataFrame) -> go.Figure:
     df = ORIGINAL_DF[ORIGINAL_DF[AutoEvalColumn.row_id.name].isin(df_filtered[AutoEvalColumn.row_id.name])]
     df = df[["model_name_for_query", "Few-shot"] + list(TASK_AVG_NAME_MAP.keys())]