diff --git a/mteb/leaderboard/table.py b/mteb/leaderboard/table.py index b3215dd067..9856493c74 100644 --- a/mteb/leaderboard/table.py +++ b/mteb/leaderboard/table.py @@ -106,6 +106,14 @@ def failsafe_get_model_meta(model_name): 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 scores_to_tables( scores_long: list[dict], search_query: str | None = None ) -> tuple[gr.DataFrame, gr.DataFrame]: @@ -145,9 +153,7 @@ def scores_to_tables( joint_table.insert( 1, "Max Tokens", - model_metas.map( - lambda m: str(int(m.max_tokens)) if m.max_tokens else "Unknown" - ), + model_metas.map(lambda m: format_max_tokens(m.max_tokens)), ) joint_table.insert( 1, diff --git a/mteb/load_results/task_results.py b/mteb/load_results/task_results.py index b3b1f8cba2..cd2a2c0847 100644 --- a/mteb/load_results/task_results.py +++ b/mteb/load_results/task_results.py @@ -511,8 +511,8 @@ def validate_and_filter_scores(self, task: AbsTask | None = None) -> AbsTask: new_scores[split].append(_scores) seen_subsets.add(_scores["hf_subset"]) if seen_subsets != hf_subsets: - raise ValueError( - f"Missing subsets {hf_subsets - seen_subsets} for split {split}" + logger.warning( + f"{task.metadata.name}: Missing subsets {hf_subsets - seen_subsets} for split {split}" ) seen_splits.add(split) if seen_splits != set(splits): diff --git a/mteb/model_meta.py b/mteb/model_meta.py index 83653ec3d1..3f2cc2843f 100644 --- a/mteb/model_meta.py +++ b/mteb/model_meta.py @@ -25,6 +25,7 @@ "TensorFlow", "API", "Tevatron", + "NumPy", ] DISTANCE_METRICS = Literal["cosine"] @@ -87,7 +88,7 @@ class ModelMeta(BaseModel): loader: Callable[..., Encoder] | None = None n_parameters: int | None = None memory_usage: float | None = None - max_tokens: int | None = None + max_tokens: float | None = None embed_dim: int | None = None license: str | None = None open_weights: bool | None = None diff --git a/mteb/models/model2vec_models.py b/mteb/models/model2vec_models.py new file mode 100644 index 0000000000..0dfdb193cd --- /dev/null +++ b/mteb/models/model2vec_models.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +import logging +from collections.abc import Sequence +from functools import partial +from typing import Any + +import numpy as np + +from mteb.encoder_interface import PromptType +from mteb.model_meta import ModelMeta + +from .wrapper import Wrapper + +logger = logging.getLogger(__name__) + + +class Model2VecWrapper(Wrapper): + def __init__( + self, + model_name: str, + **kwargs, + ) -> None: + """Wrapper for Model2Vec models. + + Args: + model_name: The Model2Vec model to load from HuggingFace Hub. + """ + try: + from model2vec import StaticModel + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + "To use the Model2Vec models `model2vec` is required. Please install it with `pip install mteb[model2vec]`." + ) from e + + self.model_name = model_name + self.static_model = StaticModel.from_pretrained(self.model_name) + + def encode( + self, + sentences: Sequence[str], + **kwargs: Any, + ) -> np.ndarray: + """Encodes the given sentences using the encoder. + + Args: + sentences: The sentences to encode. + **kwargs: Additional arguments to pass to the encoder. + + Returns: + The encoded sentences. + """ + return self.static_model.encode(sentences) + + +m2v_base_glove_subword = ModelMeta( + loader=partial( + Model2VecWrapper, + model_name="minishlab/M2V_base_glove_subword", + ), + name="minishlab/M2V_base_glove_subword", + languages=["eng_Latn"], + open_weights=True, + revision="5f4f5ca159b7321a8b39739bba0794fa0debddf4", + release_date="2024-09-21", + n_parameters=103 * 1e6, + max_tokens=np.inf, # Theoretically infinite + embed_dim=256, + license="mit", + similarity_fn_name="cosine", + framework=["NumPy"], + reference="https://huggingface.co/minishlab/M2V_base_glove_subword", + use_instructions=False, + adapted_from="BAAI/bge-base-en-v1.5", + superseded_by=None, +) + + +m2v_base_glove = ModelMeta( + loader=partial( + Model2VecWrapper, + model_name="minishlab/M2V_base_glove", + ), + name="minishlab/M2V_base_glove", + languages=["eng_Latn"], + open_weights=True, + revision="38ebd7f10f71e67fa8db898290f92b82e9cfff2b", + release_date="2024-09-21", + n_parameters=102 * 1e6, + max_tokens=np.inf, + embed_dim=256, + license="mit", + similarity_fn_name="cosine", + framework=["NumPy"], + reference="https://huggingface.co/minishlab/M2V_base_glove", + use_instructions=False, + adapted_from="BAAI/bge-base-en-v1.5", + superseded_by=None, +) + +m2v_base_output = ModelMeta( + loader=partial( + Model2VecWrapper, + model_name="minishlab/M2V_base_output", + ), + name="minishlab/M2V_base_output", + languages=["eng_Latn"], + open_weights=True, + revision="02460ae401a22b09d2c6652e23371398329551e2", + release_date="2024-09-21", + n_parameters=7.56 * 1e6, + max_tokens=np.inf, + embed_dim=256, + license="mit", + similarity_fn_name="cosine", + framework=["NumPy"], + reference="https://huggingface.co/minishlab/M2V_base_output", + use_instructions=False, + adapted_from="BAAI/bge-base-en-v1.5", + superseded_by=None, +) + +m2v_multilingual_output = ModelMeta( + loader=partial( + Model2VecWrapper, + model_name="minishlab/M2V_multilingual_output", + ), + name="minishlab/M2V_multilingual_output", + languages=["eng_Latn"], + open_weights=True, + revision="2cf4ec4e1f51aeca6c55cf9b93097d00711a6305", + release_date="2024-09-21", + n_parameters=128 * 1e6, + max_tokens=np.inf, + embed_dim=256, + license="mit", + similarity_fn_name="cosine", + framework=["NumPy"], + reference="https://huggingface.co/minishlab/M2V_multilingual_output", + use_instructions=False, + adapted_from="sentence-transformers/LaBSE", + superseded_by=None, +) + +potion_base_2m = ModelMeta( + loader=partial( + Model2VecWrapper, + model_name="minishlab/potion-base-2M", + ), + name="minishlab/potion-base-2M", + languages=["eng_Latn"], + open_weights=True, + revision="86db093558fbced2072b929eb1690bce5272bd4b", + release_date="2024-10-29", + n_parameters=2 * 1e6, + max_tokens=np.inf, + embed_dim=64, + license="mit", + similarity_fn_name="cosine", + framework=["NumPy"], + reference="https://huggingface.co/minishlab/potion-base-2M", + use_instructions=False, + adapted_from="BAAI/bge-base-en-v1.5", + superseded_by=None, +) + +potion_base_4m = ModelMeta( + loader=partial( + Model2VecWrapper, + model_name="minishlab/potion-base-4M", + ), + name="minishlab/potion-base-4M", + languages=["eng_Latn"], + open_weights=True, + revision="81b1802ada41afcd0987a37dc15e569c9fa76f04", + release_date="2024-10-29", + n_parameters=3.78 * 1e6, + max_tokens=np.inf, + embed_dim=128, + license="mit", + similarity_fn_name="cosine", + framework=["NumPy"], + reference="https://huggingface.co/minishlab/potion-base-4M", + use_instructions=False, + adapted_from="BAAI/bge-base-en-v1.5", + superseded_by=None, +) + +potion_base_8m = ModelMeta( + loader=partial( + Model2VecWrapper, + model_name="minishlab/potion-base-8M", + ), + name="minishlab/potion-base-8M", + languages=["eng_Latn"], + open_weights=True, + revision="dcbec7aa2d52fc76754ac6291803feedd8c619ce", + release_date="2024-10-29", + n_parameters=7.56 * 1e6, + max_tokens=np.inf, + embed_dim=256, + license="mit", + similarity_fn_name="cosine", + framework=["NumPy"], + reference="https://huggingface.co/minishlab/potion-base-8M", + use_instructions=False, + adapted_from="BAAI/bge-base-en-v1.5", + superseded_by=None, +) diff --git a/mteb/models/overview.py b/mteb/models/overview.py index 8341e42cdb..7418ee98fa 100644 --- a/mteb/models/overview.py +++ b/mteb/models/overview.py @@ -22,6 +22,7 @@ gte_models, jina_models, llm2vec_models, + model2vec_models, mxbai_models, nomic_models, openai_models, @@ -51,6 +52,7 @@ gte_models, llm2vec_models, mxbai_models, + model2vec_models, nomic_models, openai_models, ru_sentence_models, diff --git a/pyproject.toml b/pyproject.toml index c79112605a..79b9226e6a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,7 @@ flagembedding = ["FlagEmbedding"] jina = ["einops>=0.8.0"] flash_attention = ["flash-attn>=2.6.3"] openai = ["openai>=1.41.0", "tiktoken>=0.8.0"] +model2vec = ["model2vec>=0.3.0"] [tool.coverage.report]