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

mteb/models/vllm_wrapper.py

@@ -0,0 +1,327 @@
+from __future__ import annotations
+
+import atexit
+import gc
+import logging
+import os
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Literal
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+
+from mteb._requires_package import requires_package
+from mteb.abstasks.task_metadata import TaskMetadata
+from mteb.models import ModelMeta
+from mteb.models.abs_encoder import AbsEncoder
+from mteb.types import Array, BatchedInput, PromptType
+
+if TYPE_CHECKING:
+    from vllm.config import PoolerConfig  # type: ignore[import-not-found]
+else:
+    PoolerConfig = dict[str, Any]
+
+logger = logging.getLogger(__name__)
+
+Dtype = Literal["half", "float16", "float", "float32", "bfloat16", "auto"]
+
+
+class VllmWrapperBase:
+    """Wrapper for vllm serving engine."""
+
+    convert = "auto"
+    mteb_model_meta: ModelMeta | None = None
+
+    def __init__(
+        self,
+        model: str | ModelMeta,
+        revision: str | None = None,
+        *,
+        trust_remote_code: bool = True,
+        dtype: Dtype = "auto",
+        head_dtype: Literal["model"] | Dtype | None = None,
+        max_model_len: int | None = None,
+        max_num_batched_tokens: int | None = None,
+        max_num_seqs: int = 128,
+        tensor_parallel_size: int = 1,
+        enable_prefix_caching: bool | None = None,
+        gpu_memory_utilization: float = 0.9,
+        hf_overrides: dict[str, Any] | None = None,
+        pooler_config: PoolerConfig | None = None,
+        enforce_eager: bool = False,
+        **kwargs: Any,
+    ):
+        """Wrapper for vllm serving engine.
+
+        Args:
+            model: model name string.
+            revision: The revision of the model to use.
+            trust_remote_code: Whether to trust remote code execution when loading the model.
+                Should be True for models with custom code.
+            dtype: Data type for model weights. "auto" will automatically select appropriate
+                dtype based on hardware and model capabilities. vllm uses flash attention by
+                default, which does not support fp32. Therefore, it defaults to using fp16 for
+                inference on fp32 models. Testing has shown a relatively small drop in accuracy.
+                You can manually opt for fp32, but inference speed will be very slow.
+            head_dtype: "head" refers to the last Linear layer(s) of an LLMs, such as the score
+                or classifier in a classification model. Uses fp32 for the head by default to
+                gain extra precision.
+            max_model_len: Maximum sequence length (context window) supported by the model.
+                If None, uses the model's default maximum length.
+            max_num_batched_tokens: Maximum number of tokens to process in a single batch.
+                If None, automatically determined.
+            max_num_seqs: Maximum number of sequences to process concurrently.
+            tensor_parallel_size: Number of GPUs for tensor parallelism.
+            enable_prefix_caching: Whether to enable KV cache sharing for common prompt prefixes.
+                If None, uses the model's default setting.
+            gpu_memory_utilization: Target GPU memory utilization ratio (0.0 to 1.0).
+            hf_overrides: Dictionary mapping Hugging Face configuration keys to override values.
+            pooler_config: Controls the behavior of output pooling in pooling models.
+            enforce_eager: Whether to disable CUDA graph optimization and use eager execution.
+            **kwargs: Additional arguments to pass to the vllm serving engine model.
+        """
+        requires_package(
+            self,
+            "vllm",
+            "Wrapper for vllm serving engine",
+            install_instruction="pip install mteb[vllm]",
+        )
+
+        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+        from vllm import LLM, EngineArgs
+
+        hf_overrides = {} if hf_overrides is None else hf_overrides
+
+        if head_dtype is not None:
+            hf_overrides["head_dtype"] = head_dtype
+
+        model_name = model if isinstance(model, str) else model.name
+
+        if isinstance(model, ModelMeta):
+            logger.info(
+                "Using revision from model meta. Passed revision will be ignored"
+            )
+            revision = model.revision
+
+        args = EngineArgs(
+            model=model_name,
+            revision=revision,
+            runner="pooling",
+            convert=self.convert,  # type: ignore[arg-type]
+            max_model_len=max_model_len,
+            max_num_batched_tokens=max_num_batched_tokens,
+            max_num_seqs=max_num_seqs,
+            tensor_parallel_size=tensor_parallel_size,
+            enable_prefix_caching=enable_prefix_caching,
+            gpu_memory_utilization=gpu_memory_utilization,
+            hf_overrides=hf_overrides,
+            pooler_config=pooler_config,
+            enforce_eager=enforce_eager,
+            trust_remote_code=trust_remote_code,
+            dtype=dtype,
+            **kwargs,
+        )
+        self.llm = LLM(**vars(args))
+
+        if isinstance(model, str):
+            self.mteb_model_meta = ModelMeta.from_hub(model=model, revision=revision)
+        else:
+            self.mteb_model_meta = model
+
+        atexit.register(self.cleanup)
+
+    def cleanup(self):
+        """Clean up the VLLM distributed runtime environment and release GPU resources."""
+        if self.llm is None:
+            return
+
+        from vllm.distributed import (  # type: ignore[import-not-found]
+            cleanup_dist_env_and_memory,
+        )
+
+        self.llm = None
+        gc.collect()
+        cleanup_dist_env_and_memory()
+
+    def __del__(self):
+        try:
+            self.cleanup()
+        except Exception:
+            pass
+
+
+class VllmEncoderWrapper(AbsEncoder, VllmWrapperBase):
+    """vLLM wrapper for Encoder models.
+
+    Args:
+        model: model name string or ModelMeta.
+        revision: The revision of the model to use.
+        prompt_dict: A dictionary mapping task names to prompt strings.
+        use_instructions: Whether to use instructions from the prompt_dict.
+            When False, values from prompt_dict are used as static prompts (prefixes).
+            When True, values from prompt_dict are used as instructions to be formatted
+            using the instruction_template.
+        instruction_template: A template or callable to format instructions.
+            Can be a string with '{instruction}' placeholder or a callable that takes
+            the instruction and prompt type and returns a formatted string.
+        apply_instruction_to_documents: Whether to apply instructions to documents prompts.
+        **kwargs: Additional arguments to pass to the vllm serving engine model.
+    """
+
+    convert = "embed"
+
+    def __init__(
+        self,
+        model: str | ModelMeta,
+        revision: str | None = None,
+        prompt_dict: dict[str, str] | None = None,
+        use_instructions: bool = False,
+        instruction_template: (
+            str | Callable[[str, PromptType | None], str] | None
+        ) = None,
+        apply_instruction_to_documents: bool = True,
+        **kwargs: Any,
+    ):
+        if use_instructions and instruction_template is None:
+            raise ValueError(
+                "To use instructions, an instruction_template must be provided. "
+                "For example, `Instruction: {instruction}`"
+            )
+
+        if (
+            isinstance(instruction_template, str)
+            and "{instruction}" not in instruction_template
+        ):
+            raise ValueError(
+                "Instruction template must contain the string '{instruction}'."
+            )
+
+        self.prompts_dict = prompt_dict
+        self.use_instructions = use_instructions
+        self.instruction_template = instruction_template
+        self.apply_instruction_to_passages = apply_instruction_to_documents
+        super().__init__(
+            model,
+            revision,
+            **kwargs,
+        )
+
+    def encode(
+        self,
+        inputs: DataLoader[BatchedInput],
+        *,
+        task_metadata: TaskMetadata,
+        hf_split: str,
+        hf_subset: str,
+        prompt_type: PromptType | None = None,
+        **kwargs: Any,
+    ) -> Array:
+        """Encodes the given sentences using the encoder.
+
+        Args:
+            inputs: The sentences to encode.
+            task_metadata: The metadata of the task. Sentence-transformers uses this to
+                determine which prompt to use from a specified dictionary.
+            prompt_type: The name type of prompt. (query or passage)
+            hf_split: Split of current task
+            hf_subset: Subset of current task
+            **kwargs: Additional arguments to pass to the encoder.
+
+        Returns:
+            The encoded sentences.
+        """
+        prompt = ""
+        if self.use_instructions and self.prompts_dict is not None:
+            prompt = self.get_task_instruction(task_metadata, prompt_type)
+        elif self.prompts_dict is not None:
+            prompt_name = self.get_prompt_name(task_metadata, prompt_type)
+            if prompt_name is not None:
+                prompt = self.prompts_dict.get(prompt_name, "")
+
+        if (
+            self.use_instructions
+            and self.apply_instruction_to_passages is False
+            and prompt_type == PromptType.document
+        ):
+            logger.info(
+                f"No instruction used, because prompt type = {prompt_type.document}"
+            )
+            prompt = ""
+        else:
+            logger.info(
+                f"Using instruction: '{prompt}' for task: '{task_metadata.name}' prompt type: '{prompt_type}'"
+            )
+
+        prompts = [prompt + text for batch in inputs for text in batch["text"]]
+        outputs = self.llm.encode(
+            prompts, pooling_task="embed", truncate_prompt_tokens=-1
+        )
+        embeddings = torch.stack([output.outputs.data for output in outputs])
+        return embeddings
+
+
+class VllmCrossEncoderWrapper(VllmWrapperBase):
+    """vLLM wrapper for CrossEncoder models."""
+
+    convert = "classify"
+
+    def __init__(
+        self,
+        model: str | ModelMeta,
+        revision: str | None = None,
+        query_prefix: str = "",
+        document_prefix: str = "",
+        **kwargs: Any,
+    ):
+        super().__init__(
+            model,
+            revision,
+            **kwargs,
+        )
+        self.query_prefix = query_prefix
+        self.document_prefix = document_prefix
+
+    def predict(
+        self,
+        inputs1: DataLoader[BatchedInput],
+        inputs2: DataLoader[BatchedInput],
+        *,
+        task_metadata: TaskMetadata,
+        hf_split: str,
+        hf_subset: str,
+        prompt_type: PromptType | None = None,
+        **kwargs: Any,
+    ) -> Array:
+        """Predicts relevance scores for pairs of inputs. Note that, unlike the encoder, the cross-encoder can compare across inputs.
+
+        Args:
+            inputs1: First Dataloader of inputs to encode. For reranking tasks, these are queries (for text only tasks `QueryDatasetType`).
+            inputs2: Second Dataloader of inputs to encode. For reranking, these are documents (for text only tasks `RetrievalOutputType`).
+            task_metadata: Metadata of the current task.
+            hf_split: Split of current task, allows to know some additional information about current split.
+                E.g. Current language
+            hf_subset: Subset of current task. Similar to `hf_split` to get more information
+            prompt_type: The name type of prompt. (query or passage)
+            **kwargs: Additional arguments to pass to the cross-encoder.
+
+        Returns:
+            The predicted relevance scores for each inputs pair.
+        """
+        queries = [
+            self.query_prefix + text for batch in inputs1 for text in batch["text"]
+        ]
+        corpus = [
+            self.document_prefix + text for batch in inputs2 for text in batch["text"]
+        ]
+        # TODO: support score prompt
+
+        outputs = self.llm.score(
+            queries,
+            corpus,
+            truncate_prompt_tokens=-1,
+            use_tqdm=False,
+        )
+        scores = np.array([output.outputs.score for output in outputs])
+        return scores

