fix-image-pooling (#9)

- fix: image pooling (725b8ba6ba8cff17579843ca46e5eb21f7d5ea37)
- chore: remove prints (660fe4c4d743be00c7bdcb17c740414c21c53374)

modeling_jina_embeddings_v4.py

@@ -216,22 +216,21 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         Project the hidden states to single-vector embeddings.
         """
         if self._input_has_image(input_ids[0]):  # got document image
-            img_start_pos = torch.where(
-                input_ids[0] == self.config.vision_start_token_id
-            )[0][0]
-            img_end_pos = torch.where(input_ids[0] == self.config.vision_end_token_id)[
-                0
-            ][0]
-            pooled_output = (
-                hidden_states[0][img_start_pos : img_end_pos + 1]
-                .mean(dim=0)
-                .unsqueeze(0)
-            )
+            img_start_positions = torch.where(input_ids == self.config.vision_start_token_id)[1]
+            img_end_positions = torch.where(input_ids == self.config.vision_end_token_id)[1]
+            
+            batch_size, seq_len = input_ids.shape
+            position_indices = torch.arange(seq_len, device=input_ids.device).expand(batch_size, -1)
+            image_mask = (position_indices >= img_start_positions.unsqueeze(1)) & (position_indices <= img_end_positions.unsqueeze(1))
+            
+            masked_hidden_states = hidden_states * image_mask.unsqueeze(-1)
+            pooled_output = masked_hidden_states.sum(dim=1) / image_mask.sum(dim=1, keepdim=True)
 
         else:  # got query text
             pooled_output = torch.sum(
                 hidden_states * attention_mask.unsqueeze(-1), dim=1
             ) / torch.sum(attention_mask, dim=1, keepdim=True)
+
         single_vec_emb = self.single_vector_projector(pooled_output)
         return torch.nn.functional.normalize(single_vec_emb, dim=-1)
 
@@ -317,7 +316,6 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
                             embeddings = embeddings[:, :truncate_dim]
                     else:
                         embeddings = embeddings.multi_vec_emb
-
                     results.append(
                         embeddings.cpu()
                         if return_numpy