OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-26 18:26:48 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

enh: sentence transformers env vars

Browse Source 
Co-Authored-By: DrZoidberg09 <96449693+drzoidberg09@users.noreply.github.com>
This commit is contained in:
Timothy Jaeryang Baek
2025-04-24 01:55:18 +09:00
parent 3ec6652f99
commit 732d7aee70
2 changed files with 56 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
backend/open_webui/routers/retrieval.py
View File

@@ -91,7 +91,12 @@ from open_webui.env import (
SRC_LOG_LEVELS,
DEVICE_TYPE,
DOCKER,
SENTENCE_TRANSFORMERS_BACKEND,
SENTENCE_TRANSFORMERS_MODEL_KWARGS,
SENTENCE_TRANSFORMERS_CROSS_ENCODER_BACKEND,
SENTENCE_TRANSFORMERS_CROSS_ENCODER_MODEL_KWARGS,
)
from open_webui.constants import ERROR_MESSAGES
log = logging.getLogger(__name__)
@@ -118,6 +123,8 @@ def get_ef(
get_model_path(embedding_model, auto_update),
device=DEVICE_TYPE,
trust_remote_code=RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
backend=SENTENCE_TRANSFORMERS_BACKEND,
model_kwargs=SENTENCE_TRANSFORMERS_MODEL_KWARGS,
)
except Exception as e:
log.debug(f"Error loading SentenceTransformer: {e}")
@@ -151,6 +158,8 @@ def get_rf(
get_model_path(reranking_model, auto_update),
device=DEVICE_TYPE,
trust_remote_code=RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
backend=SENTENCE_TRANSFORMERS_CROSS_ENCODER_BACKEND,
model_kwargs=SENTENCE_TRANSFORMERS_CROSS_ENCODER_MODEL_KWARGS,
)
except Exception as e:
log.error(f"CrossEncoder: {e}")