ashrielbrian
/

jina-embeddings-v4

Visual Document Retrieval

sentence-transformers

feature-extraction

multimodal-embedding

multilingual-embedding

Text-to-Visual Document (T→VD) retrieval

sentence-similarity

Model card Files Files and versions

Brian Tang commited on 3 days ago

Commit

8f0a794

·

1 Parent(s): 49ebb9c

Adds flash attention check with the device type

Files changed (1) hide show

modeling_jina_embeddings_v4.py +2 -1

modeling_jina_embeddings_v4.py CHANGED Viewed

@@ -569,7 +569,8 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             kwargs["torch_dtype"] = "auto"
         kwargs["key_mapping"] = super()._checkpoint_conversion_mapping
-        if not is_flash_attn_2_available():
             kwargs["attn_implementation"] = "sdpa"
         base_model = super().from_pretrained(

             kwargs["torch_dtype"] = "auto"
         kwargs["key_mapping"] = super()._checkpoint_conversion_mapping
+        device = kwargs.get("device", "auto")
+        if not is_flash_attn_2_available() or device == "cpu":
             kwargs["attn_implementation"] = "sdpa"
         base_model = super().from_pretrained(