fix: remove the padding tokens when a list of multivectors are returned

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(
@@ -340,12 +343,20 @@ 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
+                    if return_multivector:
+                        # Get valid token mask from attention_mask
+                        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
+                        embeddings = torch.stack(embeddings)
                     results.append(
                         embeddings.cpu()
                         if return_numpy