Clean up

app.py

@@ -338,15 +338,17 @@ def plot_size_vs_score(df: pd.DataFrame, hidden_df: pd.DataFrame) -> go.Figure:
     return fig
 
 
+TASK_AVG_COLUMNS = [c.name for c in fields(AutoEvalColumn) if c.average and c.task_type != TaskType.AVG]
+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: pd.DataFrame, hidden_df: pd.DataFrame) -> go.Figure:
     df2 = hidden_df.iloc[df.index]
-    df2 = df2[["model_name_for_query"] + [column for column in df.columns if column.startswith("AVG ")]]
+    df2 = df2[["model_name_for_query"] + TASK_AVG_COLUMNS]
     df2 = df2.rename(columns={"model_name_for_query": "Model"})
-    df2 = df2.rename(
-        columns={
-            column: column.replace("AVG (", "").replace(")", "") for column in df2.columns if column.startswith("AVG ")
-        }
-    )
+    df2 = df2.rename(columns=TASK_AVG_NAME_MAP)
     df2 = df2.set_index("Model").astype(float)
 
     fig = go.Figure()