feat-matryoshka-embeddings

configuration_clip.py

@@ -6,7 +6,7 @@
 
 import os
 from copy import deepcopy
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from transformers import PretrainedConfig, logging
 
@@ -157,6 +157,8 @@ class JinaCLIPConfig(PretrainedConfig):
         logit_scale_init_value: float = 2.6592,
         use_text_flash_attn: Optional[bool] = None,
         use_vision_xformers: Optional[bool] = None,
+        matryoshka_dimensions: Optional[List[int]] = None,
+        truncate_dim: Optional[int] = None,
         **kwargs,
     ):
         # If `_config_dict` exist, we use them for the backward compatibility.
@@ -167,6 +169,8 @@ class JinaCLIPConfig(PretrainedConfig):
         vision_config_dict: Optional[Dict] = kwargs.pop('vision_config_dict', None)
         self.use_text_flash_attn = use_text_flash_attn
         self.use_vision_xformers = use_vision_xformers
+        self.matryoshka_dimensions = matryoshka_dimensions
+        self.truncate_dim = truncate_dim
 
         super().__init__(**kwargs)
 

modeling_clip.py

@@ -4,12 +4,13 @@
 # https://github.com/huggingface/transformers/blob/main/src/transformers/models/clip/modeling_clip.py
 # and adjusted for Jina CLIP
 
+import base64
 from functools import partial
-from typing import List, Optional, Tuple, Union
 from io import BytesIO
-import requests
-import base64
+from typing import List, Optional, Tuple, Union
+
 import numpy as np
+import requests
 import torch
 import torch.nn.functional as f
 import torch.utils.checkpoint
@@ -39,9 +40,14 @@ except ImportError:
 from .configuration_clip import JinaCLIPConfig, JinaCLIPTextConfig, JinaCLIPVisionConfig
 from .eva_model import EVAVisionTransformer
 from .hf_model import HFTextEncoder
+
 # needed for HF to correctly import in cache
 from .rope_embeddings import VisionRotaryEmbeddingFast  # noqa: F401
-from .transform import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD, image_transform  # noqa: F401
+from .transform import (  # noqa: F401
+    OPENAI_DATASET_MEAN,
+    OPENAI_DATASET_STD,
+    image_transform,
+)
 
 logger = logging.get_logger(__name__)
 
@@ -280,6 +286,20 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
         )
         return self.visual_projection(self.vision_model(x=x))
 
+    def truncate_embeddings(self, embeddings, truncate_dim):
+        if not self.config.matryoshka_dimensions:
+            logger.warning(
+                "Matryoshka embeddings are not supported, so dimension truncation will not be performed."
+            )
+            return embeddings
+        elif truncate_dim in self.config.matryoshka_dimensions:
+            return embeddings[:, :truncate_dim]
+        else:
+            raise ValueError(
+                f"The provided `truncate_dim` value of {truncate_dim} is not supported. "
+                f"Supported dimensions are {self.config.matryoshka_dimensions}."
+            )
+
     @torch.inference_mode()
     def encode_text(
         self,
@@ -290,6 +310,7 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
         convert_to_tensor: bool = False,
         device: Optional[torch.device] = None,
         normalize_embeddings: bool = True,
+        truncate_dim: Optional[int] = None,
         **tokenizer_kwargs,
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
@@ -315,6 +336,8 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
                  If set to true, returned vectors will have length 1. In that case,
                  the faster dot-product (util.dot_score) instead of cosine similarity
                  can be used.
+             truncate_dim(`int`, *optional*, defaults to None):
+                The dimension to truncate sentence embeddings to. `None` does no truncation.
              tokenizer_kwargs(`Dict[str, Any]`, *optional*, defaults to {}):
                  Keyword arguments for the tokenizer
          Returns:
@@ -364,6 +387,7 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
         else:
             range_iter = range(0, len(sentences), batch_size)
 
+        truncate_dim = truncate_dim or self.config.truncate_dim
         for i in range_iter:
             encoded_input = self.tokenizer(
                 sentences[i : i + batch_size],
@@ -372,6 +396,9 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
             ).to(self.device)
 
             embeddings = self.get_text_features(input_ids=encoded_input)
+
+            if truncate_dim:
+                embeddings = self.truncate_embeddings(embeddings, truncate_dim)
             if normalize_embeddings:
                 embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
             if convert_to_numpy:
@@ -406,6 +433,7 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
         convert_to_tensor: bool = False,
         device: Optional[torch.device] = None,
         normalize_embeddings: bool = True,
+        truncate_dim: Optional[int] = None,
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
         Computes image embeddings.
@@ -431,6 +459,8 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
                 If set to true, returned vectors will have length 1. In that case,
                 the faster dot-product (util.dot_score) instead of cosine similarity
                 can be used.
+            truncate_dim(`int`, *optional*, defaults to None):
+                The dimension to truncate sentence embeddings to. `None` does no truncation.
         Returns:
             By default, a list of tensors is returned.
             If convert_to_tensor, a stacked tensor is returned.
@@ -476,7 +506,8 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
             range_iter = range(0, len(images), batch_size)
 
         from PIL import Image
-    
+
+        truncate_dim = truncate_dim or self.config.truncate_dim
         for i in range_iter:
             batch_images = images[i:i+batch_size]
             processed_inputs = []
@@ -501,6 +532,8 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
             processed_inputs = processed_inputs.to(self.device)
             embeddings = self.get_image_features(processed_inputs)
             
+            if truncate_dim:
+                embeddings = self.truncate_embeddings(embeddings, truncate_dim)
             if normalize_embeddings:
                 embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
             if convert_to_numpy: