feat: truncation option during init

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

configuration_xlm_roberta.py

@@ -32,6 +32,7 @@ class XLMRobertaFlashConfig(PretrainedConfig):
             torch_dtype=None,
             emb_pooler=None,
             matryoshka_dimensions=None,
+            truncate_dim=None,
             **kwargs,
     ):
         super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
@@ -61,6 +62,7 @@ class XLMRobertaFlashConfig(PretrainedConfig):
         self.use_flash_attn = use_flash_attn
         self.emb_pooler = emb_pooler
         self.matryoshka_dimensions = matryoshka_dimensions
+        self.truncate_dim = truncate_dim
         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

@@ -578,6 +578,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
 
         all_embeddings = [all_embeddings[idx] for idx in inverse_permutation]
 
+        truncate_dim = truncate_dim or self.config.truncate_dim
         if truncate_dim:
             all_embeddings = self.truncate_embeddings(all_embeddings, truncate_dim)