mteb/results/benchmark_results.py

@@ -1,9 +1,12 @@
+from __future__ import annotations
+
+import functools
 import json
 import logging
 import warnings
-from collections.abc import Callable, Iterable, Iterator, Sequence
+from collections.abc import Callable, Iterable, Iterator
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 import pandas as pd
 from packaging.version import InvalidVersion, Version
@@ -15,6 +18,7 @@ from mteb.abstasks.task_metadata import (
     TaskDomain,
     TaskType,
 )
+from mteb.benchmarks.benchmark import Benchmark
 from mteb.models import ModelMeta
 from mteb.models.get_model_meta import get_model_metas
 from mteb.types import (
@@ -31,6 +35,25 @@ from .model_result import ModelResult, _aggregate_and_pivot
 logger = logging.getLogger(__name__)
 
 
+@functools.lru_cache
+def _get_cached_model_metas() -> dict[str, str | None]:
+    """Cache model metas to avoid repeated calls."""
+    return {
+        meta.name: meta.revision for meta in get_model_metas() if meta.name is not None
+    }
+
+
+@functools.lru_cache(maxsize=10000)
+def _parse_version_cached(version_str: str | None) -> Version | None:
+    """Cache version parsing to avoid repeated parsing."""
+    if version_str is None:
+        return None
+    try:
+        return Version(version_str)
+    except (InvalidVersion, TypeError):
+        return None
+
+
 class BenchmarkResults(BaseModel):
     """Data class to hold the benchmark results of a model.
 
@@ -39,10 +62,10 @@ class BenchmarkResults(BaseModel):
     """
 
     model_results: list[ModelResult]
-    model_config = (
-        ConfigDict(  # to free up the name model_results which is otherwise protected
-            protected_namespaces=(),
-        )
+    benchmark: Benchmark | None = None
+    model_config = ConfigDict(
+        protected_namespaces=(),  # to free up the name model_results which is otherwise protected
+        arbitrary_types_allowed=True,  # Benchmark is dataclasses.dataclass
     )
 
     def __repr__(self) -> str:
@@ -57,10 +80,10 @@ class BenchmarkResults(BaseModel):
         task_names: list[str] | None = None,
         languages: list[str] | None = None,
         domains: list[TaskDomain] | None = None,
-        task_types: list[TaskType] | None = None,  # type: ignore
+        task_types: list[TaskType] | None = None,
         modalities: list[Modalities] | None = None,
         is_public: bool | None = None,
-    ) -> Self:
+    ) -> BenchmarkResults:
         # TODO: Same as filter_models
         model_results = [
             res._filter_tasks(
@@ -77,7 +100,7 @@ class BenchmarkResults(BaseModel):
             model_results=[res for res in model_results if res.task_results]
         )
 
-    def select_tasks(self, tasks: Sequence[AbsTask]) -> Self:
+    def select_tasks(self, tasks: Iterable[AbsTask]) -> BenchmarkResults:
         """Select tasks from the benchmark results.
 
         Args:
@@ -95,7 +118,7 @@ class BenchmarkResults(BaseModel):
         self,
         names: list[str] | list[ModelMeta],
         revisions: list[str | None] | None = None,
-    ) -> Self:
+    ) -> BenchmarkResults:
         """Get models by name and revision.
 
         Args:
@@ -108,7 +131,7 @@ class BenchmarkResults(BaseModel):
         models_res = []
         _revisions = revisions if revisions is not None else [None] * len(names)
 
-        name_rev = {}
+        name_rev: dict[str, str | None] = {}
 
         if len(names) != len(_revisions):
             raise ValueError(
@@ -117,9 +140,12 @@ class BenchmarkResults(BaseModel):
 
         for name, revision in zip(names, _revisions):
             if isinstance(name, ModelMeta):
+                if name.name is None:
+                    raise ValueError("name in ModelMeta is None. It must be a string.")
                 name_rev[name.name] = name.revision
             else:
-                name_rev[name] = revision
+                name_ = cast(str, name)
+                name_rev[name_] = revision
 
         for model_res in self.model_results:
             model_name = model_res.model_name
@@ -139,7 +165,7 @@ class BenchmarkResults(BaseModel):
         n_parameters_range: tuple[int | None, int | None] = (None, None),
         use_instructions: bool | None = None,
         zero_shot_on: list[AbsTask] | None = None,
-    ) -> Self:
+    ) -> BenchmarkResults:
         # mostly a utility function for the leaderboard app.
         # I would probably move the filtering of the models outside of this call. No need to call get_model_metas inside the filter.
         # interface would then be the same as the get_models function
