From 1c287c3c65b865c1b27a247b9c59bb67a8468429 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Fri, 12 Dec 2025 15:44:42 +0300 Subject: [PATCH 1/3] fix nv embed wrapper --- .../model_implementations/nvidia_models.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/mteb/models/model_implementations/nvidia_models.py b/mteb/models/model_implementations/nvidia_models.py index b7c232e791..fc45fa91cd 100644 --- a/mteb/models/model_implementations/nvidia_models.py +++ b/mteb/models/model_implementations/nvidia_models.py @@ -100,8 +100,24 @@ def instruction_template( "MrTidyRetrieval", } + +def _nvembed_v2_wrapper( + *args, + **kwargs: Any, +) -> InstructSentenceTransformerModel: + required_transformers_version = "4.42.4" + + if Version(transformers_version) != Version(required_transformers_version): + raise RuntimeError( + f"transformers version {transformers_version} is not match with required " + f"install version {required_transformers_version} to run `nvidia/NV-Embed-v2`" + ) + + return InstructSentenceTransformerModel(*args, **kwargs) + + NV_embed_v2 = ModelMeta( - loader=InstructSentenceTransformerModel, + loader=_nvembed_v2_wrapper, loader_kwargs=dict( instruction_template=instruction_template, trust_remote_code=True, From 217db9822885c62143bfc2fec187ddca30d28ee8 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 16 Dec 2025 17:13:50 +0300 Subject: [PATCH 2/3] try to fix --- .../model_implementations/nvidia_models.py | 82 +++++++++++++++---- 1 file changed, 67 insertions(+), 15 deletions(-) diff --git a/mteb/models/model_implementations/nvidia_models.py b/mteb/models/model_implementations/nvidia_models.py index fc45fa91cd..9918017508 100644 --- a/mteb/models/model_implementations/nvidia_models.py +++ b/mteb/models/model_implementations/nvidia_models.py @@ -1,4 +1,5 @@ import logging +from collections.abc import Callable from typing import Any import torch @@ -29,7 +30,7 @@ }""" -def instruction_template( +def _instruction_template( instruction: str, prompt_type: PromptType | None = None ) -> str: return f"Instruct: {instruction}\nQuery: " if instruction else "" @@ -101,25 +102,76 @@ def instruction_template( } -def _nvembed_v2_wrapper( - *args, - **kwargs: Any, -) -> InstructSentenceTransformerModel: - required_transformers_version = "4.42.4" +class _NVEmbedWrapper(InstructSentenceTransformerModel): + """Inherited, because nvembed requires `sbert==2`, but it doesn't have tokenizers kwargs""" - if Version(transformers_version) != Version(required_transformers_version): - raise RuntimeError( - f"transformers version {transformers_version} is not match with required " - f"install version {required_transformers_version} to run `nvidia/NV-Embed-v2`" + def __init__( + self, + model_name: str, + revision: str, + instruction_template: str + | Callable[[str, PromptType | None], str] + | None = None, + max_seq_length: int | None = None, + apply_instruction_to_passages: bool = True, + padding_side: str | None = None, + add_eos_token: bool = False, + prompts_dict: dict[str, str] | None = None, + **kwargs: Any, + ): + from sentence_transformers import __version__ as sbert_version + + required_transformers_version = "4.42.4" + required_sbert_version = "2.7.2" + + if Version(transformers_version) != Version(required_transformers_version): + raise RuntimeError( + f"transformers version {transformers_version} is not match with required " + f"install version {required_transformers_version} to run `nvidia/NV-Embed-v2`" + ) + + if Version(sbert_version) != Version(required_sbert_version): + raise RuntimeError( + f"sbert version {sbert_version} is not match with required " + f"install version {required_sbert_version} to run `nvidia/NV-Embed-v2`" + ) + + requires_package( + self, "flash_attn", model_name, "pip install 'mteb[flash_attention]'" ) - return InstructSentenceTransformerModel(*args, **kwargs) + from sentence_transformers import SentenceTransformer + + if ( + isinstance(instruction_template, str) + and "{instruction}" not in instruction_template + ): + raise ValueError( + "Instruction template must contain the string '{instruction}'." + ) + if instruction_template is None: + logger.warning( + "No instruction template provided. Instructions will be used as-is." + ) + + self.instruction_template = instruction_template + + self.model_name = model_name + self.model = SentenceTransformer(model_name, revision=revision, **kwargs) + self.model.tokenizer.padding_side = padding_side + self.model.tokenizer.add_eos_token = add_eos_token + + if max_seq_length: + # https://github.com/huggingface/sentence-transformers/issues/3575 + self.model.max_seq_length = max_seq_length + self.apply_instruction_to_passages = apply_instruction_to_passages + self.prompts_dict = prompts_dict NV_embed_v2 = ModelMeta( - loader=_nvembed_v2_wrapper, + loader=_NVEmbedWrapper, loader_kwargs=dict( - instruction_template=instruction_template, + instruction_template=_instruction_template, trust_remote_code=True, max_seq_length=32768, padding_side="right", @@ -147,9 +199,9 @@ def _nvembed_v2_wrapper( ) NV_embed_v1 = ModelMeta( - loader=InstructSentenceTransformerModel, + loader=_NVEmbedWrapper, loader_kwargs=dict( - instruction_template=instruction_template, + instruction_template=_instruction_template, trust_remote_code=True, max_seq_length=32768, padding_side="right", From 699a90bc45a285e7bf51339de4faf91e58d6b3d1 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Tue, 16 Dec 2025 17:26:36 +0300 Subject: [PATCH 3/3] fix sbert version --- mteb/models/model_implementations/nvidia_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mteb/models/model_implementations/nvidia_models.py b/mteb/models/model_implementations/nvidia_models.py index 9918017508..9df229958d 100644 --- a/mteb/models/model_implementations/nvidia_models.py +++ b/mteb/models/model_implementations/nvidia_models.py @@ -122,7 +122,7 @@ def __init__( from sentence_transformers import __version__ as sbert_version required_transformers_version = "4.42.4" - required_sbert_version = "2.7.2" + required_sbert_version = "2.7.0" if Version(transformers_version) != Version(required_transformers_version): raise RuntimeError(