[MIEB] Fix get_fused_emebddings (#1612)

* Fix fused

* fix vlm2vec

* Fix lint

mteb/models/align_models.py

@@ -101,9 +101,6 @@ def get_fused_embeddings(
         self,
         texts: list[str] = None,
         images: list[Image.Image] | DataLoader = None,
-        task_name: str | None = None,
-        prompt_type: PromptType | None = None,
-        batch_size: int = 32,
         fusion_mode="sum",
         **kwargs: Any,
     ):
@@ -114,10 +111,10 @@ def get_fused_embeddings(
         image_embeddings = None
 
         if texts is not None:
-            text_embeddings = self.get_text_embeddings(texts, batch_size)
+            text_embeddings = self.get_text_embeddings(texts, **kwargs)
 
         if images is not None:
-            image_embeddings = self.get_image_embeddings(images, batch_size)
+            image_embeddings = self.get_image_embeddings(images, **kwargs)
 
         if text_embeddings is not None and image_embeddings is not None:
             if len(text_embeddings) != len(image_embeddings):

mteb/models/blip2_models.py

@@ -174,9 +174,6 @@ def get_fused_embeddings(
             self,
             texts: list[str] = None,
             images: list[Image.Image] | DataLoader = None,
-            task_name: str | None = None,
-            prompt_type: PromptType | None = None,
-            batch_size: int = 32,
             fusion_mode="sum",
             **kwargs: Any,
         ):
@@ -188,10 +185,10 @@ def get_fused_embeddings(
             image_embeddings = None
 
             if texts is not None:
-                text_embeddings = self.get_text_embeddings(texts, batch_size)
+                text_embeddings = self.get_text_embeddings(texts, **kwargs)
 
             if images is not None:
-                image_embeddings = self.get_image_embeddings(images, batch_size)
+                image_embeddings = self.get_image_embeddings(images, **kwargs)
 
             if text_embeddings is not None and image_embeddings is not None:
                 if len(text_embeddings) != len(image_embeddings):
@@ -202,7 +199,7 @@ def get_fused_embeddings(
                     fused_embeddings = text_embeddings + image_embeddings
                 elif fusion_mode == "multimodal":
                     fused_embeddings = self.get_multimodal_embeddings(
-                        texts, images, batch_size
+                        texts, images, kwargs.get("batch_size", 32)
                     )
                 else:
                     # to do: add other fusion mode

mteb/models/blip_models.py

@@ -121,9 +121,6 @@ def get_fused_embeddings(
         self,
         texts: list[str] = None,
         images: list[Image.Image] | DataLoader = None,
-        task_name: str | None = None,
-        prompt_type: PromptType | None = None,
-        batch_size: int = 32,
         fusion_mode="sum",
         **kwargs: Any,
     ):
@@ -134,10 +131,10 @@ def get_fused_embeddings(
         image_embeddings = None
 
         if texts is not None:
-            text_embeddings = self.get_text_embeddings(texts, batch_size)
+            text_embeddings = self.get_text_embeddings(texts, **kwargs)
 
         if images is not None:
-            image_embeddings = self.get_image_embeddings(images, batch_size)
+            image_embeddings = self.get_image_embeddings(images, **kwargs)
 
         if text_embeddings is not None and image_embeddings is not None:
             if len(text_embeddings) != len(image_embeddings):

mteb/models/clip_models.py

@@ -105,9 +105,6 @@ def get_fused_embeddings(
         self,
         texts: list[str] = None,
         images: list[Image.Image] | DataLoader = None,
-        task_name: str | None = None,
-        prompt_type: PromptType | None = None,
-        batch_size: int = 32,
         fusion_mode="sum",
         **kwargs: Any,
     ):
@@ -118,10 +115,10 @@ def get_fused_embeddings(
         image_embeddings = None
 
         if texts is not None:
-            text_embeddings = self.get_text_embeddings(texts, batch_size)
+            text_embeddings = self.get_text_embeddings(texts, **kwargs)
 
         if images is not None:
-            image_embeddings = self.get_image_embeddings(images, batch_size)
+            image_embeddings = self.get_image_embeddings(images, **kwargs)
 
         if text_embeddings is not None and image_embeddings is not None:
             if len(text_embeddings) != len(image_embeddings):

mteb/models/cohere_v.py

@@ -145,9 +145,6 @@ def get_fused_embeddings(
             self,
             texts: list[str] = None,
             images: list[Image.Image] | DataLoader = None,
-            task_name: str | None = None,
-            prompt_type: PromptType | None = None,
-            batch_size: int = 32,
             fusion_mode="sum",
             **kwargs: Any,
         ):
@@ -158,10 +155,10 @@ def get_fused_embeddings(
             image_embeddings = None
 
             if texts is not None:
-                text_embeddings = self.get_text_embeddings(texts, batch_size)
+                text_embeddings = self.get_text_embeddings(texts, **kwargs)
 
             if images is not None:
-                image_embeddings = self.get_image_embeddings(images, batch_size)
+                image_embeddings = self.get_image_embeddings(images, **kwargs)
 
             if text_embeddings is not None and image_embeddings is not None:
                 if len(text_embeddings) != len(image_embeddings):

mteb/models/dino_models.py

@@ -97,9 +97,6 @@ def get_fused_embeddings(
         self,
         texts: list[str] = None,
         images: list[Image.Image] | DataLoader = None,
-        task_name: str | None = None,
-        prompt_type: PromptType | None = None,
-        batch_size: int = 32,
         fusion_mode="sum",
         **kwargs: Any,
     ):
@@ -110,10 +107,10 @@ def get_fused_embeddings(
         image_embeddings = None
 
         if texts is not None:
-            text_embeddings = self.get_text_embeddings(texts, batch_size)
+            text_embeddings = self.get_text_embeddings(texts, **kwargs)
 
         if images is not None:
-            image_embeddings = self.get_image_embeddings(images, batch_size)
+            image_embeddings = self.get_image_embeddings(images, **kwargs)
 
         if text_embeddings is not None and image_embeddings is not None:
             raise ValueError("DINO models only support image encoding.")

mteb/models/e5_v.py

@@ -120,16 +120,14 @@ def get_fused_embeddings(
         self,
         texts: list[str] = None,
         images: list[Image.Image] = None,
-        *,
-        task_name: str | None = None,
-        prompt_type: PromptType | None = None,
         batch_size: int = 8,
         **kwargs: Any,
     ):
         if texts is None and images is None:
             raise ValueError("Either texts or images must be provided")
 
         all_fused_embeddings = []
+        kwargs.update(batch_size=batch_size)
 
         if texts is not None and images is not None:
             with torch.no_grad():
@@ -168,10 +166,10 @@ def get_fused_embeddings(
                         all_fused_embeddings.append(outputs.cpu())
             return torch.cat(all_fused_embeddings, dim=0)
         elif texts is not None:
-            text_embeddings = self.get_text_embeddings(texts, batch_size)
+            text_embeddings = self.get_text_embeddings(texts, **kwargs)
             return text_embeddings
         elif images is not None:
-            image_embeddings = self.get_image_embeddings(images, batch_size)
+            image_embeddings = self.get_image_embeddings(images, **kwargs)
             return image_embeddings
 
 

mteb/models/evaclip_models.py

@@ -129,10 +129,6 @@ def get_fused_embeddings(
             self,
             texts: list[str] = None,
             images: list[Image.Image] | DataLoader = None,
-            *,
-            task_name: str | None = None,
-            prompt_type: PromptType | None = None,
-            batch_size: int = 32,
             fusion_mode="sum",
             **kwargs: Any,
         ):
@@ -143,10 +139,10 @@ def get_fused_embeddings(
             image_embeddings = None
 
             if texts is not None:
-                text_embeddings = self.get_text_embeddings(texts, batch_size)
+                text_embeddings = self.get_text_embeddings(texts, **kwargs)
 
             if images is not None:
-                image_embeddings = self.get_image_embeddings(images, batch_size)
+                image_embeddings = self.get_image_embeddings(images, **kwargs)
 
             if text_embeddings is not None and image_embeddings is not None:
                 if len(text_embeddings) != len(image_embeddings):

mteb/models/jina_clip.py

@@ -101,10 +101,6 @@ def get_fused_embeddings(
         self,
         texts: list[str] = None,
         images: list[Image.Image] = None,
-        *,
-        task_name: str | None = None,
-        prompt_type: PromptType | None = None,
-        batch_size: int = 32,
         fusion_mode="sum",
         **kwargs: Any,
     ):
@@ -116,18 +112,12 @@ def get_fused_embeddings(
 
         if texts is not None:
             text_embeddings = self.get_text_embeddings(
-                texts,
-                batch_size=batch_size,
-                convert_to_numpy=False,
-                convert_to_tensor=True,
+                texts, convert_to_numpy=False, convert_to_tensor=True, **kwargs
             )
 
         if images is not None:
             image_embeddings = self.get_image_embeddings(
-                images,
-                batch_size=batch_size,
-                convert_to_numpy=False,
-                convert_to_tensor=True,
+                images, convert_to_numpy=False, convert_to_tensor=True, **kwargs
             )
 
         if text_embeddings is not None and image_embeddings is not None:

mteb/models/nomic_models_vision.py

@@ -128,10 +128,6 @@ def get_fused_embeddings(
         self,
         texts: list[str] = None,
         images: list[Image.Image] | DataLoader = None,
-        *,
-        task_name: str | None = None,
-        prompt_type: PromptType | None = None,
-        batch_size: int = 32,
         fusion_mode="sum",
         **kwargs: Any,
     ):
@@ -142,10 +138,10 @@ def get_fused_embeddings(
         image_embeddings = None
 
         if texts is not None:
-            text_embeddings = self.get_text_embeddings(texts, batch_size)
+            text_embeddings = self.get_text_embeddings(texts, **kwargs)
 
         if images is not None:
-            image_embeddings = self.get_image_embeddings(images, batch_size)
+            image_embeddings = self.get_image_embeddings(images, **kwargs)
 
         if text_embeddings is not None and image_embeddings is not None:
             if len(text_embeddings) != len(image_embeddings):

mteb/models/openclip_models.py

@@ -114,10 +114,6 @@ def get_fused_embeddings(
             self,
             texts: list[str] = None,
             images: list[Image.Image] | DataLoader = None,
-            *,
-            task_name: str | None = None,
-            prompt_type: PromptType | None = None,
-            batch_size: int = 32,
             fusion_mode="sum",
             **kwargs: Any,
         ):
@@ -128,10 +124,10 @@ def get_fused_embeddings(
             image_embeddings = None
 
             if texts is not None:
-                text_embeddings = self.get_text_embeddings(texts, batch_size)
+                text_embeddings = self.get_text_embeddings(texts, **kwargs)
 
             if images is not None:
-                image_embeddings = self.get_image_embeddings(images, batch_size)
+                image_embeddings = self.get_image_embeddings(images, **kwargs)
 
             if text_embeddings is not None and image_embeddings is not None:
                 if len(text_embeddings) != len(image_embeddings):

mteb/models/siglip_models.py

@@ -123,10 +123,6 @@ def get_fused_embeddings(
         self,
         texts: list[str] = None,
         images: list[Image.Image] | DataLoader = None,
-        *,
-        task_name: str | None = None,
-        prompt_type: PromptType | None = None,
-        batch_size: int = 32,
         fusion_mode="sum",
         **kwargs: Any,
     ):
@@ -137,10 +133,10 @@ def get_fused_embeddings(
         image_embeddings = None
 
         if texts is not None:
-            text_embeddings = self.get_text_embeddings(texts, batch_size)
+            text_embeddings = self.get_text_embeddings(texts, **kwargs)
 
         if images is not None:
-            image_embeddings = self.get_image_embeddings(images, batch_size)
+            image_embeddings = self.get_image_embeddings(images, **kwargs)
 
         if text_embeddings is not None and image_embeddings is not None:
             if len(text_embeddings) != len(image_embeddings):

mteb/models/vista_models.py

@@ -176,7 +176,6 @@ def get_fused_embeddings(
             self,
             texts: list[str] = None,
             images: list[Image.Image] | DataLoader = None,
-            *,
             task_name: str | None = None,
             prompt_type: PromptType | None = None,
             batch_size: int = 32,

mteb/models/vlm2vec_models.py

@@ -267,13 +267,16 @@ def get_fused_embeddings(
 
         text_embeddings = None
         image_embeddings = None
+        kwargs.update(
+            task_name=task_name, prompt_type=prompt_type, batch_size=batch_size
+        )
 
         if texts is not None and images is None:
-            text_embeddings = self.get_text_embeddings(texts, batch_size)
+            text_embeddings = self.get_text_embeddings(texts, **kwargs)
             return text_embeddings
 
         if images is not None and texts is None:
-            image_embeddings = self.get_image_embeddings(images, batch_size)
+            image_embeddings = self.get_image_embeddings(images, **kwargs)
             return image_embeddings
 
         # text_embeddings is not None and image_embeddings is not None

mteb/models/voyage_v.py

@@ -1,18 +1,17 @@
 from __future__ import annotations
 
+import logging
 import os
 from functools import partial
 from typing import Any
 
-import logging
 import torch
 from PIL import Image
 from torch.utils.data import DataLoader
 from torchvision import transforms
 from tqdm import tqdm
 
 import mteb
-from mteb.model_meta import ModelMeta
 from mteb.encoder_interface import PromptType
 from mteb.model_meta import ModelMeta
 
@@ -23,9 +22,7 @@
 def downsample_image(
     image: Image.Image, max_pixels: int = 16000000, target_longest_side: int = 4000
 ) -> Image.Image:
-    """
-    if image pixel > max_pixels, downsample it to target_longest_side while keeping the width height ratio.
-    """
+    """If image pixel > max_pixels, downsample it to target_longest_side while keeping the width height ratio."""
     width, height = image.size
     pixels = width * height