@@ -162,7 +188,7 @@ class BenchmarkResults(BaseModel):
 
         return type(self).model_construct(model_results=new_model_results)
 
-    def join_revisions(self) -> Self:
+    def join_revisions(self) -> BenchmarkResults:
         """Join revisions of the same model.
 
         In case of conflicts, the following rules are applied:
@@ -173,40 +199,6 @@ class BenchmarkResults(BaseModel):
         Returns:
             A new BenchmarkResults object with the revisions joined.
         """
-
-        def parse_version(version_str: str) -> Version | None:
-            try:
-                return Version(version_str)
-            except (InvalidVersion, TypeError):
-                return None
-
-        def keep_best(group: pd.DataFrame) -> pd.DataFrame:
-            # Filtering out task_results where no scores are present
-            group = group[group["has_scores"]]
-            is_main_revision = group["revision"] == group["main_revision"]
-            # If the main revision is present we select that
-            if is_main_revision.sum() > 0:
-                return group[is_main_revision].head(n=1)
-            unique_revisions = group["revision"].unique()
-
-            # ensure None/NA/"external" revisions is filtered out
-            group.loc[group["revision"].isna(), "revision"] = "no_revision_available"
-            group.loc[group["revision"] == "external", "revision"] = (
-                "no_revision_available"
-            )
-
-            # Filtering out no_revision_available if other revisions are present
-            if (len(unique_revisions) > 1) and (
-                "no_revision_available" in unique_revisions
-            ):
-                group = group[group["revision"] != "no_revision_available"]
-            # If there are any not-NA mteb versions, we select the latest one
-            if group["mteb_version"].notna().any():
-                group = group.dropna(subset=["mteb_version"])
-                group = group.sort_values("mteb_version", ascending=False)
-                return group.head(n=1)
-            return group.head(n=1)
-
         records = []
         for model_result in self:
             for task_result in model_result.task_results:
