diff --git a/mteb/models/arctic_models.py b/mteb/models/arctic_models.py index 5f3d41a97..26af9c8cc 100644 --- a/mteb/models/arctic_models.py +++ b/mteb/models/arctic_models.py @@ -5,7 +5,7 @@ from mteb.model_meta import ModelMeta, sentence_transformers_loader arctic_m_v1_5 = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="Snowflake/snowflake-arctic-embed-m-v1.5", revision="97eab2e17fcb7ccb8bb94d6e547898fa1a6a0f47", diff --git a/mteb/models/bge_models.py b/mteb/models/bge_models.py index 5ab429479..b643b4dfb 100644 --- a/mteb/models/bge_models.py +++ b/mteb/models/bge_models.py @@ -7,7 +7,7 @@ model_prompts = {"query": "Represent this sentence for searching relevant passages: "} bge_small_en_v1_5 = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="BAAI/bge-small-en-v1.5", revision="5c38ec7c405ec4b44b94cc5a9bb96e735b38267a", @@ -30,7 +30,7 @@ ) bge_base_en_v1_5 = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="BAAI/bge-base-en-v1.5", revision="a5beb1e3e68b9ab74eb54cfd186867f64f240e1a", @@ -53,7 +53,7 @@ ) bge_large_en_v1_5 = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="BAAI/bge-large-en-v1.5", revision="d4aa6901d3a41ba39fb536a557fa166f842b0e09", diff --git a/mteb/models/cohere_models.py b/mteb/models/cohere_models.py index 26eb5e92e..b41043301 100644 --- a/mteb/models/cohere_models.py +++ b/mteb/models/cohere_models.py @@ -74,7 +74,7 @@ def encode( } cohere_mult_3 = ModelMeta( - loader=partial( + loader=partial( # type: ignore CohereTextEmbeddingModel, model_name="embed-multilingual-v3.0", model_prompts=model_prompts, @@ -95,7 +95,7 @@ def encode( ) cohere_eng_3 = ModelMeta( - loader=partial( + loader=partial( # type: ignore CohereTextEmbeddingModel, model_name="embed-multilingual-v3.0", model_prompts=model_prompts, diff --git a/mteb/models/e5_instruct.py b/mteb/models/e5_instruct.py index 5d5b1f3ad..f26d78ed6 100644 --- a/mteb/models/e5_instruct.py +++ b/mteb/models/e5_instruct.py @@ -12,15 +12,14 @@ MISTRAL_LANGUAGES = ["eng_Latn", "fra_Latn", "deu_Latn", "ita_Latn", "spa_Latn"] -def e5_instruction(instruction: str) -> str: - return f"Instruct: {instruction}\nQuery: " +E5_INSTRUCTION = "Instruct: {instruction}\nQuery: " e5_instruct = ModelMeta( - loader=partial( + loader=partial( # type: ignore instruct_wrapper, model_name_or_path="intfloat/multilingual-e5-large-instruct", - instruction_template=e5_instruction, + instruction_template=E5_INSTRUCTION, attn="cccc", pooling_method="mean", mode="embedding", @@ -44,10 +43,10 @@ def e5_instruction(instruction: str) -> str: ) e5_mistral = ModelMeta( - loader=partial( + loader=partial( # type: ignore instruct_wrapper, model_name_or_path="intfloat/e5-mistral-7b-instruct", - instruction_template=e5_instruction, + instruction_template=E5_INSTRUCTION, attn="cccc", pooling_method="lasttoken", mode="embedding", diff --git a/mteb/models/e5_models.py b/mteb/models/e5_models.py index 612130ed6..943c46870 100644 --- a/mteb/models/e5_models.py +++ b/mteb/models/e5_models.py @@ -114,7 +114,7 @@ } e5_mult_small = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="intfloat/multilingual-e5-small", revision="fd1525a9fd15316a2d503bf26ab031a61d056e98", @@ -137,7 +137,7 @@ ) e5_mult_base = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="intfloat/multilingual-e5-base", model_prompts=model_prompts, @@ -159,7 +159,7 @@ ) e5_mult_large = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="intfloat/multilingual-e5-large", revision="ab10c1a7f42e74530fe7ae5be82e6d4f11a719eb", @@ -182,7 +182,7 @@ ) e5_eng_small_v2 = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="intfloat/e5-small-v2", model_prompts=model_prompts, @@ -204,7 +204,7 @@ ) e5_eng_small = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="intfloat/e5-small", revision="e272f3049e853b47cb5ca3952268c6662abda68f", @@ -227,7 +227,7 @@ ) e5_eng_base_v2 = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="intfloat/e5-base-v2", revision="1c644c92ad3ba1efdad3f1451a637716616a20e8", @@ -250,7 +250,7 @@ ) e5_eng_large_v2 = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="intfloat/e5-large-v2", revision="b322e09026e4ea05f42beadf4d661fb4e101d311", diff --git a/mteb/models/google_models.py b/mteb/models/google_models.py index 688680abc..8066599b0 100644 --- a/mteb/models/google_models.py +++ b/mteb/models/google_models.py @@ -78,7 +78,7 @@ def encode( name = "text-embedding-004" google_emb_004 = ModelMeta( - loader=partial( + loader=partial( # type: ignore GoogleTextEmbeddingModel, model_name=name, model_prompts={ diff --git a/mteb/models/gritlm_models.py b/mteb/models/gritlm_models.py index b1c4882bc..1e9c85eb8 100644 --- a/mteb/models/gritlm_models.py +++ b/mteb/models/gritlm_models.py @@ -17,7 +17,7 @@ def gritlm_instruction(instruction: str = "") -> str: gritlm7b = ModelMeta( - loader=partial( + loader=partial( # type: ignore instruct_wrapper, model_name_or_path="GritLM/GritLM-7B", instruction_template=gritlm_instruction, @@ -40,7 +40,7 @@ def gritlm_instruction(instruction: str = "") -> str: use_instructions=True, ) gritlm8x7b = ModelMeta( - loader=partial( + loader=partial( # type: ignore instruct_wrapper, model_name_or_path="GritLM/GritLM-8x7B", instruction_template=gritlm_instruction, diff --git a/mteb/models/gte_models.py b/mteb/models/gte_models.py index 2358ef6d5..82213b978 100644 --- a/mteb/models/gte_models.py +++ b/mteb/models/gte_models.py @@ -2,15 +2,21 @@ from functools import partial +import torch + from mteb.model_meta import ModelMeta +from mteb.models.instruct_wrapper import instruct_wrapper + + +def instruction_template(instruction: str) -> str: + return f"Instruct: {instruction}\nQuery: " if instruction else "" -from .instruct_wrapper import instruct_wrapper gte_Qwen2_7B_instruct = ModelMeta( - loader=partial( + loader=partial( # type: ignore instruct_wrapper, model_name_or_path="Alibaba-NLP/gte-Qwen2-7B-instruct", - instruction_template="Instruct: {instruction}\nQuery: ", + instruction_template=instruction_template, attn="cccc", pooling_method="lasttoken", mode="embedding", @@ -33,3 +39,59 @@ framework=["Sentence Transformers", "PyTorch"], use_instructions=True, ) + + +gte_Qwen1_5_7B_instruct = ModelMeta( + loader=partial( # type: ignore + instruct_wrapper, + model_name_or_path="Alibaba-NLP/gte-Qwen1.5-7B-instruct", + instruction_template=instruction_template, + attn="cccc", + pooling_method="lasttoken", + mode="embedding", + torch_dtype="auto", + normalized=True, + ), + name="Alibaba-NLP/gte-Qwen1.5-7B-instruct", + languages=["eng_Latn"], + open_weights=True, + revision="07d27e5226328010336563bc1b564a5e3436a298", + release_date="2024-04-20", # initial commit of hf model. + n_parameters=7_720_000_000, + memory_usage=None, + embed_dim=4096, + license="apache-2.0", + max_tokens=32768, + reference="https://huggingface.co/Alibaba-NLP/gte-Qwen1.5-7B-instruct", + similarity_fn_name="cosine", + framework=["Sentence Transformers", "PyTorch"], + use_instructions=True, +) + + +gte_Qwen2_1_5B_instruct = ModelMeta( + loader=partial( # type: ignore + instruct_wrapper, + model_name_or_path="Alibaba-NLP/gte-Qwen2-1.5B-instruct", + instruction_template=instruction_template, + attn="cccc", + pooling_method="lasttoken", + mode="embedding", + torch_dtype="auto", + normalized=True, + ), + name="Alibaba-NLP/gte-Qwen2-1.5B-instruct", + languages=["eng_Latn"], + open_weights=True, + revision="c6c1b92f4a3e1b92b326ad29dd3c8433457df8dd", + release_date="2024-07-29", # initial commit of hf model. + n_parameters=1_780_000_000, + memory_usage=None, + embed_dim=8960, + license="apache-2.0", + max_tokens=131072, + reference="https://huggingface.co/Alibaba-NLP/gte-Qwen2-1.5B-instruct", + similarity_fn_name="cosine", + framework=["Sentence Transformers", "PyTorch"], + use_instructions=True, +) diff --git a/mteb/models/jina_models.py b/mteb/models/jina_models.py index b103e174e..7705daec8 100644 --- a/mteb/models/jina_models.py +++ b/mteb/models/jina_models.py @@ -191,7 +191,7 @@ def encode( jina_embeddings_v3 = ModelMeta( - loader=partial( + loader=partial( # type: ignore JinaWrapper, model="jinaai/jina-embeddings-v3", revision="215a6e121fa0183376388ac6b1ae230326bfeaed", diff --git a/mteb/models/linq_models.py b/mteb/models/linq_models.py new file mode 100644 index 000000000..48e86ac8d --- /dev/null +++ b/mteb/models/linq_models.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +from functools import partial + +import torch + +from mteb.model_meta import ModelMeta +from mteb.models.instruct_wrapper import instruct_wrapper + + +def instruction_template(instruction: str) -> str: + return f"Instruct: {instruction}\nQuery: " if instruction else "" + + +Linq_Embed_Mistral = ModelMeta( + loader=partial( # type: ignore + instruct_wrapper, + model_name_or_path="Linq-AI-Research/Linq-Embed-Mistral", + instruction_template=instruction_template, + attn="cccc", + pooling_method="lasttoken", + mode="embedding", + torch_dtype=torch.bfloat16, + normalized=True, + ), + name="Linq-AI-Research/Linq-Embed-Mistral", + languages=["eng_Latn"], + open_weights=True, + revision="0c1a0b0589177079acc552433cad51d7c9132379", + release_date="2024-05-29", # initial commit of hf model. + n_parameters=7_110_000_000, + memory_usage=None, + embed_dim=4096, + license="cc-by-nc-4.0", + max_tokens=32768, + reference="https://huggingface.co/Linq-AI-Research/Linq-Embed-Mistral", + similarity_fn_name="cosine", + framework=["Sentence Transformers", "PyTorch"], + use_instructions=True, +) diff --git a/mteb/models/mxbai_models.py b/mteb/models/mxbai_models.py index ce7d1808b..5dfb9dc42 100644 --- a/mteb/models/mxbai_models.py +++ b/mteb/models/mxbai_models.py @@ -5,7 +5,7 @@ from mteb.model_meta import ModelMeta, sentence_transformers_loader mxbai_embed_large_v1 = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="mixedbread-ai/mxbai-embed-large-v1", revision="990580e27d329c7408b3741ecff85876e128e203", diff --git a/mteb/models/nvidia_models.py b/mteb/models/nvidia_models.py new file mode 100644 index 000000000..b3b11f559 --- /dev/null +++ b/mteb/models/nvidia_models.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +import logging +from collections.abc import Sequence +from functools import partial +from typing import Any + +import numpy as np +import torch +from sentence_transformers import CrossEncoder, SentenceTransformer + +from mteb.encoder_interface import PromptType +from mteb.model_meta import ModelMeta +from mteb.models.sentence_transformer_wrapper import SentenceTransformerWrapper + +logger = logging.getLogger(__name__) + + +def instruction_template(instruction: str) -> str: + return f"Instruct: {instruction}\nQuery: " if instruction else "" + + +class NvEmbedWrapper(SentenceTransformerWrapper): + def __init__( + self, + model: str | SentenceTransformer | CrossEncoder, + revision: str | None = None, + model_prompts: dict[str, str] | None = None, + **kwargs, + ) -> None: + super().__init__(model, revision, model_prompts, **kwargs) + self.model.max_seq_length = 32768 + self.model.tokenizer.padding_side = "right" + + def encode( + self, + sentences: Sequence[str], + *, + task_name: str, + prompt_type: PromptType | None = None, + **kwargs: Any, + ) -> np.ndarray: + # Add eos token to each input example + sentences = [example + self.model.tokenizer.eos_token for example in sentences] + + instruction = "" + if prompt_type == PromptType.query: + instruction = self.get_instruction(task_name, prompt_type) + + prompt = instruction_template(instruction) + + if prompt: + logger.info(f"Using {prompt=} for task={task_name} {prompt_type=}") + else: + logger.info(f"No model prompts found for task={task_name} {prompt_type=}") + + logger.info(f"Encoding {len(sentences)} sentences.") + + embeddings = self.model.encode( + sentences, + prompt=prompt, + normalize_embeddings=True, + **kwargs, + ) + if isinstance(embeddings, torch.Tensor): + embeddings = embeddings.cpu().detach().float().numpy() + return embeddings + + +NV_embed_v2 = ModelMeta( + loader=partial( # type: ignore + NvEmbedWrapper, + model="nvidia/NV-Embed-v2", + trust_remote_code=True, + ), + name="nvidia/NV-Embed-v2", + languages=["eng_Latn"], + open_weights=True, + revision="7604d305b621f14095a1aa23d351674c2859553a", + release_date="2024-09-09", # initial commit of hf model. + n_parameters=7_850_000_000, + memory_usage=None, + embed_dim=4096, + license="cc-by-nc-4.0", + max_tokens=32768, + reference="https://huggingface.co/nvidia/NV-Embed-v2", + similarity_fn_name="cosine", + framework=["Sentence Transformers", "PyTorch"], + use_instructions=True, +) + +NV_embed_v1 = ModelMeta( + loader=partial( # type: ignore + NvEmbedWrapper, + model="nvidia/NV-Embed-v1", + trust_remote_code=True, + ), + name="nvidia/NV-Embed-v1", + languages=["eng_Latn"], + open_weights=True, + revision="570834afd5fef5bf3a3c2311a2b6e0a66f6f4f2c", + release_date="2024-09-13", # initial commit of hf model. + n_parameters=7_850_000_000, + memory_usage=None, + embed_dim=4096, + license="cc-by-nc-4.0", + max_tokens=32768, + reference="https://huggingface.co/nvidia/NV-Embed-v1", + similarity_fn_name="cosine", + framework=["Sentence Transformers", "PyTorch"], + use_instructions=True, +) diff --git a/mteb/models/overview.py b/mteb/models/overview.py index 1db236969..8d98b9a02 100644 --- a/mteb/models/overview.py +++ b/mteb/models/overview.py @@ -18,9 +18,11 @@ gritlm_models, gte_models, jina_models, + linq_models, llm2vec_models, mxbai_models, nomic_models, + nvidia_models, openai_models, promptriever_models, repllama_models, @@ -43,24 +45,26 @@ e5_instruct, e5_models, google_models, + google_models, gritlm_models, gte_models, + jina_models, + linq_models, llm2vec_models, mxbai_models, nomic_models, + nvidia_models, openai_models, + promptriever_models, + repllama_models, + rerankers_custom, + rerankers_monot5_based, ru_sentence_models, salesforce_models, sentence_transformers_models, - voyage_models, - google_models, - repllama_models, - promptriever_models, - jina_models, - uae_models, stella_models, - rerankers_monot5_based, - rerankers_custom, + uae_models, + voyage_models, ] MODEL_REGISTRY = {} diff --git a/mteb/models/rerankers_custom.py b/mteb/models/rerankers_custom.py index dc354a550..40977f1e0 100644 --- a/mteb/models/rerankers_custom.py +++ b/mteb/models/rerankers_custom.py @@ -175,7 +175,7 @@ def loader_inner(**kwargs: Any) -> Encoder: monobert_large = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=MonoBERTReranker, model_name_or_path="castorini/monobert-large-msmarco", @@ -190,7 +190,7 @@ def loader_inner(**kwargs: Any) -> Encoder: # languages unclear: https://huggingface.co/jinaai/jina-reranker-v2-base-multilingual/discussions/28 jina_reranker_multilingual = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=JinaReranker, model_name_or_path="jinaai/jina-reranker-v2-base-multilingual", @@ -204,7 +204,7 @@ def loader_inner(**kwargs: Any) -> Encoder: ) bge_reranker_v2_m3 = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=BGEReranker, model_name_or_path="BAAI/bge-reranker-v2-m3", diff --git a/mteb/models/rerankers_monot5_based.py b/mteb/models/rerankers_monot5_based.py index aef4a19e7..b96897ee5 100644 --- a/mteb/models/rerankers_monot5_based.py +++ b/mteb/models/rerankers_monot5_based.py @@ -276,7 +276,7 @@ def get_prediction_tokens(self, *args, **kwargs): monot5_small = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=MonoT5Reranker, model_name_or_path="castorini/monot5-small-msmarco-10k", @@ -290,7 +290,7 @@ def get_prediction_tokens(self, *args, **kwargs): ) monot5_base = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=MonoT5Reranker, model_name_or_path="castorini/monot5-base-msmarco-10k", @@ -304,7 +304,7 @@ def get_prediction_tokens(self, *args, **kwargs): ) monot5_large = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=MonoT5Reranker, model_name_or_path="castorini/monot5-large-msmarco-10k", @@ -318,7 +318,7 @@ def get_prediction_tokens(self, *args, **kwargs): ) monot5_3b = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=MonoT5Reranker, model_name_or_path="castorini/monot5-3b-msmarco-10k", @@ -332,7 +332,7 @@ def get_prediction_tokens(self, *args, **kwargs): ) flant5_base = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=FLANT5Reranker, model_name_or_path="google/flan-t5-base", @@ -346,7 +346,7 @@ def get_prediction_tokens(self, *args, **kwargs): ) flant5_large = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=FLANT5Reranker, model_name_or_path="google/flan-t5-large", @@ -360,7 +360,7 @@ def get_prediction_tokens(self, *args, **kwargs): ) flant5_xl = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=FLANT5Reranker, model_name_or_path="google/flan-t5-xl", @@ -374,7 +374,7 @@ def get_prediction_tokens(self, *args, **kwargs): ) flant5_xxl = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=FLANT5Reranker, model_name_or_path="google/flan-t5-xxl", @@ -389,7 +389,7 @@ def get_prediction_tokens(self, *args, **kwargs): llama2_7b = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=LlamaReranker, model_name_or_path="meta-llama/Llama-2-7b-hf", @@ -403,7 +403,7 @@ def get_prediction_tokens(self, *args, **kwargs): ) llama2_7b_chat = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=LlamaReranker, model_name_or_path="meta-llama/Llama-2-7b-chat-hf", @@ -417,7 +417,7 @@ def get_prediction_tokens(self, *args, **kwargs): ) mistral_7b = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=MistralReranker, model_name_or_path="mistralai/Mistral-7B-Instruct-v0.2", @@ -431,7 +431,7 @@ def get_prediction_tokens(self, *args, **kwargs): ) followir_7b = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=FollowIRReranker, model_name_or_path="jhu-clsp/FollowIR-7B", @@ -550,7 +550,7 @@ def get_prediction_tokens(self, *args, **kwargs): ] mt5_base_mmarco_v2 = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=MonoT5Reranker, model_name_or_path="unicamp-dl/mt5-base-mmarco-v2", @@ -564,7 +564,7 @@ def get_prediction_tokens(self, *args, **kwargs): ) mt5_13b_mmarco_100k = ModelMeta( - loader=partial( + loader=partial( # type: ignore _loader, wrapper=MonoT5Reranker, model_name_or_path="unicamp-dl/mt5-13b-mmarco-100k", diff --git a/mteb/models/ru_sentence_models.py b/mteb/models/ru_sentence_models.py index cfe896516..301892e07 100644 --- a/mteb/models/ru_sentence_models.py +++ b/mteb/models/ru_sentence_models.py @@ -75,7 +75,7 @@ ) user_base_ru = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="deepvk/USER-base", revision="436a489a2087d61aa670b3496a9915f84e46c861", @@ -218,7 +218,7 @@ rosberta_ru_en = ModelMeta( - loader=partial( + loader=partial( # type: ignore sentence_transformers_loader, model_name="ai-forever/ru-en-RoSBERTa", revision="89fb1651989adbb1cfcfdedafd7d102951ad0555", diff --git a/mteb/models/salesforce_models.py b/mteb/models/salesforce_models.py index eabc4352a..5ef104f23 100644 --- a/mteb/models/salesforce_models.py +++ b/mteb/models/salesforce_models.py @@ -5,23 +5,22 @@ import torch from mteb.model_meta import ModelMeta +from mteb.models.instruct_wrapper import instruct_wrapper -from .instruct_wrapper import instruct_wrapper - -def sfr_instruction(instruction: str) -> str: - return f"Instruct: {instruction}\nQuery: " +def instruction_template(instruction: str) -> str: + return f"Instruct: {instruction}\nQuery: " if instruction else "" SFR_Embedding_2_R = ModelMeta( - loader=partial( + loader=partial( # type: ignore instruct_wrapper, model_name_or_path="Salesforce/SFR-Embedding-2_R", - instruction_template=sfr_instruction, + instruction_template=instruction_template, attn="cccc", pooling_method="lasttoken", mode="embedding", - torch_dtype=torch.bfloat16, + torch_dtype="auto", # The ST script does not normalize while the HF one does so unclear what to do # https://huggingface.co/Salesforce/SFR-Embedding-2_R normalized=True, @@ -41,3 +40,31 @@ def sfr_instruction(instruction: str) -> str: framework=["Sentence Transformers", "PyTorch"], use_instructions=True, ) + + +SFR_Embedding_Mistral = ModelMeta( + loader=partial( # type: ignore + instruct_wrapper, + model_name_or_path="Salesforce/SFR-Embedding-Mistral", + instruction_template=instruction_template, + attn="cccc", + pooling_method="lasttoken", + mode="embedding", + torch_dtype="auto", + normalized=True, + ), + name="Salesforce/SFR-Embedding-Mistral", + languages=["eng_Latn"], + open_weights=True, + revision="938c560d1c236aa563b2dbdf084f28ab28bccb11", + release_date="2024-01-24", # initial commit of hf model. + n_parameters=7_110_000_000, + memory_usage=None, + embed_dim=4096, + license="cc-by-nc-4.0", + max_tokens=32768, + reference="https://huggingface.co/Salesforce/SFR-Embedding-Mistral", + similarity_fn_name="cosine", + framework=["Sentence Transformers", "PyTorch"], + use_instructions=True, +) diff --git a/mteb/models/stella_models.py b/mteb/models/stella_models.py index 8fc19fd06..153ee6aa9 100644 --- a/mteb/models/stella_models.py +++ b/mteb/models/stella_models.py @@ -7,7 +7,7 @@ stella_en_400M = ModelMeta( # https://huggingface.co/dunzhang/stella_en_400M_v5/discussions/21#671a6205ac1e2416090f2bf4 - loader=partial( + loader=partial( # type: ignore instruct_wrapper, model_name_or_path="dunzhang/stella_en_400M_v5", attn="cccc", @@ -31,7 +31,7 @@ ) stella_en_1_5b = ModelMeta( - loader=partial( + loader=partial( # type: ignore instruct_wrapper, model_name_or_path="dunzhang/stella_en_1.5B_v5", attn="cccc", diff --git a/mteb/models/uae_models.py b/mteb/models/uae_models.py index b18240b47..942d50894 100644 --- a/mteb/models/uae_models.py +++ b/mteb/models/uae_models.py @@ -52,7 +52,7 @@ def encode( uae_large_v1 = ModelMeta( - loader=partial( + loader=partial( # type: ignore UAEWrapper, model="WhereIsAI/UAE-Large-V1", revision="369c368f70f16a613f19f5598d4f12d9f44235d4", diff --git a/mteb/models/voyage_models.py b/mteb/models/voyage_models.py index ea6b25bde..d63f132cf 100644 --- a/mteb/models/voyage_models.py +++ b/mteb/models/voyage_models.py @@ -145,7 +145,7 @@ def _batched_encode( revision="1", release_date="2024-05-05", languages=None, # supported languages not specified - loader=partial( + loader=partial( # type: ignore VoyageWrapper, model_name="voyage-large-2-instruct", model_prompts=model_prompts, @@ -167,7 +167,7 @@ def _batched_encode( revision="1", release_date="2024-05-30", languages=None, # supported languages not specified - loader=partial( + loader=partial( # type: ignore VoyageWrapper, model_name="voyage-finance-2", model_prompts=model_prompts, @@ -189,7 +189,7 @@ def _batched_encode( revision="1", release_date="2024-04-15", languages=None, # supported languages not specified - loader=partial( + loader=partial( # type: ignore VoyageWrapper, model_name="voyage-law-2", model_prompts=model_prompts, @@ -211,7 +211,7 @@ def _batched_encode( revision="1", release_date="2024-01-23", languages=None, # supported languages not specified - loader=partial( + loader=partial( # type: ignore VoyageWrapper, model_name="voyage-code-2", model_prompts=model_prompts, @@ -233,7 +233,7 @@ def _batched_encode( revision="1", release_date="2023-10-29", languages=None, # supported languages not specified - loader=partial( + loader=partial( # type: ignore VoyageWrapper, model_name="voyage-large-2", model_prompts=model_prompts, @@ -255,7 +255,7 @@ def _batched_encode( revision="1", release_date="2023-10-29", languages=None, # supported languages not specified - loader=partial( + loader=partial( # type: ignore VoyageWrapper, model_name="voyage-2", model_prompts=model_prompts, @@ -276,7 +276,7 @@ def _batched_encode( revision="1", release_date="2024-06-10", languages=None, # supported languages not specified - loader=partial( + loader=partial( # type: ignore VoyageWrapper, model_name="voyage-multilingual-2", model_prompts=model_prompts,