Update modeling_minicpmv.py

modeling_minicpmv.py

@@ -425,9 +425,8 @@ def transform_image_mp(img_list, transform, device, max_workers=None):
 
 
 @dataclass
-class BaseModelOutputWithAttentionMask(ModelOutput):
-    last_hidden_state: torch.FloatTensor = None
-    attention_mask: Optional[torch.Tensor] = None
+class MiniCPMVEmbeddingOutput(ModelOutput):
+    reps: torch.FloatTensor = None
 
 class MiniCPMVEmbedding(MiniCPMV): # MiniCPMVEmbedding -> MiniCPMV ->  Ultimately a CausalLM -> last_hidden_state for information retrieval
     def fused_tokenize(
@@ -524,12 +523,19 @@ class MiniCPMVEmbedding(MiniCPMV): # MiniCPMVEmbedding -> MiniCPMV ->  Ultimatel
         )
         
         last_hidden_state = vlm_outputs.last_hidden_state
+
+        # pooling, weighted mean (same in training)
+        attention_mask = model_inputs["attention_mask"]
+        attention_mask_ = attention_mask * attention_mask.cumsum(dim=1) # [0,1,1,1,0,0] -> [0,1,2,3,0,0]
+        s = torch.sum(last_hidden_state * attention_mask_.unsqueeze(-1).float(), dim=1)
+        d = attention_mask_.sum(dim=1, keepdim=True).float()
+        reps = s / d
+
+        # normalize representation (same in training)
+        reps_normalized = F.normalize(reps, dim=1)
         
-        last_hidden_state_normalized = F.normalize(last_hidden_state, dim=1)
-        
-        return BaseModelOutputWithAttentionMask(
-            last_hidden_state=last_hidden_state_normalized,
-            attention_mask=model_inputs.attention_mask
+        return MiniCPMVEmbeddingOutput(
+            reps=reps_normalized
         )