feat-remove-paddings-0623 (#22)

- fix: remove the padding tokens when a list of multivectors are returned (ef1876f5b9dbe290d7a58ff16fc37367217d32c5)
- fix: fix the bug when return_numpy is false (77d5a29ef8ef2396f56106d8ed882e322b7dc9be)
- fix: fix the bug when return_numpy is false (205b18f42bb9bd3ee57bab31cdd1b3116b6d762b)
- fix: fix the bug when return_numpy is false (6bb8cf2b7575b11b19c8f9780efda6f0b1a61708)
- fix: fix the bug (3ad717f7eaca1d26701063a16cfbe5f40ebaf551)

modeling_jina_embeddings_v4.py

@@ -127,11 +127,13 @@ class JinaEmbeddingsV4ModelOutput:
         vlm_last_hidden_states (torch.Tensor, optional): Last hidden states of the VLM.
         single_vec_emb (torch.Tensor, optional): Single-vector embeddings.
         multi_vec_emb (torch.Tensor, optional): Multi-vector embeddings.
+        attention_mask (torch.Tensor, optional): Attention mask.
     """
 
     vlm_last_hidden_states: Optional[torch.Tensor] = None
     single_vec_emb: Optional[torch.Tensor] = None
     multi_vec_emb: Optional[torch.Tensor] = None
+    attention_mask: Optional[torch.Tensor] = None
 
 
 class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
@@ -312,6 +314,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             ),
             single_vec_emb=single_vec_emb,
             multi_vec_emb=multi_vec_emb,
+            attention_mask=attention_mask,
         )
 
     def _process_batches(
@@ -331,6 +334,8 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             shuffle=False,
             collate_fn=processor_fn,
         )
+        if return_multivector and len(data) > 1:
+            assert not return_numpy, "`return_numpy` is not supported when `return_multivector=True` and more than one data is encoded"
         results = []
         self.eval()
         for batch in tqdm(dataloader, desc=desc):
@@ -340,17 +345,23 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
                     device_type=torch.device(self.device).type, dtype=torch.bfloat16
                 ):
                     embeddings = self(**batch, task_label=task_label)
+                    attention_mask = embeddings.attention_mask
                     if not return_multivector:
                         embeddings = embeddings.single_vec_emb
                         if truncate_dim is not None:
                             embeddings = embeddings[:, :truncate_dim]
                     else:
                         embeddings = embeddings.multi_vec_emb
-                    results.append(
-                        embeddings.cpu()
-                        if return_numpy
-                        else list(torch.unbind(embeddings))
-                    )
+                    if return_multivector and not return_numpy:
+                        valid_tokens = attention_mask.bool()
+                        embeddings = [emb[mask] for emb, mask in zip(embeddings, valid_tokens)]
+                        results.append(embeddings)
+                    else:
+                        results.append(
+                            embeddings.cpu()
+                            if return_numpy
+                            else list(torch.unbind(embeddings))
+                        )
         if return_numpy:
             return np.concatenate([result.numpy() for result in results], axis=0)
         return [item for sublist in results for item in sublist]
@@ -436,6 +447,12 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         )
 
         return_list = isinstance(texts, list)
+
+        # If return_multivector is True and encoding multiple texts, ignore return_numpy
+        if return_multivector and return_list and len(texts) > 1:
+            if return_numpy:
+                print("Warning: `return_numpy` is ignored when `return_multivector=True` and `len(texts) > 1`")
+            return_numpy = False
         
         if isinstance(texts, str):
             texts = [texts]
@@ -484,7 +501,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             images: image(s) to encode, can be PIL Image(s), URL(s), or local file path(s)
             batch_size: Number of images to process at once
             return_multivector: Whether to return multi-vector embeddings instead of single-vector embeddings
-            return_numpy: Whether to return numpy arrays instead of torch tensors
+            return_numpy: Whether to return numpy arrays instead of torch tensors. If `return_multivector` is `True` and more than one image is encoded, this parameter is ignored.
             truncate_dim: Dimension to truncate embeddings to (128, 256, 512, or 1024)
             max_pixels: Maximum number of pixels to process per image
 
@@ -501,6 +518,12 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         
         return_list = isinstance(images, list)
 
+        # If return_multivector is True and encoding multiple images, ignore return_numpy
+        if return_multivector and return_list and len(images) > 1:
+            if return_numpy:
+                print("Warning: `return_numpy` is ignored when `return_multivector=True` and `len(images) > 1`")
+            return_numpy = False
+
         # Convert single image to list
         if isinstance(images, (str, Image.Image)):
             images = [images]