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

mteb/models/model_implementations/voyage_models.py

@@ -25,6 +25,9 @@ VOYAGE_DTYPE_TRANSLATION = {
 
 # Total token limits per model based on VoyageAI documentation
 VOYAGE_TOTAL_TOKEN_LIMITS = {
+    "voyage-4-large": 120_000,
+    "voyage-4": 320_000,
+    "voyage-4-lite": 1_000_000,
     "voyage-3.5-lite": 1_000_000,
     "voyage-3.5": 320_000,
     "voyage-2": 320_000,
@@ -206,6 +209,84 @@ model_prompts = {
     PromptType.document.value: "document",
 }
 
+voyage_4 = ModelMeta(
+    name="voyageai/voyage-4",
+    model_type=["dense"],
+    revision="1",
+    release_date="2026-01-15",
+    languages=None,  # supported languages not specified
+    loader=VoyageModel,
+    loader_kwargs=dict(
+        max_tokens=32000,
+        model_prompts=model_prompts,
+    ),
+    max_tokens=32000,
+    embed_dim=1024,
+    open_weights=False,
+    n_parameters=None,
+    memory_usage_mb=None,
+    license=None,
+    reference="https://blog.voyageai.com/2026/01/15/voyage-4/",
+    similarity_fn_name="cosine",
+    framework=["API"],
+    use_instructions=True,
+    training_datasets=VOYAGE_TRAINING_DATA,
+    public_training_code=None,
+    public_training_data=None,
+)
+
+voyage_4_lite = ModelMeta(
+    name="voyageai/voyage-4-lite",
+    model_type=["dense"],
+    revision="1",
+    release_date="2026-01-15",
+    languages=None,  # supported languages not specified
+    loader=VoyageModel,
+    loader_kwargs=dict(
+        max_tokens=32000,
+        model_prompts=model_prompts,
+    ),
+    max_tokens=32000,
+    embed_dim=1024,
+    open_weights=False,
+    n_parameters=None,
+    memory_usage_mb=None,
+    license=None,
+    reference="https://blog.voyageai.com/2026/01/15/voyage-4/",
+    similarity_fn_name="cosine",
+    framework=["API"],
+    use_instructions=True,
+    training_datasets=VOYAGE_TRAINING_DATA,
+    public_training_code=None,
+    public_training_data=None,
+)
+
+voyage_4_large = ModelMeta(
+    name="voyageai/voyage-4-large",
+    model_type=["dense"],
+    revision="1",
+    release_date="2026-01-15",
+    languages=None,  # supported languages not specified
+    loader=VoyageModel,
+    loader_kwargs=dict(
+        max_tokens=32000,
+        model_prompts=model_prompts,
+    ),
+    max_tokens=32000,
+    embed_dim=1024,
+    open_weights=False,
+    n_parameters=None,
+    memory_usage_mb=None,
+    license=None,
+    reference="https://blog.voyageai.com/2026/01/15/voyage-4/",
+    similarity_fn_name="cosine",
+    framework=["API"],
+    use_instructions=True,
+    training_datasets=VOYAGE_TRAINING_DATA,
+    public_training_code=None,
+    public_training_data=None,
+)
+
 voyage_3_large = ModelMeta(
     name="voyageai/voyage-3-large",  # Date of publication of this post https://blog.voyageai.com/2025/01/07/voyage-3-large/
     model_type=["dense"],
@@ -230,6 +311,7 @@ voyage_3_large = ModelMeta(
     training_datasets=VOYAGE_TRAINING_DATA,
     public_training_code=None,
     public_training_data=None,
+    superseded_by="voyageai/voyage-4-large",
 )
 
 
@@ -257,6 +339,7 @@ voyage_3_5 = ModelMeta(
     training_datasets=VOYAGE_TRAINING_DATA,
     public_training_code=None,
     public_training_data=None,
+    superseded_by="voyageai/voyage-4",
 )
 
 voyage_3_5_int8 = ModelMeta(
@@ -571,6 +654,7 @@ voyage_3_lite = ModelMeta(
     training_datasets=VOYAGE_TRAINING_DATA,
     public_training_code=None,
     public_training_data=None,
+    superseded_by="voyageai/voyage-4-lite",
 )
 
 voyage_3_exp = ModelMeta(

mteb/models/model_implementations/voyage_v.py

@@ -16,6 +16,8 @@ from mteb.types import Array, BatchedInput, PromptType
 if TYPE_CHECKING:
     from PIL import Image
 
+logger = logging.getLogger(__name__)
+
 
 def _downsample_image(
     image: Image.Image, max_pixels: int = 16000000, target_longest_side: int = 4000
@@ -37,18 +39,18 @@ def _downsample_image(
             new_width = int(width * (target_longest_side / height))
 
         new_size = (new_width, new_height)
-        logging.info(
+        logger.info(
             f"Downsampling image from {width}x{height} to {new_width}x{new_height}"
         )
-        return image.resize(new_size, Image.LANCZOS)  # type: ignore
+        return image.resize(new_size, Image.LANCZOS)
     if width > height:
         if width > 10000:
-            logging.error("Processing extremely wide images.")
-            return image.resize((10000, height), Image.LANCZOS)  # type: ignore
+            logger.error("Processing extremely wide images.")
+            return image.resize((10000, height), Image.LANCZOS)
     else:
         if height > 10000:
-            logging.error("Processing extremely high images.")
-            return image.resize((width, 10000), Image.LANCZOS)  # type: ignore
+            logger.error("Processing extremely high images.")
+            return image.resize((width, 10000), Image.LANCZOS)
     return image
 
 
@@ -202,7 +204,7 @@ def voyage_v_loader(model_name, **kwargs):
 
 
 voyage_v = ModelMeta(
-    loader=voyage_v_loader,  # type: ignore
+    loader=voyage_v_loader,
     name="voyageai/voyage-multimodal-3",
     model_type=["dense"],
     languages=[],  # Unknown

mteb/models/model_implementations/youtu_models.py

@@ -127,7 +127,7 @@ Youtu_Embedding_V1 = ModelMeta(
     max_tokens=8192,
     reference="https://huggingface.co/tencent/Youtu-Embedding",
     similarity_fn_name="cosine",
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "Transformers", "safetensors"],
     use_instructions=True,
     public_training_code=None,
     public_training_data=None,

mteb/models/model_implementations/yuan_models.py

@@ -26,7 +26,7 @@ yuan_embedding_2_zh = ModelMeta(
     max_tokens=512,
     reference="https://huggingface.co/IEITYuan/Yuan-embedding-2.0-zh",
     similarity_fn_name="cosine",
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "safetensors"],
     use_instructions=False,
     public_training_code=None,
     public_training_data=None,

mteb/models/model_implementations/yuan_models_en.py

@@ -49,7 +49,7 @@ yuan_embedding_2_en = ModelMeta(
     license="apache-2.0",
     reference="https://huggingface.co/IEITYuan/Yuan-embedding-2.0-en",
     similarity_fn_name="cosine",
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "safetensors"],
     use_instructions=True,
     public_training_code=None,
     public_training_data=None,

mteb/models/model_meta.py

@@ -17,11 +17,13 @@ from huggingface_hub import (
     get_safetensors_metadata,
     hf_hub_download,
     list_repo_commits,
+    model_info,
     repo_exists,
 )
 from huggingface_hub.errors import (
     EntryNotFoundError,
     GatedRepoError,
+    HFValidationError,
     NotASafetensorsRepoError,
     RepositoryNotFoundError,
     SafetensorsParsingError,
@@ -55,6 +57,10 @@ FRAMEWORKS = Literal[
     "PyLate",
     "ColBERT",
     "ColPali",
+    "GGUF",
+    "safetensors",
+    "ONNX",
+    "Transformers",
 ]
 
 MODEL_TYPES = Literal["dense", "cross-encoder", "late-interaction"]
@@ -81,9 +87,6 @@ def _get_loader_name(
     return loader.__name__
 
 
-_SENTENCE_TRANSFORMER_LIB_NAME = "Sentence Transformers"
-
-
 class ModelMeta(BaseModel):
     """The model metadata object.
 
@@ -250,7 +253,7 @@ class ModelMeta(BaseModel):
             )
         return v
 
-    def load_model(self, **kwargs: Any) -> MTEBModels:
+    def load_model(self, device: str | None = None, **kwargs: Any) -> MTEBModels:
         """Loads the model using the specified loader function."""
         if self.loader is None:
             raise NotImplementedError(
@@ -262,11 +265,11 @@ class ModelMeta(BaseModel):
         # Allow overwrites
         _kwargs = self.loader_kwargs.copy()
         _kwargs.update(kwargs)
+        if device is not None:
+            _kwargs["device"] = device
 
-        model: EncoderProtocol = self.loader(
-            self.name, revision=self.revision, **_kwargs
-        )
-        model.mteb_model_meta = self  # type: ignore
+        model: MTEBModels = self.loader(self.name, revision=self.revision, **_kwargs)
+        model.mteb_model_meta = self  # type: ignore[misc]
         return model
 
     def model_name_as_path(self) -> str:
@@ -307,7 +310,7 @@ class ModelMeta(BaseModel):
         embedding_dim = None
         max_tokens = None
 
-        if model_name and compute_metadata and repo_exists(model_name):
+        if model_name and compute_metadata and _repo_exists(model_name):
             reference = "https://huggingface.co/" + model_name
             card = ModelCard.load(model_name)
             card_data: ModelCardData = card.data
@@ -318,22 +321,17 @@ class ModelMeta(BaseModel):
                 model_config = None
                 logger.warning(f"Can't get configuration for {model_name}. Error: {e}")
 
-            if (
-                card_data.library_name == _SENTENCE_TRANSFORMER_LIB_NAME
-                or _SENTENCE_TRANSFORMER_LIB_NAME in card_data.tags
-            ):
-                frameworks.append(_SENTENCE_TRANSFORMER_LIB_NAME)
-            else:
-                msg = "Model library not recognized, defaulting to Sentence Transformers loader."
-                logger.warning(msg)
-                warnings.warn(msg)
+            hf_frameworks = (
+                cls._get_frameworks_from_hf_tags(model_name) if model_name else []
+            )
+            frameworks.extend(hf_frameworks)
 
             if revision is None:
                 revisions = _get_repo_commits(model_name, "model")
                 revision = revisions[0].commit_id if revisions else None
 
             release_date = cls.fetch_release_date(model_name)
-            model_license = card_data.license
+            model_license = card_data.license if card_data.license != "other" else None
             n_parameters = cls._calculate_num_parameters_from_hub(model_name)
             memory_usage_mb = cls._calculate_memory_usage_mb(model_name, n_parameters)
             if model_config and hasattr(model_config, "hidden_size"):
@@ -386,8 +384,6 @@ class ModelMeta(BaseModel):
             else model.model_card_data.base_model
         )
         meta = cls._from_hub(name, revision, compute_metadata)
-        if _SENTENCE_TRANSFORMER_LIB_NAME not in meta.framework:
-            meta.framework.append("Sentence Transformers")
         meta.revision = model.model_card_data.base_model_revision or meta.revision
         meta.max_tokens = model.max_seq_length
         meta.embed_dim = model.get_sentence_embedding_dimension()
@@ -413,11 +409,9 @@ class ModelMeta(BaseModel):
             The generated ModelMeta.
         """
         meta = cls._from_hub(model, revision, compute_metadata)
-        if _SENTENCE_TRANSFORMER_LIB_NAME not in meta.framework:
-            meta.framework.append("Sentence Transformers")
         meta.modalities = ["text"]
 
-        if model and compute_metadata and repo_exists(model):
+        if model and compute_metadata and _repo_exists(model):
             # have max_seq_length field
             sbert_config = _get_json_from_hub(
                 model, "sentence_bert_config.json", "model", revision=revision
@@ -435,7 +429,7 @@ class ModelMeta(BaseModel):
                 and config_sbert.get("similarity_fn_name") is not None
             ):
                 meta.similarity_fn_name = ScoringFunction.from_str(
-                    config_sbert.get("similarity_fn_name")
+                    config_sbert["similarity_fn_name"]
                 )
             else:
                 meta.similarity_fn_name = ScoringFunction.COSINE
@@ -461,8 +455,6 @@ class ModelMeta(BaseModel):
         from mteb.models import CrossEncoderWrapper
 
         meta = cls._from_hub(model.model.name_or_path, revision, compute_metadata)
-        if _SENTENCE_TRANSFORMER_LIB_NAME not in meta.framework:
-            meta.framework.append("Sentence Transformers")
         meta.revision = model.config._commit_hash or meta.revision
         meta.loader = CrossEncoderWrapper
         meta.embed_dim = None
@@ -511,10 +503,12 @@ class ModelMeta(BaseModel):
                 if adapted_training_datasets is not None:
                     training_datasets |= adapted_training_datasets
             except (ValueError, KeyError) as e:
-                logger.warning(f"Could not get source model: {e} in MTEB")
+                msg = f"Could not get source model: {e} in MTEB"
+                logger.warning(msg)
+                warnings.warn(msg)
 
         return_dataset = training_datasets.copy()
-        visited = set()
+        visited: set[str] = set()
 
         for dataset in training_datasets:
             similar_tasks = _collect_similar_tasks(dataset, visited)
@@ -548,6 +542,8 @@ class ModelMeta(BaseModel):
 
     @staticmethod
     def _calculate_num_parameters_from_hub(model_name: str | None = None) -> int | None:
+        if not model_name:
+            return None
         try:
             safetensors_metadata = get_safetensors_metadata(model_name)
             if len(safetensors_metadata.parameter_count) >= 0:
@@ -561,7 +557,7 @@ class ModelMeta(BaseModel):
             logger.warning(
                 f"Can't calculate number of parameters for {model_name}. Got error {e}"
             )
-            return None
+        return None
 
     def calculate_num_parameters_from_hub(self) -> int | None:
         """Calculates the number of parameters in the model.
@@ -624,7 +620,7 @@ class ModelMeta(BaseModel):
         if "API" in self.framework or self.name is None:
             return None
 
-        return self._calculate_memory_usage_mb(self.model_name, self.n_parameters)
+        return self._calculate_memory_usage_mb(self.name, self.n_parameters)
 
     @staticmethod
     def fetch_release_date(model_name: str) -> StrDate | None:
@@ -640,6 +636,43 @@ class ModelMeta(BaseModel):
             return release_date
         return None
 
+    @staticmethod
+    def _get_frameworks_from_hf_tags(model_name: str) -> list[FRAMEWORKS]:
+        """Extract frameworks supported by the model from HuggingFace model tags.
+
+        Args:
+            model_name: HuggingFace model name
+
+        Returns:
+            List of framework names found in tags. Defaults to empty list if no frameworks found.
+        """
+        try:
+            info = model_info(model_name)
+            if not info.tags:
+                return []
+        except Exception as e:
+            logger.warning(
+                f"Failed to fetch frameworks from HuggingFace tags for {model_name}: {e}"
+            )
+            return []
+
+        # Mapping from HuggingFace tags to MTEB framework names
+        tag_to_framework: dict[str, FRAMEWORKS] = {
+            "sentence-transformers": "Sentence Transformers",
+            "transformers": "Transformers",
+            "onnx": "ONNX",
+            "safetensors": "safetensors",
+            "gguf": "GGUF",
+        }
+
+        frameworks: list[FRAMEWORKS] = []
+
+        for framework_tag in tag_to_framework.keys():
+            if framework_tag in info.tags:
+                frameworks.append(tag_to_framework[framework_tag])
+
+        return frameworks
+
     def to_python(self) -> str:
         """Returns a string representation of the model."""
         return _pydantic_instance_to_code(self)
@@ -784,3 +817,19 @@ def _get_file_on_hub(
     except (GatedRepoError, RepositoryNotFoundError, EntryNotFoundError) as e:
         logger.warning(f"Can't get file {file_name} of {repo_id}: {e}")
         return None
+
+
+def _repo_exists(repo_id: str, repo_type: str | None = None) -> bool:
+    """Checks if a repository exists on HuggingFace Hub.
+
+    Repo exists will raise HFValidationError for invalid local paths
+
+    Args:
+        repo_id: The repository ID.
+        repo_type: The type of repository (e.g., "model", "dataset", "space").
+    """
+    try:
+        return repo_exists(repo_id=repo_id, repo_type=repo_type)
+    except HFValidationError as e:
+        logger.warning(f"Can't check existence of {repo_id}: {e}")
+        return False

mteb/models/models_protocols.py

@@ -1,12 +1,14 @@
 from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 
 from torch.utils.data import DataLoader
+from typing_extensions import Unpack
 
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.types import (
     Array,
     BatchedInput,
     CorpusDatasetType,
+    EncodeKwargs,
     PromptType,
     QueryDatasetType,
     RetrievalOutputType,
@@ -28,7 +30,7 @@ class SearchProtocol(Protocol):
         task_metadata: TaskMetadata,
         hf_split: str,
         hf_subset: str,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
     ) -> None:
         """Index the corpus for retrieval.
 
@@ -49,7 +51,7 @@ class SearchProtocol(Protocol):
         hf_split: str,
         hf_subset: str,
         top_k: int,
-        encode_kwargs: dict[str, Any],
+        encode_kwargs: EncodeKwargs,
         top_ranked: TopRankedDocumentsType | None = None,
     ) -> RetrievalOutputType:
         """Search the corpus using the given queries.
@@ -83,12 +85,19 @@ class EncoderProtocol(Protocol):
     In general the interface is kept aligned with sentence-transformers interface. In cases where exceptions occurs these are handled within MTEB.
     """
 
-    def __init__(self, model_name: str, revision: str | None, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        model_name: str,
+        revision: str | None,
+        device: str | None = None,
+        **kwargs: Any,
+    ) -> None:
         """The initialization function for the encoder. Used when calling it from the mteb run CLI.
 
         Args:
             model_name: Name of the model
             revision: revision of the model
+            device: Device used to load the model
             kwargs: Any additional kwargs
         """
         ...
@@ -101,7 +110,7 @@ class EncoderProtocol(Protocol):
         hf_split: str,
         hf_subset: str,
         prompt_type: PromptType | None = None,
-        **kwargs: Any,
+        **kwargs: Unpack[EncodeKwargs],
     ) -> Array:
         """Encodes the given sentences using the encoder.
 
@@ -181,12 +190,19 @@ class CrossEncoderProtocol(Protocol):
     In general the interface is kept aligned with sentence-transformers interface. In cases where exceptions occurs these are handled within MTEB.
     """
 
-    def __init__(self, model_name: str, revision: str | None, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        model_name: str,
+        revision: str | None,
+        device: str | None = None,
+        **kwargs: Any,
+    ) -> None:
         """The initialization function for the encoder. Used when calling it from the mteb run CLI.
 
         Args:
             model_name: Name of the model
             revision: revision of the model
+            device: Device used to load the model
             kwargs: Any additional kwargs
         """
         ...
@@ -200,7 +216,7 @@ class CrossEncoderProtocol(Protocol):
         hf_split: str,
         hf_subset: str,
         prompt_type: PromptType | None = None,
-        **kwargs: Any,
+        **kwargs: Unpack[EncodeKwargs],
     ) -> Array:
         """Predicts relevance scores for pairs of inputs. Note that, unlike the encoder, the cross-encoder can compare across inputs.
 

mteb/models/search_encoder_index/search_indexes/faiss_search_index.py

@@ -1,4 +1,5 @@
 import logging
+import warnings
 from collections.abc import Callable
 
 import numpy as np
@@ -108,7 +109,7 @@ class FaissSearchIndex:
             ids = ids.tolist()
 
         if issubclass(self.index_type, faiss.IndexFlatL2):
-            similarities = -np.sqrt(np.maximum(similarities, 0))
+            similarities = (-np.sqrt(np.maximum(similarities, 0))).tolist()
 
         return similarities, ids
 
@@ -116,8 +117,8 @@ class FaissSearchIndex:
         self,
         embeddings: Array,
         top_k: int,
-        top_ranked: TopRankedDocumentsType | None = None,
-        query_idx_to_id: dict[int, str] | None = None,
+        top_ranked: TopRankedDocumentsType,
+        query_idx_to_id: dict[int, str],
     ) -> tuple[list[list[float]], list[list[int]]]:
         doc_id_to_idx = {doc_id: i for i, doc_id in enumerate(self.idxs)}
         scores_all: list[list[float]] = []
@@ -127,15 +128,17 @@ class FaissSearchIndex:
             query_id = query_idx_to_id[query_idx]
             ranked_ids = top_ranked.get(query_id)
             if not ranked_ids:
-                logger.warning(f"No top-ranked documents for query {query_id}")
+                msg = f"No top-ranked documents for query {query_id}"
+                logger.warning(msg)
+                warnings.warn(msg)
                 scores_all.append([])
                 idxs_all.append([])
                 continue
 
             candidate_indices = [doc_id_to_idx[doc_id] for doc_id in ranked_ids]
-            d = self.index.d
+            d = self.index.d  # type: ignore[union-attr]
             candidate_embs = np.vstack(
-                [self.index.reconstruct(idx) for idx in candidate_indices]
+                [self.index.reconstruct(idx) for idx in candidate_indices]  # type: ignore[union-attr]
             )
             sub_reranking_index = self.index_type(d)
             sub_reranking_index.add(candidate_embs)