mteb 2.5.2__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,10 +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
@@ -33,11 +35,12 @@ from .model_result import ModelResult, _aggregate_and_pivot
 logger = logging.getLogger(__name__)
 
 
-# Global cache for model metas and version parsing
 @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()}
+    return {
+        meta.name: meta.revision for meta in get_model_metas() if meta.name is not None
+    }
 
 
 @functools.lru_cache(maxsize=10000)
@@ -77,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(
@@ -97,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:
@@ -115,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:
@@ -128,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(
@@ -137,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
@@ -159,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
@@ -182,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:
@@ -212,10 +218,10 @@ class BenchmarkResults(BaseModel):
 
         # Use cached model metas
         model_to_main_revision = _get_cached_model_metas()
-        task_df["main_revision"] = task_df["model"].map(model_to_main_revision)  # type: ignore
+        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)  # type: ignore
+        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"]]
@@ -259,8 +265,8 @@ class BenchmarkResults(BaseModel):
         # 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)
@@ -291,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:
@@ -364,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
@@ -402,7 +410,7 @@ class BenchmarkResults(BaseModel):
 
         return self.benchmark._create_summary_table(self)
 
-    def __iter__(self) -> Iterator[ModelResult]:
+    def __iter__(self) -> Iterator[ModelResult]:  # type: ignore[override]
         return iter(self.model_results)
 
     def __getitem__(self, index: int) -> ModelResult:
@@ -424,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.

mteb/results/model_result.py

@@ -1,12 +1,14 @@
+from __future__ import annotations
+
 import logging
 import warnings
-from collections.abc import Callable, Iterable, Sequence
-from typing import Any, Literal
+from collections.abc import Callable, Iterable
+from typing import Any, Literal, cast
 
 import numpy as np
 import pandas as pd
 from pydantic import BaseModel, ConfigDict, Field
-from typing_extensions import Self
+from typing_extensions import overload
 
 from mteb.abstasks.abstask import AbsTask
 from mteb.abstasks.task_metadata import (
@@ -58,7 +60,7 @@ def _aggregate_and_pivot(
             index=index_columns,
             columns=columns,
             values="score",
-            aggfunc=aggregation_fn,
+            aggfunc=aggregation_fn,  # type: ignore[arg-type]
         ).reset_index()
     elif format == "long":
         return (
@@ -81,7 +83,7 @@ class ModelResult(BaseModel):
     model_revision: str | None
     task_results: list[TaskResult]
     default_modalities: list[Modalities] = Field(
-        default_factory=lambda: ["text"], alias="modalities"
+        default_factory=lambda: [cast(Modalities, "text")], alias="modalities"
     )
     model_config = (
         ConfigDict(  # to free up the name model_* which is otherwise protected
@@ -95,16 +97,17 @@ class ModelResult(BaseModel):
         return f"ModelResult(model_name={self.model_name}, model_revision={self.model_revision}, task_results=[...](#{n_entries}))"
 
     @classmethod
-    def from_validated(cls, **data: dict[str, Any]) -> Self:
+    def from_validated(cls, **data: dict[str, Any]) -> ModelResult:
         """Create a ModelResult from validated data.
 
         Args:
             data: The validated data.
         """
-        data["task_results"] = [
-            TaskResult.from_validated(**res) for res in data["task_results"]
+        data["task_results"] = [  # type: ignore[assignment]
+            TaskResult.from_validated(**res)  # type: ignore[arg-type]
+            for res in data["task_results"]
         ]
-        return cls.model_construct(**data)
+        return cls.model_construct(**data)  # type: ignore[arg-type]
 
     def _filter_tasks(
         self,
@@ -114,7 +117,7 @@ class ModelResult(BaseModel):
         task_types: list[TaskType] | None = None,
         modalities: list[Modalities] | None = None,
         is_public: bool | None = None,
-    ) -> Self:
+    ) -> ModelResult:
         new_task_results = []
         for task_result in self.task_results:
             if (task_names is not None) and (task_result.task_name not in task_names):
@@ -142,7 +145,7 @@ class ModelResult(BaseModel):
             task_results=new_task_results,
         )
 
-    def select_tasks(self, tasks: Sequence[AbsTask]) -> Self:
+    def select_tasks(self, tasks: Iterable[AbsTask]) -> ModelResult:
         """Select tasks from the ModelResult based on a list of AbsTask objects.
 
         Args:
@@ -160,6 +163,28 @@ class ModelResult(BaseModel):
             task_results=new_task_results,
         )
 
+    @overload
+    def _get_scores(
+        self,
+        splits: list[SplitName] | None = None,
+        languages: list[ISOLanguage | ISOLanguageScript] | None = None,
+        scripts: list[ISOLanguageScript] | None = None,
+        getter: Callable[[ScoresDict], Score] | None = None,
+        aggregation: Callable[[list[Score]], Any] | None = None,
+        format: Literal["wide"] = "wide",
+    ) -> dict: ...
+
+    @overload
+    def _get_scores(
+        self,
+        splits: list[SplitName] | None = None,
+        languages: list[ISOLanguage | ISOLanguageScript] | None = None,
+        scripts: list[ISOLanguageScript] | None = None,
+        getter: Callable[[ScoresDict], Score] | None = None,
+        aggregation: Callable[[list[Score]], Any] | None = None,
+        format: Literal["long"] = "long",
+    ) -> list: ...
+
     def _get_scores(
         self,
         splits: list[SplitName] | None = None,
@@ -177,21 +202,24 @@ class ModelResult(BaseModel):
             aggregation = aggregation if aggregation is not None else np.mean
         else:
             use_fast = True
+        aggregation = cast(Callable[[list[Score]], Any], aggregation)
+        getter = cast(Callable[[ScoresDict], Score], getter)
+
         if format == "wide":
             scores = {}
             for res in self.task_results:
                 try:
                     if use_fast:
                         scores[res.task_name] = res._get_score_fast(
-                            splits=splits,  # type: ignore
-                            languages=languages,  # type: ignore
+                            splits=splits,
+                            languages=languages,
                         )
                     else:
                         scores[res.task_name] = res.get_score(
                             splits=splits,
                             languages=languages,
-                            aggregation=aggregation,  # type: ignore
-                            getter=getter,  # type: ignore
+                            aggregation=aggregation,
+                            getter=getter,
                             scripts=scripts,
                         )
                 except Exception as e:
@@ -206,14 +234,14 @@ class ModelResult(BaseModel):
                     if use_fast:
                         score = task_res._get_score_fast(
                             splits=splits,
-                            languages=languages,  # type: ignore
+                            languages=languages,
                         )
                     else:
                         score = task_res.get_score(
                             splits=splits,
                             languages=languages,
-                            aggregation=aggregation,  # type: ignore
-                            getter=getter,  # type: ignore
+                            aggregation=aggregation,
+                            getter=getter,
                             scripts=scripts,
                         )
                     entry = dict(
@@ -292,7 +320,9 @@ class ModelResult(BaseModel):
         scores_data = self._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
@@ -315,7 +345,7 @@ class ModelResult(BaseModel):
     def __hash__(self) -> int:
         return id(self)
 
-    def __iter__(self) -> Iterable[TaskResult]:
+    def __iter__(self) -> Iterable[TaskResult]:  # type: ignore[override]
         return iter(self.task_results)
 
     def __getitem__(self, index) -> TaskResult:
@@ -368,13 +398,13 @@ class ModelResult(BaseModel):
         return [task_res.task_name for task_res in self.task_results]
 
     @property
-    def modalities(self) -> list[str]:
+    def modalities(self) -> list[Modalities]:
         """Get all modalities in the task results.
 
         Returns:
             A list of modalities in the task results.
         """
-        mods = []
+        mods: list[Modalities] = []
         for task_res in self.task_results:
             task_modalities = getattr(task_res, "modalities", [])
             mods.extend(task_modalities)