@@ -223,21 +215,58 @@ class BenchmarkResults(BaseModel):
         if not records:
             return BenchmarkResults.model_construct(model_results=[])
         task_df = pd.DataFrame.from_records(records)
-        model_to_main_revision = {
-            meta.name: meta.revision for meta in get_model_metas()
-        }
-        task_df["main_revision"] = task_df["model"].map(model_to_main_revision)  # type: ignore
-        task_df["mteb_version"] = task_df["mteb_version"].map(parse_version)  # type: ignore
-        task_df = (
-            task_df.groupby(["model", "task_name"])
-            .apply(keep_best)
-            .reset_index(drop=True)
+
+        # Use cached model metas
+        model_to_main_revision = _get_cached_model_metas()
+        task_df["main_revision"] = task_df["model"].map(model_to_main_revision)
+
+        # Use cached version parsing
+        task_df["mteb_version"] = task_df["mteb_version"].map(_parse_version_cached)
+
+        # Filter out rows without scores first
+        task_df = task_df[task_df["has_scores"]]
+
+        # Optimize groupby with vectorized operations
+        # Sort by priority: main_revision match, then mteb_version (descending), then revision
+        task_df["is_main_revision"] = task_df["revision"] == task_df["main_revision"]
+
+        # Handle None/NA/external revisions
+        task_df["revision_clean"] = task_df["revision"].copy()
+        task_df.loc[task_df["revision"].isna(), "revision_clean"] = (
+            "no_revision_available"
+        )
+        task_df.loc[task_df["revision"] == "external", "revision_clean"] = (
+            "no_revision_available"
+        )
+
+        # Create a priority column for sorting
+        # Higher priority = better to keep
+        # Priority: main_revision (1000), has valid mteb_version (100), has valid revision (10)
+        task_df["priority"] = 0
+        task_df.loc[task_df["is_main_revision"], "priority"] += 1000
+        task_df.loc[task_df["mteb_version"].notna(), "priority"] += 100
+        task_df.loc[
+            task_df["revision_clean"] != "no_revision_available", "priority"
+        ] += 10
+
+        # Sort by priority (desc), mteb_version (desc), and take first per group
+        task_df = task_df.sort_values(
+            ["model", "task_name", "priority", "mteb_version"],
+            ascending=[True, True, False, False],
+            na_position="last",
         )
+
+        task_df = task_df.groupby(["model", "task_name"], as_index=False).first()
+
+        # Reconstruct model results
         model_results = []
+        # Group by original revision to maintain deterministic behavior
+        # After the first() selection above, each (model, task_name) is unique,
+        # so grouping by original revision ensures consistent ModelResult creation
         for (model, model_revision), group in task_df.groupby(["model", "revision"]):
             model_result = ModelResult.model_construct(
-                model_name=model,
-                model_revision=model_revision,
+                model_name=model,  # type: ignore[arg-type]
+                model_revision=model_revision,  # type: ignore[arg-type]
                 task_results=list(group["task_result"]),
             )
             model_results.append(model_result)
@@ -268,7 +297,7 @@ class BenchmarkResults(BaseModel):
                         {
                             "model": model_res.model_name,
                             "revision": model_res.model_revision,
-                            **model_scores,  # type: ignore
+                            **model_scores,
                         }
                     )
                 except Exception as e:
