mteb 2.1.4__py3-none-any.whl → 2.7.2__py3-none-any.whl

mteb/abstasks/zeroshot_classification.py

@@ -1,13 +1,14 @@
 import logging
 from pathlib import Path
-from typing import Any, TypedDict
+from typing import TypedDict
 
 import torch
 from datasets import Dataset
 from sklearn import metrics
 
 from mteb._evaluators import ZeroShotClassificationEvaluator
-from mteb.models import EncoderProtocol
+from mteb.models import EncoderProtocol, MTEBModels
+from mteb.types import EncodeKwargs
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -111,15 +112,18 @@ class AbsTaskZeroShotClassification(AbsTask):
 
     def _evaluate_subset(
         self,
-        model: EncoderProtocol,
+        model: MTEBModels,
         data_split: Dataset,
         *,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         **kwargs,
     ) -> ZeroShotClassificationMetrics:
+        if not isinstance(model, EncoderProtocol):
+            raise TypeError("Expected model to be an instance of EncoderProtocol")
+
         candidate_labels = self.get_candidate_labels()
         data_split = data_split.select_columns(
             [self.input_column_name, self.label_column_name]

mteb/benchmarks/_create_table.py

@@ -1,6 +1,6 @@
-import math
 import re
 from collections import defaultdict
+from typing import Literal
 
 import numpy as np
 import pandas as pd
@@ -32,26 +32,18 @@ def _split_on_capital(s: str) -> str:
     return " ".join(re.findall(r"[A-Z]?[a-z]+|[A-Z]+(?=[A-Z]|$)", s))
 
 
-def _format_n_parameters(n_parameters) -> str:
-    if (n_parameters is None) or (not int(n_parameters)):
-        return "Unknown"
-    n_thousand = int(n_parameters // 1e3)
-    if n_thousand < 1:
-        return str(int(n_parameters))
-    n_zeros = math.log10(n_thousand)
-    if n_zeros >= 6:
-        return str(n_thousand // (10**6)) + "B"
-    if n_zeros >= 3:
-        return str(n_thousand // (10**3)) + "M"
-    return str(n_thousand) + "K"
+def _format_n_parameters(n_parameters) -> float | None:
+    """Format n_parameters to be in billions with decimals down to 1 million. I.e. 7M -> 0.007B, 1.5B -> 1.5B, None -> None"""
+    if n_parameters:
+        n_parameters = float(n_parameters)
+        return round(n_parameters / 1e9, 3)
+    return None
 
 
-def _format_max_tokens(max_tokens: float | None) -> str:
-    if max_tokens is None:
-        return "Unknown"
-    if max_tokens == np.inf:
-        return "Infinite"
-    return str(int(max_tokens))
+def _format_max_tokens(max_tokens: float | None) -> float | None:
+    if max_tokens is None or max_tokens == np.inf:
+        return None
+    return float(max_tokens)
 
 
 def _get_means_per_types(per_task: pd.DataFrame):
@@ -144,18 +136,18 @@ def _create_summary_table_from_benchmark_results(
     joint_table.insert(
         1,
         "Embedding Dimensions",
-        model_metas.map(lambda m: str(int(m.embed_dim)) if m.embed_dim else "Unknown"),
+        model_metas.map(lambda m: int(m.embed_dim) if m.embed_dim else None),
     )
     joint_table.insert(
         1,
-        "Number of Parameters",
+        "Number of Parameters (B)",
         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"
+            lambda m: int(m.memory_usage_mb) if m.memory_usage_mb else None
         ),
     )
 
@@ -250,6 +242,65 @@ def _create_per_task_table_from_benchmark_results(
     return per_task
 
 
+def _create_per_language_table_from_benchmark_results(
+    benchmark_results: BenchmarkResults,
+    language_view: list[str] | Literal["all"],
+) -> pd.DataFrame:
+    """Create per-language table from BenchmarkResults.
+
+    Returns a DataFrame with one row per model and one column per language.
+
+    Args:
+        benchmark_results: BenchmarkResults object containing model results
+        language_view: List of languages to include in the per-language table, or "all" for all languages present in the results
+    Returns:
+        DataFrame with per-language scores, ready for styling in the leaderboard
+    """
+    if language_view != "all" and not isinstance(language_view, list):
+        raise ValueError("language_view must be a list of languages or 'all'")
+
+    data = benchmark_results.to_dataframe(aggregation_level="language", format="long")
+
+    if data.empty:
+        no_results_frame = pd.DataFrame(
+            {"No results": ["You can try relaxing your criteria"]}
+        )
+        return no_results_frame
+
+    if language_view != "all":
+        data = data[data["language"].isin(language_view)]
+
+    per_language = data.pivot_table(
+        index="model_name", columns="language", values="score", aggfunc="mean"
+    )
+
+    to_remove = per_language.isna().all(axis="columns")
+    if to_remove.all():
+        no_results_frame = pd.DataFrame(
+            {"No results": ["You can try relaxing your criteria"]}
+        )
+        return no_results_frame
+
+    models_to_remove = list(per_language[to_remove].index)
+    per_language = per_language.drop(models_to_remove, axis=0)
+
+    per_language["borda_rank"] = _get_borda_rank(per_language)
+    per_language = per_language.sort_values("borda_rank", ascending=True)
+    per_language = per_language.drop(columns=["borda_rank"])
+    per_language = per_language.reset_index()
+
+    per_language["model_name"] = per_language["model_name"].map(
+        lambda name: name.split("/")[-1]
+    )
+    per_language = per_language.rename(
+        columns={
+            "model_name": "Model",
+        }
+    )
+
+    return per_language
+
+
 def _create_summary_table_mean_public_private(
     benchmark_results: BenchmarkResults,
 ) -> pd.DataFrame:
@@ -323,18 +374,18 @@ def _create_summary_table_mean_public_private(
     joint_table.insert(
         1,
         "Embedding Dimensions",
-        model_metas.map(lambda m: str(int(m.embed_dim)) if m.embed_dim else "Unknown"),
+        model_metas.map(lambda m: int(m.embed_dim) if m.embed_dim else None),
     )
     joint_table.insert(
         1,
-        "Number of Parameters",
+        "Number of Parameters (B)",
         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"
+            lambda m: int(m.memory_usage_mb) if m.memory_usage_mb else None
         ),
     )
 
@@ -358,9 +409,7 @@ def _create_summary_table_mean_public_private(
         "mean(public)": "Mean (Public)",
         "mean(private)": "Mean (Private)",
     }
-    # For RTEB: all tasks are Retrieval type, so Retrieval column = Mean (Task)
-    if "Retrieval" in joint_table.columns:
-        rename_dict["Retrieval"] = "Mean (Task)"
+
     joint_table = joint_table.rename(columns=rename_dict)
 
     # Move borda rank to front
@@ -447,18 +496,18 @@ def _create_summary_table_mean_subset(
     joint_table.insert(
         1,
         "Embedding Dimensions",
-        model_metas.map(lambda m: str(int(m.embed_dim)) if m.embed_dim else "Unknown"),
+        model_metas.map(lambda m: int(m.embed_dim) if m.embed_dim else None),
     )
     joint_table.insert(
         1,
-        "Number of Parameters",
+        "Number of Parameters (B)",
         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"
+            lambda m: int(m.memory_usage_mb) if m.memory_usage_mb else None
         ),
     )
 
@@ -560,25 +609,23 @@ def _create_summary_table_mean_task_type(
 
     # Insert model metadata columns
     joint_table.insert(
-        1,
-        "Max Tokens",
-        model_metas.map(lambda m: _format_max_tokens(m.max_tokens)),
+        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"),
+        model_metas.map(lambda m: int(m.embed_dim) if m.embed_dim else None),
     )
     joint_table.insert(
         1,
-        "Number of Parameters",
+        "Number of Parameters (B)",
         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"
+            lambda m: int(m.memory_usage_mb) if m.memory_usage_mb else None
         ),
     )
 

mteb/benchmarks/benchmark.py

@@ -1,21 +1,16 @@
-from collections.abc import Iterable, Sequence
-from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from __future__ import annotations
+
+from collections.abc import Iterator, Sequence
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Literal
 
 import pandas as pd
 
-from mteb.benchmarks._create_table import (
-    _create_per_task_table_from_benchmark_results,
-    _create_summary_table_from_benchmark_results,
-    _create_summary_table_mean_public_private,
-    _create_summary_table_mean_subset,
-    _create_summary_table_mean_task_type,
-)
-from mteb.results import BenchmarkResults
+from mteb.abstasks.abstask import AbsTask
 from mteb.types import StrURL
 
 if TYPE_CHECKING:
-    from mteb.abstasks import AbsTask
+    from mteb.results import BenchmarkResults
 
 
 @dataclass
@@ -24,6 +19,7 @@ class Benchmark:
 
     Args:
         name: The name of the benchmark
+        aliases: Alternative names for the benchmark
         tasks: The tasks within the benchmark.
         description: A description of the benchmark, should include its intended goal and potentially a description of its construction
         reference: A link reference, to a source containing additional information typically to a paper, leaderboard or github.
@@ -42,7 +38,8 @@ class Benchmark:
     """
 
     name: str
-    tasks: Sequence["AbsTask"]
+    tasks: Sequence[AbsTask]
+    aliases: Sequence[str] = field(default_factory=tuple)
     description: str | None = None
     reference: StrURL | None = None
     citation: str | None = None
@@ -50,14 +47,15 @@ class Benchmark:
     display_on_leaderboard: bool = True
     icon: str | None = None
     display_name: str | None = None
+    language_view: list[str] | Literal["all"] = field(default_factory=list)
 
-    def __iter__(self) -> Iterable["AbsTask"]:
+    def __iter__(self) -> Iterator[AbsTask]:
         return iter(self.tasks)
 
     def __len__(self) -> int:
         return len(self.tasks)
 
-    def __getitem__(self, index: int) -> "AbsTask":
+    def __getitem__(self, index: int) -> AbsTask:
         return self.tasks[index]
 
     def _create_summary_table(
@@ -68,6 +66,10 @@ class Benchmark:
         Returns:
             A pandas DataFrame representing the summary results.
         """
+        from mteb.benchmarks._create_table import (
+            _create_summary_table_from_benchmark_results,
+        )
+
         return _create_summary_table_from_benchmark_results(benchmark_results)
 
     def _create_per_task_table(
@@ -78,8 +80,38 @@ class Benchmark:
         Returns:
             A pandas DataFrame representing the per-task results.
         """
+        from mteb.benchmarks._create_table import (
+            _create_per_task_table_from_benchmark_results,
+        )
+
         return _create_per_task_table_from_benchmark_results(benchmark_results)
 
+    def _create_per_language_table(
+        self, benchmark_results: BenchmarkResults
+    ) -> pd.DataFrame:
+        """Create per-language table. Called by the leaderboard app.
+
+        Returns:
+            A pandas DataFrame representing the per-language results.
+        """
+        from mteb.benchmarks._create_table import (
+            _create_per_language_table_from_benchmark_results,
+        )
+
+        if self.language_view == "all" or len(self.language_view) > 0:
+            return _create_per_language_table_from_benchmark_results(
+                benchmark_results, self.language_view
+            )
+        else:
+            no_results_frame = pd.DataFrame(
+                {
+                    "No results": [
+                        "The per-language table is not available for this benchmark."
+                    ]
+                }
+            )
+            return no_results_frame
+
 
 class RtebBenchmark(Benchmark):
     """Wrapper for RTEB benchmark."""
@@ -87,7 +119,14 @@ class RtebBenchmark(Benchmark):
     def _create_summary_table(
         self, benchmark_results: BenchmarkResults
     ) -> pd.DataFrame:
-        return _create_summary_table_mean_public_private(benchmark_results)
+        from mteb.benchmarks._create_table import (
+            _create_summary_table_mean_public_private,
+        )
+
+        joint_table = _create_summary_table_mean_public_private(benchmark_results)
+        # For RTEB: all tasks are Retrieval type, so Retrieval column = Mean (Task)
+        joint_table = joint_table.rename(columns={"Retrieval": "Mean (Task)"})
+        return joint_table
 
 
 class HUMEBenchmark(Benchmark):
@@ -96,6 +135,8 @@ class HUMEBenchmark(Benchmark):
     def _create_summary_table(
         self, benchmark_results: BenchmarkResults
     ) -> pd.DataFrame:
+        from mteb.benchmarks._create_table import _create_summary_table_mean_subset
+
         return _create_summary_table_mean_subset(benchmark_results)
 
 
@@ -105,4 +146,24 @@ class MIEBBenchmark(Benchmark):
     def _create_summary_table(
         self, benchmark_results: BenchmarkResults
     ) -> pd.DataFrame:
+        from mteb.benchmarks._create_table import _create_summary_table_mean_task_type
+
         return _create_summary_table_mean_task_type(benchmark_results)
+
+
+class VidoreBenchmark(Benchmark):
+    """Wrapper for Vidore3 benchmark."""
+
+    def _create_summary_table(
+        self, benchmark_results: BenchmarkResults
+    ) -> pd.DataFrame:
+        from mteb.benchmarks._create_table import (
+            _create_summary_table_mean_public_private,
+        )
+
+        joint_table = _create_summary_table_mean_public_private(benchmark_results)
+        # For ViDoRe (V1, V2, V3): all tasks are Document Understanding type, so Document Understanding column = Mean (Task)
+        joint_table = joint_table.rename(
+            columns={"Document Understanding": "Mean (Task)"}
+        )
+        return joint_table

mteb/benchmarks/benchmarks/__init__.py

@@ -6,12 +6,16 @@ from mteb.benchmarks.benchmarks.benchmarks import (
     BUILT_MTEB,
     C_MTEB,
     CHEMTEB,
+    CHEMTEB_V1_1,
     CODE_RAG,
     ENCODECHKA,
     FA_MTEB,
     FA_MTEB_2,
     HUME,
     JINA_VDR,
+    JMTEB_LITE_V1,
+    JMTEB_V2,
+    KOVIDORE_V2,
     LONG_EMBED,
     MIEB_ENG,
     MIEB_IMG,
@@ -38,10 +42,12 @@ from mteb.benchmarks.benchmarks.benchmarks import (
     SEB,
     VIDORE,
     VIDORE_V2,
+    VIDORE_V3,
     VISUAL_DOCUMENT_RETRIEVAL,
     VN_MTEB,
     CoIR,
     MTEB_code,
+    MTEB_MAIN_RU_v1_1,
     MTEB_multilingual_v1,
     MTEB_multilingual_v2,
     RAR_b,
@@ -65,6 +71,7 @@ __all__ = [
     "BRIGHT_LONG",
     "BUILT_MTEB",
     "CHEMTEB",
+    "CHEMTEB_V1_1",
     "CODE_RAG",
     "C_MTEB",
     "ENCODECHKA",
@@ -73,6 +80,9 @@ __all__ = [
     "HUME",
     "HUME",
     "JINA_VDR",
+    "JMTEB_LITE_V1",
+    "JMTEB_V2",
+    "KOVIDORE_V2",
     "LONG_EMBED",
     "MIEB_ENG",
     "MIEB_IMG",
@@ -108,9 +118,11 @@ __all__ = [
     "SEB",
     "VIDORE",
     "VIDORE_V2",
+    "VIDORE_V3",
     "VISUAL_DOCUMENT_RETRIEVAL",
     "VN_MTEB",
     "CoIR",
+    "MTEB_MAIN_RU_v1_1",
     "MTEB_code",
     "MTEB_multilingual_v1",
     "MTEB_multilingual_v2",