Openai embedding fix to support jina-embeddings-v2 (#4642)

oobabooga · Nov 18, 2023 · af76fbe · af76fbe
1 parent baab894
commit af76fbe
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/extensions/openai/embeddings.py b/extensions/openai/embeddings.py
@@ -1,6 +1,7 @@
 import os
 
 import numpy as np
+from transformers import AutoModel
 
 from extensions.openai.errors import ServiceUnavailableError
 from extensions.openai.utils import debug_msg, float_list_to_base64
@@ -41,7 +42,12 @@ def load_embedding_model(model: str):
     global embeddings_device, embeddings_model
     try:
         print(f"Try embedding model: {model} on {embeddings_device}")
-        embeddings_model = SentenceTransformer(model, device=embeddings_device)
+        if 'jina-embeddings' in model:
+            embeddings_model = AutoModel.from_pretrained(model, trust_remote_code=True)  # trust_remote_code is needed to use the encode method
+            embeddings_model = embeddings_model.to(embeddings_device)
+        else:
+            embeddings_model = SentenceTransformer(model, device=embeddings_device)
+
         print(f"Loaded embedding model: {model}")
     except Exception as e:
         embeddings_model = None