fix: fix the bug when return_numpy is false

modeling_jina_embeddings_v4.py

@@ -334,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):
@@ -344,23 +346,18 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
                 ):
                     embeddings = self(**batch, task_label=task_label)
                     attention_mask = embeddings.attention_mask
-                    if not return_multivector:
+                    if return_multivector and not return_numpy:
                         embeddings = embeddings.single_vec_emb
                         if truncate_dim is not None:
                             embeddings = embeddings[:, :truncate_dim]
                     else:
                         embeddings = embeddings.multi_vec_emb
-                    if return_multivector and not return_numpy:
-                        # Get valid token mask from attention_mask
+                    if return_multivector:
                         valid_tokens = attention_mask.bool()
-                        # Remove padding by selecting only valid tokens for each sequence
                         embeddings = [emb[mask] for emb, mask in zip(embeddings, valid_tokens)]
-                        # Stack back into tensor with variable sequence lengths
                         results.append(embeddings)
                     else:
                         results.append(
-                            # If return_numpy is True, move embeddings to CPU for numpy conversion
-                            # Otherwise, unbind the tensor into a list of individual tensors along dim=0
                             embeddings.cpu()
                             if return_numpy
                             else list(torch.unbind(embeddings))
@@ -450,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]
@@ -498,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
 
@@ -515,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]