refactor: multi-vec, st truncation, etc

Signed-off-by: jupyterjazz <saba.sturua@jina.ai>

config.json

@@ -55,6 +55,6 @@
   "vocab_size": 151936,
   "truncate_dim": null,
   "task_names": ["retrieval", "text-matching", "code"],
-  "matryoshka_dims": [128, 256, 512, 1024],
+  "matryoshka_dims": [128, 256, 512, 1024, 2048],
   "_attn_implementation": "flash_attention_2"
 }

custom_st.py

@@ -103,7 +103,7 @@ class Transformer(nn.Module):
         return encoding
 
     def forward(
-        self, features: Dict[str, torch.Tensor], task: Optional[str] = None
+        self, features: Dict[str, torch.Tensor], task: Optional[str] = None, truncate_dim: Optional[int] = None
     ) -> Dict[str, torch.Tensor]:
         self.model.eval()
 
@@ -136,8 +136,8 @@ class Transformer(nn.Module):
                     text_embeddings = self.model(
                         **text_batch, task_label=task
                     ).single_vec_emb
-                    if self.config.truncate_dim:
-                        text_embeddings = text_embeddings[:, : self.config.truncate_dim]
+                    if truncate_dim:
+                        text_embeddings = text_embeddings[:, : truncate_dim]
                         text_embeddings = torch.nn.functional.normalize(text_embeddings, p=2, dim=-1)
                 for i, embedding in enumerate(text_embeddings):
                     all_embeddings.append((text_indices[i], embedding))
@@ -154,8 +154,8 @@ class Transformer(nn.Module):
                     img_embeddings = self.model(
                         **image_batch, task_label=task
                     ).single_vec_emb
-                    if self.config.truncate_dim:
-                        img_embeddings = img_embeddings[:, : self.config.truncate_dim]
+                    if truncate_dim:
+                        img_embeddings = img_embeddings[:, : truncate_dim]
                         img_embeddings = torch.nn.functional.normalize(img_embeddings, p=2, dim=-1)
 
                 for i, embedding in enumerate(img_embeddings):

modeling_jina_embeddings_v4.py

@@ -127,13 +127,11 @@ 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):
@@ -314,7 +312,6 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             ),
             single_vec_emb=single_vec_emb,
             multi_vec_emb=multi_vec_emb,
-            attention_mask=attention_mask,
         )
 
     def _process_batches(
@@ -345,17 +342,18 @@ 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]
-                            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=-1)
                     else:
                         embeddings = embeddings.multi_vec_emb
+                    
                     if return_multivector and not return_numpy:
-                        valid_tokens = attention_mask.bool()
-                        embeddings = [emb[mask] for emb, mask in zip(embeddings, valid_tokens)]
+                        valid_tokens = batch["attention_mask"].bool()
+                        embeddings = [
+                            emb[mask] for emb, mask in zip(embeddings, valid_tokens)
+                        ]
                         results.append(embeddings)
                     else:
                         results.append(

modules.json

@@ -4,6 +4,6 @@
         "name": "transformer",
         "path": "",
         "type": "custom_st.Transformer",
-        "kwargs": ["task"]
+        "kwargs": ["task", "truncate_dim"]
     }
 ]