support flash attn in from_pretrained

configuration_clip.py

@@ -155,6 +155,7 @@ class JinaCLIPConfig(PretrainedConfig):
         add_projections: bool = False,
         projection_dim: int = 768,
         logit_scale_init_value: float = 2.6592,
+        use_flash_attn: bool = False,
         **kwargs,
     ):
         # If `_config_dict` exist, we use them for the backward compatibility.
@@ -163,6 +164,7 @@ 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_flash_attn = use_flash_attn
 
         super().__init__(**kwargs)
 
@@ -259,6 +261,9 @@ class JinaCLIPConfig(PretrainedConfig):
                 'with default values.'
             )
 
+        if use_flash_attn:
+            text_config.hf_model_config_kwargs.use_flash_attn = use_flash_attn
+
         self.text_config = JinaCLIPTextConfig(**text_config)
         self.vision_config = JinaCLIPVisionConfig(**vision_config)