@@ -296,7 +325,7 @@ class BenchmarkResults(BaseModel):
 
     def to_dataframe(
         self,
-        aggregation_level: Literal["subset", "split", "task"] = "task",
+        aggregation_level: Literal["subset", "split", "task", "language"] = "task",
         aggregation_fn: Callable[[list[Score]], Any] | None = None,
         include_model_revision: bool = False,
         format: Literal["wide", "long"] = "wide",
@@ -321,6 +350,7 @@ class BenchmarkResults(BaseModel):
                 - "subset"/None: No aggregation will be done. The DataFrame will have one row per model, task, split and subset.
                 - "split": Aggregates the scores by split. The DataFrame will have one row per model, task and split.
                 - "task": Aggregates the scores by task. The DataFrame will have one row per model and task.
+                - "language": Aggregates the scores by language. The DataFrame will have one row per model and language.
             aggregation_fn: The function to use for aggregation. If None, the mean will be used.
             include_model_revision: If True, the model revision will be included in the DataFrame. If False, it will be excluded.
                 If there are multiple revisions for the same model, they will be joined using the `join_revisions` method.
@@ -340,7 +370,9 @@ class BenchmarkResults(BaseModel):
             scores_data.extend(model_result._get_score_for_table())
 
         if not scores_data:
-            logger.warning("No scores data available. Returning empty DataFrame.")
+            msg = "No scores data available. Returning empty DataFrame."
+            logger.warning(msg)
+            warnings.warn(msg)
             return pd.DataFrame()
 
         # Create DataFrame
@@ -361,7 +393,24 @@ class BenchmarkResults(BaseModel):
             format=format,
         )
 
