Refactor split create_tables into static Benchmark methods

mteb/benchmarks/benchmark.py

@@ -72,3 +72,19 @@ def load_results(
         results = base_results.select_tasks(self.tasks)
         self.results_cache[base_results] = results
         return results
+
+    @staticmethod
+    def create_summary_table(scores_long: list[dict], search_query: str | None = None):
+        """create_summary_table"""
+        # Avoid circular references
+        from mteb.leaderboard.table import create_summary_table
+
+        return create_summary_table(scores_long, search_query)
+
+    @staticmethod
+    def create_per_task_table(scores_long: list[dict], search_query: str | None = None):
+        """create_per_task_table"""
+        # Avoid circular references
+        from mteb.leaderboard.table import create_per_task_table
+
+        return create_per_task_table(scores_long, search_query)

mteb/leaderboard/app.py

@@ -23,7 +23,6 @@
     make_selector,
 )
 from mteb.leaderboard.figures import performance_size_plot, radar_chart
-from mteb.leaderboard.table import create_tables
 from mteb.leaderboard.text_segments import ACKNOWLEDGEMENT, FAQ
 
 logger = logging.getLogger(__name__)
@@ -218,10 +217,12 @@ def get_leaderboard_app() -> gr.Blocks:
         max_model_size=MAX_MODEL_SIZE,
         zero_shot_setting="allow_all",
     )
+    default_filtered_scores = [
+        entry for entry in default_scores if entry["model_name"] in filtered_models
+    ]
+    summary_table = default_benchmark.create_summary_table(default_filtered_scores)
+    per_task_table = default_benchmark.create_per_task_table(default_filtered_scores)
 
