set use_flash_attn at different position

configuration_clip.py

@@ -263,11 +263,6 @@ class JinaCLIPConfig(PretrainedConfig):
                 'with default values.'
             )
 
-        if use_text_flash_attn:
-            text_config.hf_model_config_kwargs.use_flash_attn = use_text_flash_attn
-        if use_vision_xformers:
-            vision_config.x_attention = use_vision_xformers
-
         self.text_config = JinaCLIPTextConfig(**text_config)
         self.vision_config = JinaCLIPVisionConfig(**vision_config)
 

modeling_clip.py

@@ -39,6 +39,9 @@ except ImportError:
 from .configuration_clip import JinaCLIPConfig, JinaCLIPTextConfig, JinaCLIPVisionConfig
 from .eva_model import EVAVisionTransformer
 from .hf_model import HFTextEncoder
+from .rope_embeddings import rx
+from .transform import rt
+from .processing_clip import rp
 
 logger = logging.get_logger(__name__)
 
@@ -210,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