mteb 2.5.3__py3-none-any.whl → 2.5.5__py3-none-any.whl

mteb/models/model_implementations/openclip_models.py

@@ -120,7 +120,7 @@ def openclip_loader(model_name, **kwargs):
 
 
 CLIP_ViT_L_14_DataComp_XL_s13B_b90K = ModelMeta(
-    loader=openclip_loader,  # type: ignore
+    loader=openclip_loader,
     name="laion/CLIP-ViT-L-14-DataComp.XL-s13B-b90K",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -146,7 +146,7 @@ CLIP_ViT_L_14_DataComp_XL_s13B_b90K = ModelMeta(
 )
 
 CLIP_ViT_B_32_DataComp_XL_s13B_b90K = ModelMeta(
-    loader=openclip_loader,  # type: ignore
+    loader=openclip_loader,
     name="laion/CLIP-ViT-B-32-DataComp.XL-s13B-b90K",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -172,7 +172,7 @@ CLIP_ViT_B_32_DataComp_XL_s13B_b90K = ModelMeta(
 )
 
 CLIP_ViT_B_16_DataComp_XL_s13B_b90K = ModelMeta(
-    loader=openclip_loader,  # type: ignore
+    loader=openclip_loader,
     name="laion/CLIP-ViT-B-16-DataComp.XL-s13B-b90K",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -198,7 +198,7 @@ CLIP_ViT_B_16_DataComp_XL_s13B_b90K = ModelMeta(
 )
 
 CLIP_ViT_bigG_14_laion2B_39B_b160k = ModelMeta(
-    loader=openclip_loader,  # type: ignore
+    loader=openclip_loader,
     name="laion/CLIP-ViT-bigG-14-laion2B-39B-b160k",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -224,7 +224,7 @@ CLIP_ViT_bigG_14_laion2B_39B_b160k = ModelMeta(
 )
 
 CLIP_ViT_g_14_laion2B_s34B_b88K = ModelMeta(
-    loader=openclip_loader,  # type: ignore
+    loader=openclip_loader,
     name="laion/CLIP-ViT-g-14-laion2B-s34B-b88K",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -250,7 +250,7 @@ CLIP_ViT_g_14_laion2B_s34B_b88K = ModelMeta(
 )
 
 CLIP_ViT_H_14_laion2B_s32B_b79K = ModelMeta(
-    loader=openclip_loader,  # type: ignore
+    loader=openclip_loader,
     name="laion/CLIP-ViT-H-14-laion2B-s32B-b79K",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -276,7 +276,7 @@ CLIP_ViT_H_14_laion2B_s32B_b79K = ModelMeta(
 )
 
 CLIP_ViT_L_14_laion2B_s32B_b82K = ModelMeta(
-    loader=openclip_loader,  # type: ignore
+    loader=openclip_loader,
     name="laion/CLIP-ViT-L-14-laion2B-s32B-b82K",
     model_type=["dense"],
     languages=["eng-Latn"],

mteb/models/model_implementations/random_baseline.py

@@ -68,7 +68,7 @@ _common_mock_metadata = dict(
     license="mit",
     max_tokens=np.inf,
     reference=None,
-    similarity_fn_name="cosine",  # type: ignore
+    similarity_fn_name="cosine",
     framework=[],
     use_instructions=False,
     public_training_code=None,  # No training code, as this is a random baseline
@@ -187,7 +187,7 @@ class RandomEncoderBaseline:
 
 
 random_encoder_baseline = ModelMeta(
-    loader=RandomEncoderBaseline,  # type: ignore
+    loader=RandomEncoderBaseline,
     name="baseline/random-encoder-baseline",
     model_type=["dense"],
     modalities=["text", "image"],
@@ -232,7 +232,7 @@ class RandomCrossEncoderBaseline:
 
 
 random_cross_encoder_baseline = ModelMeta(
-    loader=RandomCrossEncoderBaseline,  # type: ignore
+    loader=RandomCrossEncoderBaseline,
     name="baseline/random-cross-encoder-baseline",
     model_type=["cross-encoder"],
     modalities=["text", "image"],

mteb/models/model_implementations/rasgaard_models.py

@@ -4,7 +4,7 @@ from mteb.models.model_implementations.model2vec_models import Model2VecModel
 from mteb.models.model_meta import ModelMeta, ScoringFunction
 
 potion_base_8m = ModelMeta(
-    loader=Model2VecModel,  # type: ignore
+    loader=Model2VecModel,
     name="rasgaard/m2v-dfm-large",
     model_type=["dense"],
     languages=["dan-Latn"],

mteb/models/model_implementations/repllama_models.py

@@ -154,7 +154,7 @@ REPLLAMA_CITATION = """
 """
 
 repllama_llama2_original = ModelMeta(
-    loader=RepLLaMAModel,  # type: ignore
+    loader=RepLLaMAModel,
     loader_kwargs=dict(
         base_model_name_or_path="meta-llama/Llama-2-7b-hf",
         device_map="auto",
@@ -187,7 +187,7 @@ repllama_llama2_original = ModelMeta(
 
 
 repllama_llama2_reproduced = ModelMeta(
-    loader=RepLLaMAModel,  # type: ignore
+    loader=RepLLaMAModel,
     loader_kwargs=dict(
         base_model_name_or_path="meta-llama/Llama-2-7b-hf",
         device_map="auto",

mteb/models/model_implementations/rerankers_custom.py

@@ -214,7 +214,7 @@ class JinaReranker(RerankerWrapper):
 
 
 monobert_large = ModelMeta(
-    loader=MonoBERTReranker,  # type: ignore
+    loader=MonoBERTReranker,
     loader_kwargs=dict(
         fp_options="float16",
     ),
@@ -239,7 +239,7 @@ monobert_large = ModelMeta(
 
 # languages unclear: https://huggingface.co/jinaai/jina-reranker-v2-base-multilingual/discussions/28
 jina_reranker_multilingual = ModelMeta(
-    loader=JinaReranker,  # type: ignore
+    loader=JinaReranker,
     loader_kwargs=dict(
         fp_options="float16",
     ),
@@ -263,7 +263,7 @@ jina_reranker_multilingual = ModelMeta(
 )
 
 bge_reranker_v2_m3 = ModelMeta(
-    loader=BGEReranker,  # type: ignore
+    loader=BGEReranker,
     loader_kwargs=dict(
         fp_options="float16",
     ),

mteb/models/model_implementations/rerankers_monot5_based.py

@@ -343,7 +343,7 @@ monot5_small = ModelMeta(
 )
 
 monot5_base = ModelMeta(
-    loader=MonoT5Reranker,  # type: ignore
+    loader=MonoT5Reranker,
     loader_kwargs=dict(
         fp_options="float16",
     ),
@@ -442,7 +442,7 @@ monot5_3b = ModelMeta(
 )
 
 flant5_base = ModelMeta(
-    loader=FLANT5Reranker,  # type: ignore
+    loader=FLANT5Reranker,
     loader_kwargs=dict(
         fp_options="float16",
     ),
@@ -902,7 +902,7 @@ mt5_base_mmarco_v2 = ModelMeta(
 )
 
 mt5_13b_mmarco_100k = ModelMeta(
-    loader=MonoT5Reranker,  # type: ignore
+    loader=MonoT5Reranker,
     loader_kwargs=dict(
         fp_options="float16",
     ),

mteb/models/model_implementations/seed_1_6_embedding_models_1215.py

@@ -4,13 +4,15 @@ import base64
 import logging
 import os
 import time
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
 from io import BytesIO
 from typing import TYPE_CHECKING, Any
 
 import requests
 import torch
 from torch.utils.data import DataLoader
+from tqdm import tqdm
 
 from mteb._requires_package import requires_package
 from mteb.abstasks.task_metadata import TaskMetadata
@@ -26,114 +28,6 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-
-def pil_to_base64(image, format="jpeg"):
-    if image is None:
-        return None
-    buffer = BytesIO()
-    image.save(buffer, format=format)
-    img_bytes = buffer.getvalue()
-    encoded_bytes = base64.b64encode(img_bytes)
-    return encoded_bytes.decode("utf-8")
-
-
-def multimodal_embedding(image_base64=None, text_content=None):
-    auth_token = os.getenv("VOLCES_AUTH_TOKEN")
-    model_name = "doubao-embedding-vision-251215"
-    api_url = "https://ark.cn-beijing.volces.com/api/v3/embeddings/multimodal"
-
-    headers = {
-        "Authorization": f"Bearer {auth_token}",
-        "x-ark-vlm1": "true",
-        "Content-Type": "application/json",
-    }
-
-    if image_base64 is not None and text_content is None:
-        inputs = []
-        for image in image_base64:
-            image_format = "jpeg"
-            image_data = f"data:image/{image_format};base64,{image}"
-            inputs.append({"type": "image_url", "image_url": {"url": image_data}})
-
-        payload = {"model": model_name, "input": inputs}
-    elif image_base64 is None and text_content is not None:
-        payload = {
-            "model": model_name,
-            "input": [
-                {"type": "text", "text": text_content},
-            ],
-        }
-    else:
-        inputs = []
-        for image in image_base64:
-            image_format = "jpeg"
-            image_data = f"data:image/{image_format};base64,{image}"
-            inputs.append({"type": "image_url", "image_url": {"url": image_data}})
-        inputs.append({"type": "text", "text": text_content})
-        payload = {"model": model_name, "input": inputs}
-
-    try:
-        response = requests.post(url=api_url, headers=headers, json=payload, timeout=10)
-
-        response.raise_for_status()
-        return response.json()
-
-    except requests.exceptions.HTTPError as http_err:
-        logger.error(f"HTTP error ({http_err.response.status_code}): {http_err}")
-    except requests.exceptions.JSONDecodeError:
-        logger.error("Error:The response is not in valid JSON format")
-    except requests.exceptions.Timeout:
-        logger.error("Error:Request timeout")
-    except Exception as e:
-        logger.error(f"Unknown error: {str(e)}")
-
-    return None
-
-
-def multi_thread_encode(sentences, batch_size=1, max_workers=8):
-    batches = []
-    for idx in range(0, len(sentences), batch_size):
-        batches.append((idx // batch_size, sentences[idx : idx + batch_size]))
-
-    n_batches = len(batches)
-    results = [None] * n_batches  # Pre-allocated result list
-    all_embeddings = []  # Final ordered embeddings
-
-    def _process_batch(batch_idx, batch_sentences):
-        sentence = batch_sentences[0]
-
-        retries = 5
-        while retries > 0:
-            try:
-                resp = multimodal_embedding(text_content=sentence)
-                embedding = torch.tensor(resp["data"]["embedding"])
-                break
-            except Exception as e:
-                time.sleep(1)
-                logger.warning(f"Retrying... {retries} retries left. Error: {str(e)}")
-                retries -= 1
-                if retries == 0:
-                    raise e
-        return batch_idx, embedding
-
-    with ThreadPoolExecutor(max_workers=max_workers) as executor:
-        futures = {
-            executor.submit(_process_batch, idx, batch): idx for idx, batch in batches
-        }
-
-        for future in as_completed(futures):
-            batch_idx, embeddings = future.result()
-            results[batch_idx] = embeddings
-
-    for batch_embeddings in results:
-        all_embeddings.append(batch_embeddings)
-
-    all_embeddings = torch.stack(all_embeddings, dim=0)
-    all_embeddings = torch.nn.functional.normalize(all_embeddings, dim=-1)
-
-    return all_embeddings.float().cpu()
-
-
 doubao_embedding_training_data = (
     {
         "PawsXPairClassification",
@@ -166,25 +60,80 @@ class Seed16EmbeddingWrapper(AbsEncoder):
             "pip install mteb[ark]",
             "tiktoken",
         )
-        import tiktoken
 
         self._model_name = model_name
         self._max_tokens = 32768
         self._embed_dim = embed_dim
         self._available_embed_dims = [2048, 1024]
-        self._encoding = tiktoken.get_encoding(tokenizer_name)
 
-    def truncate_text_tokens(self, text: str) -> str:
-        """Truncate a string to have `max_tokens` according to the given encoding.
+    def pil_to_base64(self, image, format="jpeg"):
+        if image is None:
+            return None
+        buffer = BytesIO()
+        image.save(buffer, format=format)
+        img_bytes = buffer.getvalue()
+        encoded_bytes = base64.b64encode(img_bytes)
+        return encoded_bytes.decode("utf-8")
+
+    def multimodal_embedding(self, instruction, image_base64, text_content):
+        auth_token = os.getenv("VOLCES_AUTH_TOKEN")
+        model_name = "doubao-embedding-vision-251215"
+        api_url = "https://ark.cn-beijing.volces.com/api/v3/embeddings/multimodal"
+
+        headers = {
+            "Authorization": f"Bearer {auth_token}",
+            "x-ark-vlm1": "true",
+            "Content-Type": "application/json",
+        }
 
-        Args:
-            text: The input string to be truncated.
+        if text_content is not None and len(text_content) > self._max_tokens:
+            text_content = text_content[: self._max_tokens]
+
+        if image_base64 is not None and text_content is None:
+            inputs = []
+            for image in image_base64:
+                image_format = "jpeg"
+                image_data = f"data:image/{image_format};base64,{image}"
+                inputs.append({"type": "image_url", "image_url": {"url": image_data}})
+
+            payload = {"model": model_name, "input": inputs}
+        elif image_base64 is None and text_content is not None:
+            payload = {
+                "model": model_name,
+                "instruction": instruction,
+                "input": [
+                    {"type": "text", "text": text_content},
+                ],
+            }
+        else:
+            inputs = []
+            for image in image_base64:
+                image_format = "jpeg"
+                image_data = f"data:image/{image_format};base64,{image}"
+                inputs.append({"type": "image_url", "image_url": {"url": image_data}})
+            inputs.append({"type": "text", "text": text_content})
+            payload = {"model": model_name, "input": inputs}
+
+        max_retries = 3
+        retry_count = 0
+
+        while retry_count < max_retries:
+            response = requests.post(
+                url=api_url, headers=headers, json=payload, timeout=30
+            )
 
-        Returns:
-            The truncated string.
-        """
-        truncated_sentence = self._encoding.encode(text)[: self._max_tokens]
-        return self._encoding.decode(truncated_sentence)
+            if response.status_code != 200:
+                retry_count += 1
+                time.sleep(3)
+                continue
+
+            response_json = response.json()
+            return response_json
+
+        raise Exception(
+            f"Request failed with status code {response.status_code}. "
+            f"Response: {response.text}"
+        )
 
     def get_fused_embeddings(
         self,
@@ -204,59 +153,69 @@ class Seed16EmbeddingWrapper(AbsEncoder):
         if images is not None and texts is not None:
             assert len(texts) == len(images)
             batch_len = len(texts)
-            images_base64 = [pil_to_base64(image) for image in images]
+            images_base64 = [self.pil_to_base64(image) for image in images]
         elif images is None:
             batch_len = len(texts)
             images_base64 = [None for _ in range(batch_len)]
         elif texts is None:
             batch_len = len(images)
-            images_base64 = [pil_to_base64(image) for image in images]
+            images_base64 = [self.pil_to_base64(image) for image in images]
         else:
             raise ValueError("images and texts cannot be None at the same time")
 
-        outputs = []
-        for i in range(batch_len):
+        def process_item(
+            i, prompt_type, task_name, texts, images_base64, multimodal_embedding
+        ):
             if (
                 prompt_type == PromptType("query") or prompt_type is None
             ) and task_name in TASK_NAME_TO_INSTRUCTION:
                 instruction = TASK_NAME_TO_INSTRUCTION[task_name]
                 instruction = instruction.rstrip("{}").rstrip("\n")
-                if texts[i] != "":
-                    input_text = (
-                        "Target_modality:Text.\n Instruction:"
-                        + instruction
-                        + "\n Query:{}"
-                    ).format(texts[i])
-                else:
-                    input_text = (
-                        "Target_modality:Text.\n Instruction:"
-                        + instruction
-                        + "\n Query:"
-                    )
+                instruction = (
+                    "Target_modality:Text.\n Instruction:" + instruction + "\n Query:"
+                )
+                input_text = texts[i]
             else:
                 if texts[i] != "" and images_base64[i] is not None:
-                    instruction = "Instruction: Compress the the text and image into one word.\n Query: {}"
-                    input_text = instruction.format(texts[i])
+                    instruction = "Instruction: Compress the text and image into one word.\n Query:"
+                    input_text = texts[i]
                 elif texts[i] != "":
                     instruction = (
-                        "Instruction: Compress the the text into one word.\n Query: {}"
+                        "Instruction: Compress the text into one word.\n Query:"
                     )
-                    input_text = instruction.format(texts[i])
+                    input_text = texts[i]
                 elif images_base64[i] is not None:
                     instruction = (
-                        "Instruction: Compress the the image into one word.\n Query:"
+                        "Instruction: Compress the image into one word.\n Query:"
                     )
-                    input_text = instruction
+                    input_text = None
                 else:
                     raise ValueError("image and text are both None")
 
             resp = multimodal_embedding(
-                image_base64=[images_base64[i]], text_content=input_text
+                instruction=instruction,
+                image_base64=images_base64[i],
+                text_content=input_text,
             )
             embedding = torch.tensor(resp["data"]["embedding"])
             embedding = torch.reshape(embedding, (1, -1))
+            return embedding
+
+        outputs = []
+        process_partial = partial(
+            process_item,
+            prompt_type=prompt_type,
+            task_name=task_name,
+            texts=texts,
+            images_base64=images_base64,
+            multimodal_embedding=self.multimodal_embedding,
+        )
+        with ThreadPoolExecutor(max_workers=15) as executor:
+            futures = [executor.submit(process_partial, i) for i in range(batch_len)]
+            for future in tqdm(futures, total=batch_len, desc="Encoding"):
+                outputs.append(future.result())
 
-        outputs = torch.stack(outputs, dim=0)
+        outputs = torch.stack(outputs, dim=0).squeeze(1)
 
         if self._embed_dim is not None:
             outputs = outputs[:, : self._embed_dim]
@@ -273,13 +232,21 @@ class Seed16EmbeddingWrapper(AbsEncoder):
         prompt_type: PromptType | None = None,
         **kwargs: Any,
     ) -> Array:
-        sentences = [text for batch in inputs for text in batch["text"]]
-        images = [image for batch in inputs for image in batch["image"]]
+        if "text" in inputs.dataset.features:
+            sentences = [text for batch in inputs for text in batch["text"]]
+        else:
+            sentences = None
+
+        if "image" in inputs.dataset.features:
+            images = [image for batch in inputs for image in batch["image"]]
+        else:
+            images = None
 
         return self.get_fused_embeddings(
             texts=sentences,
             images=images,
             task_name=task_metadata.name,
+            prompt_type=prompt_type,
             **kwargs,
         )
 

mteb/models/model_implementations/siglip_models.py

@@ -123,7 +123,7 @@ siglip_training_datasets = set(
 )
 
 siglip_so400m_patch14_224 = ModelMeta(
-    loader=SiglipModelWrapper,  # type: ignore
+    loader=SiglipModelWrapper,
     name="google/siglip-so400m-patch14-224",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -147,7 +147,7 @@ siglip_so400m_patch14_224 = ModelMeta(
 )
 
 siglip_so400m_patch14_384 = ModelMeta(
-    loader=SiglipModelWrapper,  # type: ignore
+    loader=SiglipModelWrapper,
     name="google/siglip-so400m-patch14-384",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -171,7 +171,7 @@ siglip_so400m_patch14_384 = ModelMeta(
 )
 
 siglip_so400m_patch16_256_i18n = ModelMeta(
-    loader=SiglipModelWrapper,  # type: ignore
+    loader=SiglipModelWrapper,
     name="google/siglip-so400m-patch16-256-i18n",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -195,7 +195,7 @@ siglip_so400m_patch16_256_i18n = ModelMeta(
 )
 
 siglip_base_patch16_256_multilingual = ModelMeta(
-    loader=SiglipModelWrapper,  # type: ignore
+    loader=SiglipModelWrapper,
     name="google/siglip-base-patch16-256-multilingual",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -219,7 +219,7 @@ siglip_base_patch16_256_multilingual = ModelMeta(
 )
 
 siglip_base_patch16_256 = ModelMeta(
-    loader=SiglipModelWrapper,  # type: ignore
+    loader=SiglipModelWrapper,
     name="google/siglip-base-patch16-256",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -243,7 +243,7 @@ siglip_base_patch16_256 = ModelMeta(
 )
 
 siglip_base_patch16_512 = ModelMeta(
-    loader=SiglipModelWrapper,  # type: ignore
+    loader=SiglipModelWrapper,
     name="google/siglip-base-patch16-512",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -267,7 +267,7 @@ siglip_base_patch16_512 = ModelMeta(
 )
 
 siglip_base_patch16_384 = ModelMeta(
-    loader=SiglipModelWrapper,  # type: ignore
+    loader=SiglipModelWrapper,
     name="google/siglip-base-patch16-384",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -291,7 +291,7 @@ siglip_base_patch16_384 = ModelMeta(
 )
 
 siglip_base_patch16_224 = ModelMeta(
-    loader=SiglipModelWrapper,  # type: ignore
+    loader=SiglipModelWrapper,
     name="google/siglip-base-patch16-224",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -315,7 +315,7 @@ siglip_base_patch16_224 = ModelMeta(
 )
 
 siglip_large_patch16_256 = ModelMeta(
-    loader=SiglipModelWrapper,  # type: ignore
+    loader=SiglipModelWrapper,
     name="google/siglip-large-patch16-256",
     model_type=["dense"],
     languages=["eng-Latn"],
@@ -339,7 +339,7 @@ siglip_large_patch16_256 = ModelMeta(
 )
 
 siglip_large_patch16_384 = ModelMeta(
-    loader=SiglipModelWrapper,  # type: ignore
+    loader=SiglipModelWrapper,
     name="google/siglip-large-patch16-384",
     model_type=["dense"],
     languages=["eng-Latn"],

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

mteb/models/model_implementations/voyage_v.py

@@ -40,15 +40,15 @@ def _downsample_image(
         logging.info(
             f"Downsampling image from {width}x{height} to {new_width}x{new_height}"
         )
-        return image.resize(new_size, Image.LANCZOS)  # type: ignore
+        return image.resize(new_size, Image.LANCZOS)
     if width > height:
         if width > 10000:
             logging.error("Processing extremely wide images.")
-            return image.resize((10000, height), Image.LANCZOS)  # type: ignore
+            return image.resize((10000, height), Image.LANCZOS)
     else:
         if height > 10000:
             logging.error("Processing extremely high images.")
-            return image.resize((width, 10000), Image.LANCZOS)  # type: ignore
+            return image.resize((width, 10000), Image.LANCZOS)
     return image
 
 
@@ -202,7 +202,7 @@ def voyage_v_loader(model_name, **kwargs):
 
 
 voyage_v = ModelMeta(
-    loader=voyage_v_loader,  # type: ignore
+    loader=voyage_v_loader,
     name="voyageai/voyage-multimodal-3",
     model_type=["dense"],
     languages=[],  # Unknown