AI/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-10-01 01:26:03 -04:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
wizd 2023-11-19 07:24:29 +08:00 committed by GitHub
parent baab894759
commit af76fbedb8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
extensions/openai/embeddings.py
View File

@ -1,6 +1,7 @@
import os import os
import numpy as np import numpy as np
from transformers import AutoModel
from extensions.openai.errors import ServiceUnavailableError from extensions.openai.errors import ServiceUnavailableError
from extensions.openai.utils import debug_msg, float_list_to_base64 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 global embeddings_device, embeddings_model
try: try:
print(f"Try embedding model: {model} on {embeddings_device}") print(f"Try embedding model: {model} on {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) embeddings_model = SentenceTransformer(model, device=embeddings_device)
print(f"Loaded embedding model: {model}") print(f"Loaded embedding model: {model}")
except Exception as e: except Exception as e:
embeddings_model = None embeddings_model = None