support use_flash_attn in from_pretrained (#2)

- support flash attn in from_pretrained (d7c984ce33a82aa27b9bb6cf4e6a0ef775577760)
- change use_flash_attn and add x_attention attribute (ab448a5fe4db0f489546be2a56a8fd0e64f73d5b)
- set use_flash_attn at different position (a3f5a6005182cd3d5a4be6a9695c09f3952cc0d5)
- remove imports used for testing (853dc7d429ec17e8c8b8a7778453062e4cbcff16)


Co-authored-by: Michael Günther <michael-guenther@users.noreply.huggingface.co>

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

@@ -213,6 +213,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