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

mteb/models/model_implementations/slm_models.py

@@ -0,0 +1,416 @@
+"""
+SauerkrautLM Visual Document Retrieval Models - MTEB Integration
+
+This module provides MTEB wrappers for SauerkrautLM ColPali-style models:
+- SLM-ColQwen3 (Qwen3-VL backbone)
+- SLM-ColLFM2 (LFM2 backbone)
+- SLM-ColMinistral3 (Ministral3 backbone)
+
+Based on:
+- MTEB ColPali implementation: mteb/models/model_implementations/colpali_models.py
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import torch
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+
+from mteb._requires_package import (
+    requires_image_dependencies,
+    requires_package,
+)
+from mteb.abstasks.task_metadata import TaskMetadata
+from mteb.models.abs_encoder import AbsEncoder
+from mteb.models.model_implementations.colpali_models import (
+    COLPALI_CITATION,
+    COLPALI_TRAINING_DATA,
+)
+from mteb.models.model_meta import ModelMeta, ScoringFunction
+from mteb.types import Array, BatchedInput, PromptType
+
+logger = logging.getLogger(__name__)
+
+
+SUPPORTED_LANGUAGES = [
+    "eng-Latn",  # English
+    "deu-Latn",  # German
+    "fra-Latn",  # French
+    "spa-Latn",  # Spanish
+    "ita-Latn",  # Italian
+    "por-Latn",  # Portuguese
+]
+
+
+class SLMBaseWrapper(AbsEncoder):
+    """
+    Base wrapper for SauerkrautLM multi-vector embedding models.
+
+    All our models use late interaction (MaxSim) for retrieval scoring.
+    """
+
+    model_class = None
+    processor_class = None
+    model_name_prefix = "SLM"
+
+    def __init__(
+        self,
+        model_name: str,
+        revision: str | None = None,
+        device: str | None = None,
+        use_flash_attn: bool = True,
+        **kwargs,
+    ):
+        requires_image_dependencies()
+        requires_package(
+            self, "sauerkrautlm_colpali", model_name, "pip install sauerkrautlm-colpali"
+        )
+
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self._load_model_and_processor(model_name, revision, use_flash_attn, **kwargs)
+        self.mdl = self.mdl.to(self.device)
+        self.mdl.eval()
+
+    def _load_model_and_processor(self, model_name, revision, use_flash_attn, **kwargs):
+        """Override in subclasses to load specific model/processor."""
+        raise NotImplementedError
+
+    def encode(
+        self,
+        inputs: DataLoader[BatchedInput],
+        *,
+        task_metadata: TaskMetadata,
+        hf_split: str,
+        hf_subset: str,
+        prompt_type: PromptType | None = None,
+        **kwargs: Any,
+    ) -> Array:
+        text_embeddings = None
+        image_embeddings = None
+
+        if "text" in inputs.dataset.features:
+            text_embeddings = self.get_text_embeddings(inputs, **kwargs)
+        if "image" in inputs.dataset.features:
+            image_embeddings = self.get_image_embeddings(inputs, **kwargs)
+
+        if text_embeddings is not None and image_embeddings is not None:
+            if len(text_embeddings) != len(image_embeddings):
+                raise ValueError(
+                    "The number of texts and images must have the same length"
+                )
+            fused_embeddings = text_embeddings + image_embeddings
+            return fused_embeddings
+        elif text_embeddings is not None:
+            return text_embeddings
+        elif image_embeddings is not None:
+            return image_embeddings
+        raise ValueError("No text or image features found in inputs")
+
+    def encode_input(self, inputs):
+        """Forward pass through the model."""
+        return self.mdl(**inputs)
+
+    def _move_to_device(self, inputs: dict) -> dict:
+        """Move all tensor inputs to the model's device."""
+        result = {}
+        for k, v in inputs.items():
+            if isinstance(v, torch.Tensor):
+                result[k] = v.to(self.device)
+            else:
+                result[k] = v
+        return result
+
+    def get_image_embeddings(
+        self,
+        images: DataLoader,
+        batch_size: int = 32,
+        **kwargs,
+    ) -> torch.Tensor:
+        import torchvision.transforms.functional as F
+
+        all_embeds = []
+
+        with torch.no_grad():
+            for batch in tqdm(images, desc="Encoding images"):
+                from PIL import Image
+
+                imgs = [
+                    F.to_pil_image(b) if not isinstance(b, Image.Image) else b
+                    for b in batch["image"]
+                ]
+                inputs = self.processor.process_images(imgs)
+                inputs = self._move_to_device(inputs)
+                outs = self.encode_input(inputs)
+                all_embeds.extend(outs.cpu().to(torch.float32))
+
+        padded = torch.nn.utils.rnn.pad_sequence(
+            all_embeds, batch_first=True, padding_value=0
+        )
+        return padded
+
+    def get_text_embeddings(
+        self,
+        texts: DataLoader,
+        batch_size: int = 32,
+        **kwargs,
+    ) -> torch.Tensor:
+        all_embeds = []
+
+        with torch.no_grad():
+            for batch in tqdm(texts, desc="Encoding texts"):
+                inputs = self.processor.process_queries(batch["text"])
+                inputs = self._move_to_device(inputs)
+                outs = self.encode_input(inputs)
+                all_embeds.extend(outs.cpu().to(torch.float32))
+
+        padded = torch.nn.utils.rnn.pad_sequence(
+            all_embeds, batch_first=True, padding_value=0
+        )
+        return padded
+
+    def calculate_probs(
+        self,
+        text_embeddings: torch.Tensor,
+        image_embeddings: torch.Tensor,
+    ) -> torch.Tensor:
+        scores = self.similarity(text_embeddings, image_embeddings).T
+        return scores.softmax(dim=-1)
+
+    def similarity(
+        self,
+        a: torch.Tensor | list,
+        b: torch.Tensor | list,
+    ) -> torch.Tensor:
+        return self.processor.score(a, b, device=self.device)
+
+
+class SLMColQwen3Wrapper(SLMBaseWrapper):
+    """Wrapper for SLM-ColQwen3 models (Qwen3-VL backbone)."""
+
+    def _load_model_and_processor(self, model_name, revision, use_flash_attn, **kwargs):
+        from sauerkrautlm_colpali.models.qwen3.colqwen3 import (
+            ColQwen3,
+            ColQwen3Processor,
+        )
+
+        self.mdl = ColQwen3.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16,
+            attn_implementation="flash_attention_2" if use_flash_attn else "eager",
+            revision=revision,
+            **kwargs,
+        )
+
+        self.processor = ColQwen3Processor.from_pretrained(
+            model_name,
+            revision=revision,
+        )
+
+        logger.info(f"SLM-ColQwen3 loaded: dim={self.mdl.dim}, device={self.device}")
+
+
+class SLMColLFM2Wrapper(SLMBaseWrapper):
+    """Wrapper for SLM-ColLFM2 models (LFM2 backbone)."""
+
+    def _load_model_and_processor(self, model_name, revision, use_flash_attn, **kwargs):
+        from sauerkrautlm_colpali.models.lfm2.collfm2 import ColLFM2, ColLFM2Processor
+
+        self.mdl = ColLFM2.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16,
+            revision=revision,
+            **kwargs,
+        )
+
+        self.processor = ColLFM2Processor.from_pretrained(
+            model_name,
+            revision=revision,
+        )
+
+        logger.info(f"SLM-ColLFM2 loaded: dim={self.mdl.dim}, device={self.device}")
+
+
+class SLMColMinistral3Wrapper(SLMBaseWrapper):
+    """Wrapper for SLM-ColMinistral3 models (Ministral3 backbone)."""
+
+    def _load_model_and_processor(self, model_name, revision, use_flash_attn, **kwargs):
+        from sauerkrautlm_colpali.models.ministral3.colministral3 import (
+            ColMinistral3,
+            ColMinistral3Processor,
+        )
+
+        self.mdl = ColMinistral3.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16,
+        )
+
+        self.processor = ColMinistral3Processor.from_pretrained(model_name)
+
+        logger.info(
+            f"SLM-ColMinistral3 loaded: dim={self.mdl.dim}, device={self.device}"
+        )
+
+
+SAUERKRAUTLM_CITATION = """
+@misc{sauerkrautlm-colpali-2025,
+  title={SauerkrautLM-ColPali: Multi-Vector Vision Retrieval Models},
+  author={David Golchinfar},
+  organization={VAGO Solutions},
+  year={2025},
+  url={https://github.com/VAGOsolutions/sauerkrautlm-colpali}
+}
+"""
+
+
+slm_colqwen3_1_7b_turbo = ModelMeta(
+    loader=SLMColQwen3Wrapper,
+    name="VAGOsolutions/SauerkrautLM-ColQwen3-1.7b-Turbo-v0.1",
+    languages=SUPPORTED_LANGUAGES,
+    revision="19c295a18e057d6d82754f627c09408117ffdb66",
+    release_date="2025-12-20",
+    modalities=["image", "text"],
+    model_type=["late-interaction"],
+    n_parameters=1_756_572_288,
+    memory_usage_mb=3350,
+    max_tokens=262144,
+    embed_dim=128,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["ColPali"],
+    reference="https://huggingface.co/VAGOsolutions/SauerkrautLM-ColQwen3-1.7b-Turbo-v0.1",
+    similarity_fn_name=ScoringFunction.MAX_SIM,
+    use_instructions=True,
+    adapted_from="Qwen/Qwen3-VL-2B-Instruct",
+    training_datasets=COLPALI_TRAINING_DATA,
+    citation=SAUERKRAUTLM_CITATION + COLPALI_CITATION,
+)
+
+slm_colqwen3_2b = ModelMeta(
+    loader=SLMColQwen3Wrapper,
+    name="VAGOsolutions/SauerkrautLM-ColQwen3-2b-v0.1",
+    languages=SUPPORTED_LANGUAGES,
+    revision="48f699713c10af754684e12060a2af9266462cc9",
+    release_date="2025-12-20",
+    modalities=["image", "text"],
+    model_type=["late-interaction"],
+    n_parameters=2_127_794_304,
+    memory_usage_mb=4058,
+    max_tokens=262144,
+    embed_dim=128,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["ColPali"],
+    reference="https://huggingface.co/VAGOsolutions/SauerkrautLM-ColQwen3-2b-v0.1",
+    similarity_fn_name=ScoringFunction.MAX_SIM,
+    use_instructions=True,
+    adapted_from="Qwen/Qwen3-VL-2B-Instruct",
+    training_datasets=COLPALI_TRAINING_DATA,
+    citation=SAUERKRAUTLM_CITATION + COLPALI_CITATION,
+)
+
+slm_colqwen3_4b = ModelMeta(
+    loader=SLMColQwen3Wrapper,
+    name="VAGOsolutions/SauerkrautLM-ColQwen3-4b-v0.1",
+    languages=SUPPORTED_LANGUAGES,
+    revision="b635fbb3ab145f07608ed10a85def33544de1723",
+    release_date="2025-12-20",
+    modalities=["image", "text"],
+    model_type=["late-interaction"],
+    n_parameters=4_438_143_616,
+    memory_usage_mb=8465,
+    max_tokens=262144,
+    embed_dim=128,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["ColPali"],
+    reference="https://huggingface.co/VAGOsolutions/SauerkrautLM-ColQwen3-4b-v0.1",
+    similarity_fn_name=ScoringFunction.MAX_SIM,
+    use_instructions=True,
+    adapted_from="Qwen/Qwen3-VL-4B-Instruct",
+    training_datasets=COLPALI_TRAINING_DATA,
+    citation=SAUERKRAUTLM_CITATION + COLPALI_CITATION,
+)
+
+slm_colqwen3_8b = ModelMeta(
+    loader=SLMColQwen3Wrapper,
+    name="VAGOsolutions/SauerkrautLM-ColQwen3-8b-v0.1",
+    languages=SUPPORTED_LANGUAGES,
+    revision="36ac136e451a7b8d8229725d69d4ec23aa4f03c8",
+    release_date="2025-12-20",
+    modalities=["image", "text"],
+    model_type=["late-interaction"],
+    n_parameters=8_145_318_256,
+    memory_usage_mb=15536,
+    max_tokens=262144,
+    embed_dim=128,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["ColPali"],
+    reference="https://huggingface.co/VAGOsolutions/SauerkrautLM-ColQwen3-8b-v0.1",
+    similarity_fn_name=ScoringFunction.MAX_SIM,
+    use_instructions=True,
+    adapted_from="Qwen/Qwen3-VL-8B-Instruct",
+    training_datasets=COLPALI_TRAINING_DATA,
+    citation=SAUERKRAUTLM_CITATION + COLPALI_CITATION,
+)
+
+slm_collfm2_450m = ModelMeta(
+    loader=SLMColLFM2Wrapper,
+    name="VAGOsolutions/SauerkrautLM-ColLFM2-450M-v0.1",
+    languages=SUPPORTED_LANGUAGES,
+    revision="a65223fd6633f331ccff4483e47575c3c620dc60",
+    release_date="2025-12-20",
+    modalities=["image", "text"],
+    model_type=["late-interaction"],
+    n_parameters=450_953_856,
+    memory_usage_mb=860,
+    max_tokens=32768,
+    embed_dim=128,
+    license="https://huggingface.co/LiquidAI/LFM2-VL-450M/blob/main/LICENSE",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["ColPali"],
+    reference="https://huggingface.co/VAGOsolutions/SauerkrautLM-ColLFM2-450M-v0.1",
+    similarity_fn_name=ScoringFunction.MAX_SIM,
+    use_instructions=True,
+    adapted_from="LiquidAI/LFM2-VL-450M",
+    training_datasets=COLPALI_TRAINING_DATA,
+    citation=SAUERKRAUTLM_CITATION + COLPALI_CITATION,
+)
+
+slm_colministral3_3b = ModelMeta(
+    loader=SLMColMinistral3Wrapper,
+    name="VAGOsolutions/SauerkrautLM-ColMinistral3-3b-v0.1",
+    languages=SUPPORTED_LANGUAGES,
+    revision="54aa3ffbbce20471fdcc4afc07d13989c65e71b8",
+    release_date="2025-12-20",
+    modalities=["image", "text"],
+    model_type=["late-interaction"],
+    n_parameters=4_252_136_448,
+    memory_usage_mb=8110,
+    max_tokens=262144,
+    embed_dim=128,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["ColPali"],
+    reference="https://huggingface.co/VAGOsolutions/SauerkrautLM-ColMinistral3-3b-v0.1",
+    similarity_fn_name=ScoringFunction.MAX_SIM,
+    use_instructions=True,
+    adapted_from="mistralai/Ministral-3B-Instruct-2410",
+    training_datasets=COLPALI_TRAINING_DATA,
+    citation=SAUERKRAUTLM_CITATION + COLPALI_CITATION,
+)

