mteb 2.3.0__py3-none-any.whl → 2.3.1__py3-none-any.whl

mteb/models/model_implementations/colqwen_models.py

@@ -1,11 +1,19 @@
 import logging
+from typing import Any
 
 import torch
+from PIL import Image
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
 
 from mteb._requires_package import (
+    requires_image_dependencies,
     requires_package,
 )
-from mteb.models.model_meta import ModelMeta
+from mteb.abstasks.task_metadata import TaskMetadata
+from mteb.models.abs_encoder import AbsEncoder
+from mteb.models.model_meta import ModelMeta, ScoringFunction
+from mteb.types import Array, BatchedInput, PromptType
 
 from .colpali_models import (
     COLPALI_CITATION,
@@ -73,6 +81,132 @@ class ColQwen2_5Wrapper(ColPaliEngineWrapper):  # noqa: N801
         )
 
 
+class ColQwen3Wrapper(AbsEncoder):
+    """Wrapper for the ColQwen3 vision-language retrieval model."""
+
+    def __init__(
+        self,
+        model_name: str,
+        *,
+        revision: str | None = None,
+        device: str | None = None,
+        dtype: torch.dtype | str | None = torch.bfloat16,
+        **kwargs: Any,
+    ):
+        requires_image_dependencies()
+        requires_package(self, "transformers", model_name, "pip install mteb[colqwen3]")
+        from transformers import AutoModel, AutoProcessor
+
+        self.device = device or (
+            "cuda"
+            if torch.cuda.is_available()
+            else "mps"
+            if torch.backends.mps.is_available()
+            else "cpu"
+        )
+        self.model = AutoModel.from_pretrained(
+            model_name,
+            revision=revision,
+            dtype=dtype,
+            trust_remote_code=True,
+            **kwargs,
+        ).to(self.device)
+        self.model.eval()
+
+        self.processor = AutoProcessor.from_pretrained(
+            model_name,
+            revision=revision,
+            trust_remote_code=True,
+            max_num_visual_tokens=1280,
+        )
+
+    def encode(
+        self,
+        inputs: DataLoader[BatchedInput],
+        *,
+        task_metadata: TaskMetadata,
+        hf_split: str,
+        hf_subset: str,
+        prompt_type: PromptType | None = None,
+        **kwargs: Any,
+    ) -> Array:
+        if (
+            "text" not in inputs.dataset.features
+            and "image" not in inputs.dataset.features
+        ):
+            raise ValueError("No text or image features found in inputs.")
+        return self.get_fused_embeddings(inputs, **kwargs)
+
+    def _encode_inputs(self, encoded_inputs: dict[str, torch.Tensor]) -> torch.Tensor:
+        outputs = self.model(**encoded_inputs)
+        # Avoid boolean casting of tensors when checking for custom attributes.
+        embeddings = getattr(outputs, "embeddings", None)
+        if embeddings is None:
+            embeddings = outputs[0]
+        return embeddings
+
+    def get_fused_embeddings(
+        self,
+        image_texts_pairs: DataLoader[BatchedInput] | None = None,
+        batch_size: int = 32,
+        show_progress_bar: bool = True,
+        fusion_mode="concat",
+        **kwargs: Any,
+    ):
+        import torchvision.transforms.functional as F
+
+        contains_image = "image" in image_texts_pairs.dataset.features
+        contains_text = "text" in image_texts_pairs.dataset.features
+        contains_both = contains_image and contains_text
+
+        if contains_both:
+            progress_desc = "Encoding images+texts"
+        elif contains_image:
+            progress_desc = "Encoding images"
+        elif contains_text:
+            progress_desc = "Encoding texts"
+        else:
+            raise ValueError("No text or image features found in inputs.")
+
+        all_embeds: list[torch.Tensor] = []
+        with torch.no_grad():
+            for batch in tqdm(
+                image_texts_pairs,
+                disable=not show_progress_bar,
+                desc=progress_desc,
+            ):
+                if contains_image:
+                    imgs = [
+                        F.to_pil_image(b.to(self.device))
+                        if not isinstance(b, Image.Image)
+                        else b
+                        for b in batch["image"]
+                    ]
+                else:
+                    imgs = None
+                if contains_text:
+                    texts = batch["text"]
+                else:
+                    texts = None
+                if contains_both:
+                    assert len(imgs) == len(texts), (
+                        f"The number of texts and images must have the same length, got {len(imgs)} and {len(texts)}"
+                    )
+
+                inputs = self.processor(images=imgs, text=texts)
+                inputs = {k: v.to(self.device) for k, v in inputs.items()}
+                outs = self._encode_inputs(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 similarity(self, a, b):
+        return self.processor.score_multi_vector(a, b, device=self.device)
+
+
 colqwen2 = ModelMeta(
     loader=ColQwen2Wrapper,
     loader_kwargs=dict(
@@ -125,6 +259,72 @@ colqwen2_5 = ModelMeta(
     citation=COLPALI_CITATION,
 )
 
+TOMORO_TRAINING_DATA = {
+    "VDRMultilingualRetrieval",
+    # from https://huggingface.co/datasets/vidore/colpali_train_set
+    "VidoreDocVQARetrieval",
+    "VidoreInfoVQARetrieval",
+    "VidoreTatdqaRetrieval",
+    "VidoreArxivQARetrieval",
+    "VisRAG-Ret-Train-Synthetic-data",
+    "VisRAG-Ret-Train-In-domain-data",
+}
+
+TOMORO_CITATION = """
+@misc{huang2025tomoro_colqwen3_embed,
+  title={TomoroAI/tomoro-colqwen3-embed},
+  author={Xin Huang and Kye Min Tan and Albert Phelps},
+  year={2025},
+  url={https://huggingface.co/TomoroAI/tomoro-colqwen3-embed-8b}
+}
+"""
+
+colqwen3_8b = ModelMeta(
+    loader=ColQwen3Wrapper,
+    name="TomoroAI/tomoro-colqwen3-embed-8b",
+    languages=["eng-Latn"],
+    revision="0b9fe28142910e209bbac15b1efe85507c27644f",
+    release_date="2025-11-26",
+    modalities=["image", "text"],
+    n_parameters=8_000_000_000,
+    memory_usage_mb=16724,
+    max_tokens=262144,
+    embed_dim=320,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code="https://github.com/illuin-tech/colpali",
+    public_training_data=None,
+    framework=["PyTorch"],
+    reference="https://huggingface.co/TomoroAI/tomoro-colqwen3-embed-8b",
+    similarity_fn_name=ScoringFunction.MAX_SIM,
+    use_instructions=True,
+    training_datasets=TOMORO_TRAINING_DATA,
+    citation=TOMORO_CITATION,
+)
+
+colqwen3_4b = ModelMeta(
+    loader=ColQwen3Wrapper,
+    name="TomoroAI/tomoro-colqwen3-embed-4b",
+    languages=["eng-Latn"],
+    revision="6a32fb68598730bf5620fbf18d832c784235c59c",
+    release_date="2025-11-26",
+    modalities=["image", "text"],
+    n_parameters=4_000_000_000,
+    memory_usage_mb=8466,
+    max_tokens=262144,
+    embed_dim=320,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code="https://github.com/illuin-tech/colpali",
+    public_training_data=None,
+    framework=["PyTorch"],
+    reference="https://huggingface.co/TomoroAI/tomoro-colqwen3-embed-4b",
+    similarity_fn_name=ScoringFunction.MAX_SIM,
+    use_instructions=True,
+    training_datasets=TOMORO_TRAINING_DATA,
+    citation=TOMORO_CITATION,
+)
+
 colnomic_7b = ModelMeta(
     loader=ColQwen2_5Wrapper,
     loader_kwargs=dict(
@@ -223,7 +423,11 @@ colnomic_7b = ModelMeta(
 
 
 EVOQWEN_TRAINING_DATA = {
-    "colpali_train_set",
+    # "colpali_train_set",
+    "VidoreDocVQARetrieval",
+    "VidoreInfoVQARetrieval",
+    "VidoreTatdqaRetrieval",
+    "VidoreArxivQARetrieval",
     "VisRAG-Ret-Train-Synthetic-data",
     "VisRAG-Ret-Train-In-domain-data",
 }

mteb/models/model_implementations/euler_models.py

@@ -0,0 +1,25 @@
+from mteb.models.model_meta import ModelMeta
+from mteb.models.sentence_transformer_wrapper import sentence_transformers_loader
+
+Euler_Legal_Embedding_V1 = ModelMeta(
+    loader=sentence_transformers_loader,
+    name="Mira190/Euler-Legal-Embedding-V1",
+    revision="df607ed9e25e569514a99c27cdaaab16e76b6dd4",
+    release_date="2025-11-06",
+    languages=["eng-Latn"],
+    n_parameters=8000000000,
+    memory_usage_mb=15618,
+    max_tokens=1536,
+    embed_dim=4096,
+    license="apache-2.0",
+    open_weights=True,
+    public_training_code=None,
+    public_training_data=None,
+    framework=["PyTorch", "Sentence Transformers"],
+    reference="https://huggingface.co/Mira190/Euler-Legal-Embedding-V1",
+    similarity_fn_name="cosine",
+    use_instructions=False,
+    training_datasets=set(),  # final-data-new-anonymized-grok4-filtered
+    adapted_from="Qwen/Qwen3-Embedding-8B",
+    superseded_by=None,
+)

mteb/models/model_implementations/jina_models.py

@@ -1,8 +1,10 @@
 import logging
+from collections import defaultdict
 from typing import Any, ClassVar
 
 import numpy as np
 import torch
+from sentence_transformers import CrossEncoder
 from torch.utils.data import DataLoader
 
 from mteb._requires_package import requires_package
@@ -10,13 +12,92 @@ from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.languages import PROGRAMMING_LANGS
 from mteb.models.abs_encoder import AbsEncoder
 from mteb.models.model_meta import ModelMeta, ScoringFunction
-from mteb.models.sentence_transformer_wrapper import SentenceTransformerEncoderWrapper
+from mteb.models.sentence_transformer_wrapper import (
+    CrossEncoderWrapper,
+    SentenceTransformerEncoderWrapper,
+)
 from mteb.types import Array, BatchedInput, PromptType
 
 logger = logging.getLogger(__name__)
 
 MIN_SENTENCE_TRANSFORMERS_VERSION = (3, 1, 0)
 
+multilingual_langs = [
+    "afr-Latn",
+    "ara-Arab",
+    "aze-Latn",
+    "bel-Cyrl",
+    "bul-Cyrl",
+    "ben-Beng",
+    "cat-Latn",
+    "ceb-Latn",
+    "ces-Latn",
+    "cym-Latn",
+    "dan-Latn",
+    "deu-Latn",
+    "ell-Grek",
+    "eng-Latn",
+    "spa-Latn",
+    "est-Latn",
+    "eus-Latn",
+    "fas-Arab",
+    "fin-Latn",
+    "fra-Latn",
+    "glg-Latn",
+    "guj-Gujr",
+    "heb-Hebr",
+    "hin-Deva",
+    "hrv-Latn",
+    "hat-Latn",
+    "hun-Latn",
+    "hye-Armn",
+    "ind-Latn",
+    "isl-Latn",
+    "ita-Latn",
+    "jpn-Jpan",
+    "jav-Latn",
+    "kat-Geor",
+    "kaz-Cyrl",
+    "khm-Khmr",
+    "kan-Knda",
+    "kor-Hang",
+    "kir-Cyrl",
+    "lao-Laoo",
+    "lit-Latn",
+    "lav-Latn",
+    "mkd-Cyrl",
+    "mal-Mlym",
+    "mon-Cyrl",
+    "mar-Deva",
+    "msa-Latn",
+    "mya-Mymr",
+    "nep-Deva",
+    "nld-Latn",
+    "nor-Latn",
+    "nob-Latn",
+    "nno-Latn",
+    "pan-Guru",
+    "pol-Latn",
+    "por-Latn",
+    "que-Latn",
+    "ron-Latn",
+    "rus-Cyrl",
+    "sin-Sinh",
+    "slk-Latn",
+    "slv-Latn",
+    "swa-Latn",
+    "tam-Taml",
+    "tel-Telu",
+    "tha-Thai",
+    "tgl-Latn",
+    "tur-Latn",
+    "ukr-Cyrl",
+    "urd-Arab",
+    "vie-Latn",
+    "yor-Latn",
+    "zho-Hans",
+]
+
 XLMR_LANGUAGES = [
     "afr-Latn",
     "amh-Latn",
@@ -119,6 +200,28 @@ XLMR_LANGUAGES = [
     "zho-Hans",
 ]
 
+JINARerankerV3_TRAINING_DATA = {
+    "MIRACLRetrieval",
+    "MIRACLRetrievalHardNegatives",
+    "MIRACLReranking",
+    "CMedQAv1-reranking",
+    "CMedQAv2-reranking",
+    "MrTidyRetrieval",
+    "T2Reranking",
+    "MSMARCO",
+    "MSMARCOHardNegatives",
+    "NQ",
+    "NQHardNegatives",
+    "HotpotQA",
+    "HotpotQAHardNegatives",
+    "T2Retrieval",
+    "DuRetrieval",
+    "MMarcoReranking",
+    "CornStack",
+    "MultiLongDocRetrieval",
+    "StackOverflowQA",
+}
+
 JinaV4_TRAINING_DATA = {
     "MSMARCO",
     "MSMARCOHardNegatives",
@@ -139,14 +242,72 @@ JinaV4_TRAINING_DATA = {
     "CornStack",
     "VDRMultilingualRetrieval",
     # from https://huggingface.co/datasets/vidore/colpali_train_set
-    "DocVQA",
-    "InfoVQA",
-    "TATDQA",
-    "arXivQA",
+    "VidoreDocVQARetrieval",
+    "VidoreInfoVQARetrieval",
+    "VidoreTatdqaRetrieval",
+    "VidoreArxivQARetrieval",
     # "other", # inhouse dataset including synthetic datasets
 }
 
 
+class JinaRerankerV3Wrapper(CrossEncoderWrapper):
+    """Wrapper integration for MTEB."""
+
+    def __init__(
+        self,
+        model: CrossEncoder | str,
+        revision: str | None = None,
+        trust_remote_code: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        from sentence_transformers.util import get_device_name
+        from transformers import AutoModel
+
+        self.model = AutoModel.from_pretrained(
+            model, trust_remote_code=trust_remote_code, dtype="auto"
+        )
+
+        device = kwargs.get("device", None)
+        if device is None:
+            device = get_device_name()
+            logger.info(f"Use pytorch device: {device}")
+
+        self.model.to(device)
+        self.model.eval()
+
+    def predict(
+        self,
+        inputs1: DataLoader[BatchedInput],
+        inputs2: DataLoader[BatchedInput],
+        *,
+        task_metadata: TaskMetadata,
+        hf_split: str,
+        hf_subset: str,
+        prompt_type: PromptType | None = None,
+        **kwargs: Any,
+    ) -> Array:
+        all_corpus = [text for batch in inputs2 for text in batch["text"]]
+        all_queries = [text for batch in inputs1 for text in batch["text"]]
+
+        sentences_count = len(all_corpus)
+        query_groups: dict[str, list[tuple[int, str]]] = defaultdict(list)
+        for idx, (query, doc) in enumerate(zip(all_queries, all_corpus)):
+            query_groups[query].append((idx, doc))
+
+        results = np.zeros(sentences_count, dtype=np.float32)
+        for query, doc_infos in query_groups.items():
+            original_indices, docs = zip(*doc_infos)
+
+            scores = self.model.rerank(
+                query, list(docs), max_query_length=3072, max_doc_length=2048
+            )
+            for scr in scores:
+                original_idx = original_indices[scr["index"]]
+                results[original_idx] = float(scr["relevance_score"])
+
+        return results
+
+
 class JinaWrapper(SentenceTransformerEncoderWrapper):
     """following the hf model card documentation."""
 
@@ -553,6 +714,43 @@ def get_programming_task_override(
     return current_task_name
 
 
+jina_reranker_v3 = ModelMeta(
+    loader=JinaRerankerV3Wrapper,
+    loader_kwargs=dict(
+        trust_remote_code=True,
+    ),
+    name="jinaai/jina-reranker-v3",
+    languages=multilingual_langs,
+    open_weights=True,
+    revision="050e171c4f75dfec5b648ed8470a2475e5a30f30",
+    release_date="2025-09-18",  # official release date
+    modalities=["text"],
+    n_parameters=int(0.6 * 1e9),
+    memory_usage_mb=1138,
+    max_tokens=131072,
+    embed_dim=None,
+    license="cc-by-nc-4.0",
+    similarity_fn_name=None,
+    framework=["PyTorch"],
+    use_instructions=None,
+    reference="https://huggingface.co/jinaai/jina-reranker-v3",
+    is_cross_encoder=True,
+    public_training_code=None,
+    public_training_data=None,
+    training_datasets=JINARerankerV3_TRAINING_DATA,
+    adapted_from="Qwen/Qwen3-0.6B",
+    citation="""@misc{wang2025jinarerankerv3lateinteractionlistwise,
+      title={jina-reranker-v3: Last but Not Late Interaction for Listwise Document Reranking}, 
+      author={Feng Wang and Yuqing Li and Han Xiao},
+      year={2025},
+      eprint={2509.25085},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2509.25085},}
+""",
+)
+
+
 jina_embeddings_v4 = ModelMeta(
     loader=JinaV4Wrapper,
     loader_kwargs=dict(

mteb/models/model_implementations/nvidia_llama_nemoretriever_colemb.py

@@ -117,19 +117,18 @@ class LlamaNemoretrieverColembed(AbsEncoder):
 
 TRAINING_DATA = {
     # from https://huggingface.co/datasets/vidore/colpali_train_set
-    "DocVQA",
-    "InfoVQA",
-    "TATDQA",
-    "arXivQA",
-    "hotpotqa",
-    "miracl",
+    "VidoreDocVQARetrieval",
+    "VidoreInfoVQARetrieval",
+    "VidoreTatdqaRetrieval",
+    "VidoreArxivQARetrieval",
+    "HotpotQA",
+    "MIRACLRetrieval",
     "NQ",
-    "stackexchange",
+    "StackExchangeClustering",
     "SQuAD",
     "WebInstructSub",
     "docmatix-ir",
-    "vdr-multilingual-train",
-    "colpali_train_set",  # as it contains PDFs
+    "VDRMultilingualRetrieval",
     "VisRAG-Ret-Train-Synthetic-data",
     "VisRAG-Ret-Train-In-domain-data",
     "wiki-ss-nq",

mteb/models/model_implementations/vdr_models.py

@@ -43,5 +43,6 @@ vdr_2b_multi_v1 = ModelMeta(
     public_training_data="https://huggingface.co/datasets/llamaindex/vdr-multilingual-train",
     training_datasets=set(
         # llamaindex/vdr-multilingual-train
+        "VDRMultilingualRetrieval",
     ),
 )

mteb/models/model_implementations/yuan_models_en.py

@@ -0,0 +1,57 @@
+from mteb.models.instruct_wrapper import InstructSentenceTransformerModel
+from mteb.models.model_meta import ModelMeta
+from mteb.models.models_protocols import PromptType
+
+
+def instruction_template(
+    instruction: str, prompt_type: PromptType | None = None
+) -> str:
+    if not instruction or prompt_type == PromptType.document:
+        return ""
+    if isinstance(instruction, dict):
+        if prompt_type is None:
+            instruction = next(iter(instruction.values()))  # TODO
+        else:
+            instruction = instruction[prompt_type]
+    return f"Instruct: {instruction}\nQuery:"
+
+
+training_data = {
+    "T2Retrieval",
+    "DuRetrieval",
+    "MMarcoReranking",
+    "CMedQAv2-reranking",
+    "NQ",
+    "MSMARCO",
+    "HotpotQA",
+    "MrTidyRetrieval",
+    "MIRACLRetrieval",
+    "CodeSearchNet",
+}
+
+
+yuan_embedding_2_en = ModelMeta(
+    loader=InstructSentenceTransformerModel,
+    loader_kwargs=dict(
+        instruction_template=instruction_template,
+        apply_instruction_to_passages=False,
+    ),
+    name="IEITYuan/Yuan-embedding-2.0-en",
+    languages=["eng-Latn"],
+    open_weights=True,
+    revision="b2fd15da3bcae3473c8529593825c15068f09fce",
+    release_date="2025-11-27",
+    n_parameters=595776512,
+    memory_usage_mb=2272,
+    embed_dim=1024,
+    max_tokens=2048,
+    license="apache-2.0",
+    reference="https://huggingface.co/IEITYuan/Yuan-embedding-2.0-en",
+    similarity_fn_name="cosine",
+    framework=["Sentence Transformers", "PyTorch"],
+    use_instructions=True,
+    public_training_code=None,
+    public_training_data=None,
+    training_datasets=training_data,
+    adapted_from="Qwen/Qwen3-Embedding-0.6B",
+)

mteb/results/model_result.py

@@ -22,7 +22,7 @@ from mteb.types import (
     SplitName,
 )
 
-from .task_result import TaskResult
+from .task_result import TaskError, TaskResult
 
 logger = logging.getLogger(__name__)
 
@@ -82,6 +82,7 @@ class ModelResult(BaseModel):
             protected_namespaces=(),
         )
     )
+    exceptions: list[TaskError] | None = None
 
     def __repr__(self) -> str:
         n_entries = len(self.task_results)

mteb/results/task_result.py

@@ -840,3 +840,15 @@ class TaskResult(BaseModel):
                     )
                 )
         return results
+
+
+class TaskError(BaseModel):
+    """A class to represent an error that occurred during the evaluation of a task.
+
+    Attributes:
+        task_name: The name of the MTEB task.
+        exception: The error message that occurred during the evaluation.
+    """
+
+    task_name: str
+    exception: str

mteb/tasks/reranking/multilingual/__init__.py

@@ -1,6 +1,7 @@
 from .esci_reranking import ESCIReranking
 from .hume_wikipedia_reranking_multilingual import HUMEWikipediaRerankingMultilingual
 from .miracl_reranking import MIRACLReranking
+from .multi_long_doc_reranking import MultiLongDocReranking
 from .wikipedia_reranking_multilingual import WikipediaRerankingMultilingual
 from .x_glue_wpr_reranking import XGlueWPRReranking
 
@@ -8,6 +9,7 @@ __all__ = [
     "ESCIReranking",
     "HUMEWikipediaRerankingMultilingual",
     "MIRACLReranking",
+    "MultiLongDocReranking",
     "WikipediaRerankingMultilingual",
     "XGlueWPRReranking",
 ]