-    summary_table, per_task_table = create_tables(
-        [entry for entry in default_scores if entry["model_name"] in filtered_models]
-    )
     lang_select = gr.Dropdown(
         LANGUAGE,
         value=sorted(default_results.languages),
@@ -763,7 +764,8 @@ def update_tables(
                     filtered_scores.append(entry)
             else:
                 filtered_scores = scores
-            summary, per_task = create_tables(filtered_scores)
+            summary = default_benchmark.create_summary_table(filtered_scores)
+            per_task = default_benchmark.create_per_task_table(filtered_scores)
             elapsed = time.time() - start_time
             logger.debug(f"update_tables callback: {elapsed}s")
             return summary, per_task

mteb/leaderboard/table.py

@@ -344,3 +344,239 @@ def create_tables(
         joint_table, per_task, score_columns, column_types
     )
     return summary_table, per_task_table
+
+
+def _prepare_data(
+    scores_long: list[dict], search_query: str | None = None
+) -> tuple[pd.DataFrame, list[str]] | None:
+    """Prepare raw dataframe and filter invalid models.
+
+    Returns:
+        per_task dataframe, models_to_remove
+    """
+    if not scores_long:
+        return None
+
+    data = pd.DataFrame.from_records(scores_long)
+    per_task = data.pivot(index="model_name", columns="task_name", values="score")
+
+    to_remove = per_task.isna().all(axis="columns")
+    if search_query:
+        names = per_task.index.get_level_values("model_name")
+        names = pd.Series(names, index=per_task.index)
+        to_remove |= ~names.str.contains(search_query, regex=True)
+
+    if to_remove.all():
+        return None
+
+    models_to_remove = list(per_task[to_remove].index)
+    per_task = per_task.drop(models_to_remove, axis=0)
+
+    return per_task, models_to_remove
+
+
+def apply_summary_styling(
+    joint_table: pd.DataFrame, score_columns: list[str], column_types: list[str]
+) -> gr.DataFrame:
+    """Apply styling for summary (joint) table."""
+    excluded_columns = [
+        "Rank (Borda)",
+        "Model",
+        "Number of Parameters",
+        "Embedding Dimensions",
+        "Max Tokens",
+        "Memory Usage (MB)",
+    ]
+    gradient_columns = [
+        col for col in joint_table.columns if col not in excluded_columns
+    ]
+    light_green_cmap = create_light_green_cmap()
+
+    numeric_data = joint_table.copy()
+    joint_table["Zero-shot"] = joint_table["Zero-shot"].apply(format_zero_shot)
+    joint_table[score_columns] = joint_table[score_columns].map(format_scores)
+
+    joint_table_style = joint_table.style.format(
+        {**dict.fromkeys(score_columns, "{:.2f}"), "Rank (Borda)": "{:.0f}"},
+        na_rep="",
+    )
+    joint_table_style = joint_table_style.highlight_min(
+        "Rank (Borda)", props="font-weight: bold"
+    ).highlight_max(subset=score_columns, props="font-weight: bold")
+
+    # background gradient for each column
+    for col in gradient_columns:
+        if col in joint_table.columns:
+            mask = numeric_data[col].notna()
+            if col != "Zero-shot":
+                gmap_values = numeric_data[col] * 100
+                cmap = light_green_cmap
+                joint_table_style = joint_table_style.background_gradient(
+                    cmap=cmap,
+                    subset=pd.IndexSlice[mask, col],
+                    gmap=gmap_values.loc[mask],
+                )
+            else:
+                gmap_values = numeric_data[col]
+                cmap = "RdYlGn"
+                joint_table_style = joint_table_style.background_gradient(
+                    cmap=cmap,
+                    subset=pd.IndexSlice[mask, col],
+                    vmin=50,
+                    vmax=100,
+                    gmap=gmap_values.loc[mask],
+                )
+
+    column_widths = get_column_widths(joint_table_style.data)
+    column_widths[0] = "100px"
+    column_widths[1] = "250px"
+
+    return gr.DataFrame(
+        joint_table_style,
+        datatype=column_types,
+        interactive=False,
+        pinned_columns=3,
+        column_widths=column_widths,
+        wrap=True,
+        show_fullscreen_button=True,
+        show_copy_button=True,
+        show_search="filter",
+    )
+
+
+def apply_per_task_styling(per_task: pd.DataFrame) -> gr.DataFrame:
+    """Apply styling for per-task table."""
+    task_score_columns = per_task.select_dtypes("number").columns
+    per_task[task_score_columns] *= 100
+
+    per_task_style = per_task.style.format(
+        "{:.2f}", subset=task_score_columns, na_rep=""
+    ).highlight_max(subset=task_score_columns, props="font-weight: bold")
+
+    return gr.DataFrame(
+        per_task_style,
+        interactive=False,
+        pinned_columns=1,
+        show_fullscreen_button=True,
+        show_copy_button=True,
+        show_search="filter",
+    )
+
+
+def create_summary_table(
+    scores_long: list[dict], search_query: str | None = None
+) -> gr.DataFrame:
+    """create_summary_table"""
+    prepared = _prepare_data(scores_long, search_query)
+    if prepared is None:
+        no_results_frame = pd.DataFrame(
+            {"No results": ["You can try relaxing your criteria"]}
+        )
+        return gr.DataFrame(no_results_frame)
+
+    per_task, models_to_remove = prepared
+    data = pd.DataFrame.from_records(scores_long)
+
+    mean_per_type = get_means_per_types(per_task)
+    mean_per_type = mean_per_type.pivot(
+        index="model_name", columns="task_type", values="score"
+    )
+    mean_per_type.columns = [
+        split_on_capital(column) for column in mean_per_type.columns
+    ]
+    typed_mean = mean_per_type.mean(skipna=False, axis=1)
+    overall_mean = per_task.mean(skipna=False, axis=1)
+    joint_table = mean_per_type.copy()
+    joint_table = joint_table.drop(models_to_remove, axis=0)
+    joint_table.insert(0, "mean", overall_mean)
+    joint_table.insert(1, "mean_by_task_type", typed_mean)
+    joint_table["borda_rank"] = get_borda_rank(per_task)
+    joint_table = joint_table.sort_values("borda_rank", ascending=True)
+    joint_table = joint_table.reset_index()
+    model_metas = joint_table["model_name"].map(failsafe_get_model_meta)
+    joint_table = joint_table[model_metas.notna()]
+    joint_table["model_link"] = model_metas.map(lambda m: m.reference)
+    joint_table.insert(
+        1,
+        "Max Tokens",
+        model_metas.map(lambda m: format_max_tokens(m.max_tokens)),
+    )
+    joint_table.insert(
+        1,
+        "Embedding Dimensions",
+        model_metas.map(lambda m: str(int(m.embed_dim)) if m.embed_dim else "Unknown"),
+    )
+    joint_table.insert(
+        1,
+        "Number of Parameters",
+        model_metas.map(lambda m: format_n_parameters(m.n_parameters)),
+    )
+    joint_table.insert(
+        1,
+        "Memory Usage (MB)",
+        model_metas.map(
+            lambda m: str(int(m.memory_usage_mb)) if m.memory_usage_mb else "Unknown"
+        ),
+    )
+    tasks = get_tasks(tasks=list(data["task_name"].unique()))
+    joint_table.insert(
+        1, "Zero-shot", model_metas.map(lambda m: m.zero_shot_percentage(tasks))
+    )
+    joint_table["Zero-shot"] = joint_table["Zero-shot"].fillna(-1)
+    # joint_table = joint_table[joint_table["Zero-shot"].notna()]
+    # Removing HF organization from model
+    joint_table["model_name"] = joint_table["model_name"].map(
+        lambda name: name.split("/")[-1]
+    )
+    # Adding markdown link to model names
+    name_w_link = (
+        "[" + joint_table["model_name"] + "](" + joint_table["model_link"] + ")"
+    )
+    joint_table["model_name"] = joint_table["model_name"].mask(
+        joint_table["model_link"].notna(), name_w_link
+    )
+    joint_table = joint_table.drop(columns=["model_link"])
+    joint_table = joint_table.rename(
+        columns={
+            "model_name": "Model",
+            "mean_by_task_type": "Mean (TaskType)",
+            "mean": "Mean (Task)",
+        }
+    )
+
+    joint_table.insert(0, "Rank (Borda)", joint_table.pop("borda_rank"))
+    column_types = get_column_types(joint_table)
+    # setting model name column to markdown
+    column_types[1] = "markdown"
+    score_columns = ["Mean (Task)", "Mean (TaskType)", *mean_per_type.columns]
+
+    return apply_summary_styling(joint_table, score_columns, column_types)
+
+
+def create_per_task_table(
+    scores_long: list[dict], search_query: str | None = None
+) -> gr.DataFrame:
+    """create_per_task_table"""
+    prepared = _prepare_data(scores_long, search_query)
+    if prepared is None:
+        no_results_frame = pd.DataFrame(
+            {"No results": ["You can try relaxing your criteria"]}
+        )
+        return gr.DataFrame(no_results_frame)
+
+    per_task, _ = prepared
+
+    per_task["borda_rank"] = get_borda_rank(per_task)
+    per_task = per_task.sort_values("borda_rank", ascending=True)
+    per_task = per_task.drop(columns=["borda_rank"])
+    per_task = per_task.reset_index()
+    per_task["model_name"] = per_task["model_name"].map(
+        lambda name: name.split("/")[-1]
+    )
+    per_task = per_task.rename(
+        columns={
+            "model_name": "Model",
+        }
+    )
+
+    return apply_per_task_styling(per_task)