jinaai
/

jina-code-embeddings-1.5b

Feature Extraction

sentence-transformers

text-generation

text-generation-inference

🇪🇺 Region: EU

Model card Files Files and versions

dariakryvosheieva commited on Aug 29

Commit

f9464e7

·

verified ·

1 Parent(s): 47838a0

Update README.md

Files changed (1) hide show

README.md +3 -3

README.md CHANGED Viewed

@@ -102,7 +102,7 @@ model = SentenceTransformer(
     model_kwargs={
         "torch_dtype": torch.bfloat16,
         "attn_implementation": "flash_attention_2",
-        "device_map": "auto"
     }
 )
@@ -122,8 +122,8 @@ document_embeddings = model.encode(documents, prompt_name="nl2code_document")
 # Compute the (cosine) similarity between the query and document embeddings
 similarity = model.similarity(query_embeddings, document_embeddings)
 print(similarity)
-# tensor([[0.8157, 0.1222],
-#         [0.1201, 0.5500]])
 ```
 </details>

     model_kwargs={
         "torch_dtype": torch.bfloat16,
         "attn_implementation": "flash_attention_2",
+        "device_map": "cuda"
     }
 )
 # Compute the (cosine) similarity between the query and document embeddings
 similarity = model.similarity(query_embeddings, document_embeddings)
 print(similarity)
+# tensor([[0.7650, 0.1131],
+#         [0.0938, 0.6607]])
 ```
 </details>