-    def __iter__(self) -> Iterator[ModelResult]:
+    def get_benchmark_result(self) -> pd.DataFrame:
+        """Get aggregated scores for each model in the benchmark.
+
+        Uses the benchmark's summary table creation method to compute scores.
+
+        Returns:
+            A DataFrame with the aggregated benchmark scores for each model.
+        """
+        if self.benchmark is None:
+            raise ValueError(
+                "No benchmark associated with these results (self.benchmark is None). "
+                "To get benchmark results, load results with a Benchmark object. "
+                "`results = cache.load_results(tasks='MTEB(eng, v2)')`"
+            )
+
+        return self.benchmark._create_summary_table(self)
+
+    def __iter__(self) -> Iterator[ModelResult]:  # type: ignore[override]
         return iter(self.model_results)
 
     def __getitem__(self, index: int) -> ModelResult:
@@ -383,11 +432,11 @@ class BenchmarkResults(BaseModel):
             out_file.write(self.model_dump_json(indent=2))
 
     @classmethod
-    def from_validated(cls, **data) -> Self:
+    def from_validated(cls, **data: Any) -> BenchmarkResults:
         """Create BenchmarkResults from validated data.
 
         Args:
-            data: Dictionary containing the data.
+            **data: Arbitrary keyword arguments containing the data.
 
         Returns:
             An instance of BenchmarkResults.