Mirrors/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2025-12-16 11:57:51 +01:00
Code Issues Packages Projects Releases Wiki Activity

refac: lazily load sentence_transformers to reduce start up memory usage

Browse Source
This commit is contained in:
Jun Siang Cheah
2024-07-01 08:13:56 +08:00
parent 17c684369e
commit a48ac6a209
2 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
backend/apps/rag/main.py
View File

@@ -48,8 +48,6 @@ import mimetypes
import uuid
import json
import sentence_transformers
from apps.webui.models.documents import (
Documents,
DocumentForm,
@@ -190,6 +188,8 @@ def update_embedding_model(
update_model: bool = False,
):
if embedding_model and app.state.config.RAG_EMBEDDING_ENGINE == "":
import sentence_transformers
app.state.sentence_transformer_ef = sentence_transformers.SentenceTransformer(
get_model_path(embedding_model, update_model),
device=DEVICE_TYPE,
@@ -204,6 +204,8 @@ def update_reranking_model(
update_model: bool = False,
):
if reranking_model:
import sentence_transformers
app.state.sentence_transformer_rf = sentence_transformers.CrossEncoder(
get_model_path(reranking_model, update_model),
device=DEVICE_TYPE,