remove-single-vector-projection (#18)

- refactor: remove single vec projection (5bb9539014f6cef2d3662dba50b75749e0922b50)

adapters/adapter_config.json

@@ -1,7 +1,7 @@
 {
   "alpha_pattern": {},
   "auto_mapping": null,
-  "base_model_name_or_path": "jinaai/colqwen25-duo-base",
+  "base_model_name_or_path": "jinaai/jina-embeddings-v4",
   "bias": "none",
   "corda_config": null,
   "eva_config": null,

adapters/adapter_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c9799872132988d3689a35300538fb97fc5b0e02c1c42f7afd914fd1d8b59a88
-size 360118024
+oid sha256:b6b7ab4a79daa3b4f3b5274500cc99d3dc89aa8c3419e9d79f89e366685e12e5
+size 359863776

config.json

@@ -33,7 +33,6 @@
   },
   "rope_theta": 1000000.0,
   "single_vector_pool_strategy": "mean",
-  "single_vector_projector_dim": 1024,
   "sliding_window": 32768,
   "tie_word_embeddings": true,
   "torch_dtype": "bfloat16",

configuration_jina_embeddings_v4.py

@@ -9,14 +9,12 @@ class JinaEmbeddingsV4Config(Qwen2_5_VLConfig):
 
     def __init__(
         self,
-        single_vector_projector_dim: int = 1024,
         single_vector_pool_strategy: str = "mean",
         multi_vector_projector_dim: int = 128,
         pretrained_peft_model_name_or_path: Optional[str] = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
-        self.single_vector_projector_dim = single_vector_projector_dim
         self.single_vector_pool_strategy = single_vector_pool_strategy
         self.multi_vector_projector_dim = multi_vector_projector_dim
         self.pretrained_peft_model_name_or_path = pretrained_peft_model_name_or_path

model-00001-of-00002.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6b45c7afe391b4d9cc49f1ed3f6976f4a25ed40aa2165ed2ae118ff549355985
-size 4997750760
+oid sha256:abb244162956ec2f26d944b6c10cbb96afe211d2aff908b8b2f498ec27a9100b
+size 4997750728

model-00002-of-00002.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a20083234b15a57f34207bb99589241cf7531f01c09fd657110712cb634a811a
-size 2516308496
+oid sha256:5d5252a7ede6469220b0e7386af53fea9a45fa299a1d2af6fe68cb29897de3e3
+size 2512111904

model.safetensors.index.json

@@ -439,8 +439,6 @@
     "model.norm.weight": "model-00002-of-00002.safetensors",
     "multi_vector_projector.bias": "model-00002-of-00002.safetensors",
     "multi_vector_projector.weight": "model-00002-of-00002.safetensors",
-    "single_vector_projector.bias": "model-00002-of-00002.safetensors",
-    "single_vector_projector.weight": "model-00002-of-00002.safetensors",
     "visual.blocks.0.attn.proj.bias": "model-00001-of-00002.safetensors",
     "visual.blocks.0.attn.proj.weight": "model-00001-of-00002.safetensors",
     "visual.blocks.0.attn.qkv.bias": "model-00001-of-00002.safetensors",

modeling_jina_embeddings_v4.py

@@ -141,12 +141,11 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
 
     def __init__(self, config: JinaEmbeddingsV4Config):
         Qwen2_5_VLForConditionalGeneration.__init__(self, config)
-        self._init_projection_layers(config)
+        self._init_projection_layer(config)
         self.post_init()
         self.processor = JinaEmbeddingsV4Processor.from_pretrained(
             self.name_or_path, trust_remote_code=True, use_fast=True
         )
-        self.single_vector_projector_dim = config.single_vector_projector_dim
         self.multi_vector_projector_dim = config.multi_vector_projector_dim
         self._task = None
 
@@ -204,32 +203,25 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
 
         return hidden_states[-1]
 
-    def _init_projection_layers(self, config) -> None:
+    def _init_projection_layer(self, config) -> None:
         """
         Initializes projection layers.
         """
-        self.config.single_vector_projector_dim = config.single_vector_projector_dim
         self.config.multi_vector_projector_dim = config.multi_vector_projector_dim
 
-        self.single_vector_projector = nn.Linear(
-            in_features=self.config.text_config.hidden_size,
-            out_features=self.config.single_vector_projector_dim,
-        )
-
         self.multi_vector_projector = nn.Linear(
             in_features=self.config.text_config.hidden_size,
             out_features=self.config.multi_vector_projector_dim,
         )
 
-    def project_to_single_vector_embeddings(
+    def get_single_vector_embeddings(
         self,
-        task_label: Union[str, List[str]],
         hidden_states: torch.Tensor,
         attention_mask: torch.Tensor,
         input_ids: Optional[torch.LongTensor] = None,
     ) -> torch.Tensor:
         """
-        Project the hidden states to single-vector embeddings.
+        Get the single-vector embeddings from the hidden states.
         """
         if self._input_has_image(input_ids[0]):  # got document image
             img_start_positions = torch.where(
@@ -257,12 +249,9 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
                 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, task_label=task_label
-        )
-        return torch.nn.functional.normalize(single_vec_emb, dim=-1)
+        return torch.nn.functional.normalize(pooled_output, dim=-1)
 
-    def project_to_multi_vector_embeddings(
+    def get_multi_vector_embeddings(
         self,
         task_label: Union[str, List[str]],
         hidden_states: torch.Tensor,
@@ -306,13 +295,12 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             **kwargs,
         )  # (batch_size, seq_length, hidden_size)
         # Compute the embeddings
-        single_vec_emb = self.project_to_single_vector_embeddings(
+        single_vec_emb = self.get_single_vector_embeddings(
             hidden_states=hidden_states,
             attention_mask=attention_mask,
             input_ids=input_ids,
-            task_label=task_label,
         )
-        multi_vec_emb = self.project_to_multi_vector_embeddings(
+        multi_vec_emb = self.get_multi_vector_embeddings(
             hidden_states=hidden_states,
             attention_mask=attention_mask,
             task_label=task_label,