mteb/models/model_implementations/spartan8806_atles_champion.py

@@ -17,7 +17,7 @@ spartan8806_atles_champion_embedding = ModelMeta(
     embed_dim=768,
     license="apache-2.0",
     similarity_fn_name="cosine",
-    framework=["Sentence Transformers"],
+    framework=["Sentence Transformers", "safetensors", "Transformers"],
     reference="https://huggingface.co/spartan8806/atles-champion-embedding",
     use_instructions=False,
     training_datasets={"STSBenchmark"},

mteb/models/model_implementations/stella_models.py

@@ -71,7 +71,13 @@ stella_en_400m = ModelMeta(
     embed_dim=4096,
     license="mit",
     similarity_fn_name=ScoringFunction.COSINE,
-    framework=["Sentence Transformers", "PyTorch", "GritLM"],
+    framework=[
+        "Sentence Transformers",
+        "PyTorch",
+        "GritLM",
+        "safetensors",
+        "Transformers",
+    ],
     reference="https://huggingface.co/NovaSearch/stella_en_400M_v5",
     training_datasets=nvidia_training_datasets,  # also distilled from gte-qwen (but training data is unknown) #2164
     public_training_code="https://github.com/NovaSearch-Team/RAG-Retrieval/blob/c40f4638b705eb77d88305d2056901ed550f9f4b/rag_retrieval/train/embedding/README.md",
@@ -100,7 +106,14 @@ stella_en_1_5b = ModelMeta(
     embed_dim=8960,
     license="mit",
     similarity_fn_name=ScoringFunction.COSINE,
-    framework=["Sentence Transformers", "PyTorch", "GritLM"],
+    framework=[
+        "Sentence Transformers",
+        "PyTorch",
+        "GritLM",
+        "ONNX",
+        "safetensors",
+        "Transformers",
+    ],
     reference="https://huggingface.co/NovaSearch/stella_en_1.5B_v5",
     training_datasets=nvidia_training_datasets,  # also distilled from gte-qwen (but training data is unknown) #2164
     public_training_code="https://github.com/NovaSearch-Team/RAG-Retrieval/blob/c40f4638b705eb77d88305d2056901ed550f9f4b/rag_retrieval/train/embedding/README.md",
@@ -178,7 +191,7 @@ stella_mrl_large_zh_v3_5_1792d = ModelMeta(
     max_tokens=512,
     reference="https://huggingface.co/dunzhang/stella-large-zh-v3-1792d",
     similarity_fn_name=ScoringFunction.COSINE,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "safetensors"],
     use_instructions=False,
     superseded_by=None,
     adapted_from="dunzhang/stella-large-zh-v3-1792d",
@@ -202,7 +215,7 @@ zpoint_large_embedding_zh = ModelMeta(
     max_tokens=512,
     reference="https://huggingface.co/iampanda/zpoint_large_embedding_zh",
     similarity_fn_name=ScoringFunction.COSINE,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "safetensors"],
     use_instructions=False,
     superseded_by=None,
     adapted_from="dunzhang/stella-mrl-large-zh-v3.5-1792d",

mteb/models/model_implementations/tarka_models.py

@@ -331,7 +331,7 @@ tarka_embedding_150m_v1 = ModelMeta(
     max_tokens=2048,
     license="gemma",
     reference="https://huggingface.co/Tarka-AIR/Tarka-Embedding-150M-V1",
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "safetensors"],
     use_instructions=True,
     public_training_code=None,
     public_training_data=None,
@@ -367,7 +367,7 @@ tarka_embedding_350m_v1 = ModelMeta(
     license=None,
     reference="https://huggingface.co/Tarka-AIR/Tarka-Embedding-350M-V1",
     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_implementations/text2vec_models.py

@@ -27,7 +27,7 @@ text2vec_base_chinese = ModelMeta(
     max_tokens=512,
     reference="https://huggingface.co/shibing624/text2vec-base-chinese",
     similarity_fn_name=ScoringFunction.COSINE,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "ONNX", "safetensors"],
     use_instructions=False,
     superseded_by=None,
     adapted_from=None,
@@ -57,7 +57,7 @@ text2vec_base_chinese_paraphrase = ModelMeta(
     max_tokens=512,
     reference="https://huggingface.co/shibing624/text2vec-base-chinese-paraphrase",
     similarity_fn_name=ScoringFunction.COSINE,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "safetensors", "Transformers"],
     use_instructions=False,
     superseded_by=None,
     adapted_from=None,
@@ -101,7 +101,13 @@ text2vec_base_multilingual = ModelMeta(
     max_tokens=256,
     reference="https://huggingface.co/shibing624/text2vec-base-chinese-paraphrase",
     similarity_fn_name=ScoringFunction.COSINE,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=[
+        "Sentence Transformers",
+        "PyTorch",
+        "ONNX",
+        "safetensors",
+        "Transformers",
+    ],
     use_instructions=False,
     superseded_by=None,
     adapted_from="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",

mteb/models/model_implementations/ua_sentence_models.py

@@ -17,7 +17,7 @@ xlm_roberta_ua_distilled = ModelMeta(
     open_weights=True,
     public_training_code="https://github.com/panalexeu/xlm-roberta-ua-distilled/blob/main/researches/research_final.ipynb",
     similarity_fn_name="cosine",
-    framework=["Sentence Transformers"],
+    framework=["Sentence Transformers", "safetensors"],
     reference="https://github.com/panalexeu/xlm-roberta-ua-distilled/tree/main",
     languages=["eng-Latn", "ukr-Cyrl"],
     training_datasets=set(

mteb/models/model_implementations/uae_models.py

@@ -72,7 +72,13 @@ uae_large_v1 = ModelMeta(
     embed_dim=1024,
     license="mit",
     similarity_fn_name=ScoringFunction.COSINE,
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=[
+        "Sentence Transformers",
+        "PyTorch",
+        "ONNX",
+        "safetensors",
+        "Transformers",
+    ],
     reference="https://huggingface.co/WhereIsAI/UAE-Large-V1",
     use_instructions=True,
     citation="""

mteb/models/model_implementations/vdr_models.py

@@ -38,7 +38,7 @@ vdr_2b_multi_v1 = ModelMeta(
     license="apache-2.0",
     reference="https://huggingface.co/llamaindex/vdr-2b-multi-v1",
     similarity_fn_name=ScoringFunction.COSINE,
-    framework=["PyTorch", "Sentence Transformers"],
+    framework=["PyTorch", "Sentence Transformers", "safetensors", "Transformers"],
     use_instructions=True,
     public_training_code=None,
     public_training_data="https://huggingface.co/datasets/llamaindex/vdr-multilingual-train",

mteb/models/model_implementations/vi_vn_models.py

@@ -22,7 +22,7 @@ greennode_embedding_large_vn_v1 = ModelMeta(
     max_tokens=8194,
     reference="https://huggingface.co/GreenNode/GreenNode-Embedding-Large-VN-V1",
     similarity_fn_name="cosine",
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "safetensors"],
     use_instructions=False,
     public_training_code=None,
     public_training_data="https://huggingface.co/datasets/GreenNode/GreenNode-Table-Markdown-Retrieval-VN",
@@ -47,7 +47,7 @@ greennode_embedding_large_vn_mixed_v1 = ModelMeta(
     max_tokens=8194,
     reference="https://huggingface.co/GreenNode/GreenNode-Embedding-Large-VN-Mixed-V1",
     similarity_fn_name="cosine",
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "safetensors"],
     use_instructions=False,
     public_training_code=None,
     public_training_data="https://huggingface.co/datasets/GreenNode/GreenNode-Table-Markdown-Retrieval-VN",
@@ -72,7 +72,7 @@ aiteamvn_vietnamese_embeddings = ModelMeta(
     max_tokens=8194,
     reference="https://huggingface.co/AITeamVN/Vietnamese_Embedding",
     similarity_fn_name="cosine",
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "ONNX", "safetensors"],
     use_instructions=False,
     public_training_code=None,
     public_training_data=None,
@@ -104,7 +104,7 @@ hiieu_halong_embedding = ModelMeta(
     max_tokens=514,
     reference="https://huggingface.co/hiieu/halong_embedding",
     similarity_fn_name="cosine",
-    framework=["Sentence Transformers", "PyTorch"],
+    framework=["Sentence Transformers", "PyTorch", "safetensors"],
     public_training_code=None,
     public_training_data=None,
     training_datasets=None,
@@ -135,7 +135,7 @@ sup_simcse_vietnamese_phobert_base_ = ModelMeta(
     license="apache-2.0",
     public_training_code=None,
     public_training_data=None,
-    framework=["PyTorch", "Sentence Transformers"],
+    framework=["PyTorch", "Sentence Transformers", "Transformers", "safetensors"],
     reference="https://huggingface.co/VoVanPhuc/sup-SimCSE-VietNamese-phobert-base",
     similarity_fn_name="cosine",
     training_datasets=None,
@@ -173,7 +173,7 @@ bkai_foundation_models_vietnamese_bi_encoder = ModelMeta(
     license="apache-2.0",
     public_training_code=None,
     public_training_data=None,
-    framework=["PyTorch", "Sentence Transformers"],
+    framework=["PyTorch", "Sentence Transformers", "safetensors", "Transformers"],
     reference="https://huggingface.co/bkai-foundation-models/vietnamese-bi-encoder",
     similarity_fn_name="cosine",
     training_datasets=None,

mteb/models/model_implementations/vlm2vec_models.py

@@ -41,7 +41,7 @@ class VLM2VecWrapper(AbsEncoder):
             model_name,
             "pip install flash-attn --no-build-isolation",
         ):
-            import flash_attn  # noqa
+            pass
 
         requires_package(self, "peft", model_name, "pip install 'mteb[peft]'")
         from peft import LoraConfig, PeftModel
@@ -282,7 +282,7 @@ vlm2vec_lora = ModelMeta(
     open_weights=True,
     public_training_code="https://github.com/TIGER-AI-Lab/VLM2Vec",
     public_training_data="https://huggingface.co/datasets/TIGER-Lab/MMEB-train",
-    framework=["PyTorch"],
+    framework=["PyTorch", "Transformers"],
     reference="https://huggingface.co/TIGER-Lab/VLM2Vec-LoRA",
     similarity_fn_name=ScoringFunction.COSINE,
     use_instructions=True,
@@ -306,7 +306,7 @@ vlm2vec_full = ModelMeta(
     open_weights=True,
     public_training_code="https://github.com/TIGER-AI-Lab/VLM2Vec",
     public_training_data="https://huggingface.co/TIGER-Lab/VLM2Vec-Full",
-    framework=["PyTorch"],
+    framework=["PyTorch", "Transformers", "safetensors"],
     reference="https://huggingface.co/TIGER-Lab/VLM2Vec-Full",
     similarity_fn_name=ScoringFunction.COSINE,
     use_instructions=True,