feat: matryoshka embeddings

Signed-off-by: jupyterjazz <[email protected]>

configuration_xlm_roberta.py

@@ -31,6 +31,7 @@ class XLMRobertaFlashConfig(PretrainedConfig):
             use_flash_attn=True,
             torch_dtype=None,
             emb_pooler=None,
+            matryoshka_dimensions=None,
             **kwargs,
     ):
         super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
@@ -59,6 +60,7 @@ class XLMRobertaFlashConfig(PretrainedConfig):
         self.lora_main_params_trainable = lora_main_params_trainable
         self.use_flash_attn = use_flash_attn
         self.emb_pooler = emb_pooler
+        self.matryoshka_dimensions = matryoshka_dimensions
         if torch_dtype and hasattr(torch, torch_dtype) and type(getattr(torch, torch_dtype)) is torch.dtype:
             self.torch_dtype = getattr(torch, torch_dtype)
         else:

modeling_xlm_roberta.py

@@ -452,6 +452,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         convert_to_tensor: bool = False,
         device: Optional[torch.device] = None,
         normalize_embeddings: bool = False,
+        truncate_dim: int = None,
         **tokenizer_kwargs,
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
@@ -481,6 +482,8 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
                 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:
@@ -575,6 +578,17 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
 
         all_embeddings = [all_embeddings[idx] for idx in inverse_permutation]
 
+        if truncate_dim:
+            if not self.config.matryoshka_dimension:
+                logger.warning(
+                    'Matryoshka embeddings are not supported, so dimension truncation will not be performed.'
+                )
+            elif truncate_dim in self.config.matryoshka_dimension:
+                all_embeddings = [tensor[:truncate_dim] for tensor in all_embeddings]
+            else:
+                raise ValueError(f'The provided `truncate_dim` value of {truncate_dim} is not supported. '
+                                 f'Supported dimensions are {self.config.matryoshka_dimension}.')
+
         if convert_to_tensor:
             all_embeddings = torch.stack(all_embeddings)
         elif convert_to_numpy: