support use_flash_attn in from_pretrained

configuration_clip.py

@@ -155,6 +155,8 @@ class JinaCLIPConfig(PretrainedConfig):
         add_projections: bool = False,
         projection_dim: int = 768,
         logit_scale_init_value: float = 2.6592,
+        use_text_flash_attn: Optional[bool] = None,
+        use_vision_xformers: Optional[bool] = None,
         **kwargs,
     ):
         # If `_config_dict` exist, we use them for the backward compatibility.
@@ -163,6 +165,8 @@ class JinaCLIPConfig(PretrainedConfig):
 
         text_config_dict: Optional[Dict] = kwargs.pop('text_config_dict', None)
         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
 
         super().__init__(**kwargs)
 

modeling_clip.py

@@ -210,6 +210,11 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
         text_config = config.text_config
         vision_config = config.vision_config
 
+        if config.use_text_flash_attn is not None:
+            text_config.hf_model_config_kwargs['use_flash_attn'] = config.use_text_flash_attn
+        if config.use_vision_xformers is not None:
+            vision_config.x_attention = config.use_vision_xformers
+
         self.add_projections = config.add_projections
         self.projection_dim = config.projection_dim
         self.text_embed_dim = text_config.embed_dim