Selaa lähdekoodia

refac: lazy load big dependencies to reduce min memory usage

Jun Siang Cheah 7 kuukautta sitten
vanhempi
commit
318f61161e

+ 2 - 3
backend/open_webui/apps/retrieval/main.py

@@ -15,9 +15,6 @@ from pydantic import BaseModel
 
 from open_webui.apps.retrieval.vector.connector import VECTOR_DB_CLIENT
 
-# Information retrieval models
-from open_webui.apps.retrieval.models.colbert import ColBERT
-
 # Document loaders
 from open_webui.apps.retrieval.loaders.main import Loader
 
@@ -192,6 +189,8 @@ def update_reranking_model(
     if reranking_model:
         if any(model in reranking_model for model in ["jinaai/jina-colbert-v2"]):
             try:
+                from open_webui.apps.retrieval.models.colbert import ColBERT
+
                 app.state.sentence_transformer_rf = ColBERT(
                     get_model_path(reranking_model, auto_update),
                     env="docker" if DOCKER else None,

+ 4 - 4
backend/open_webui/apps/retrieval/vector/connector.py

@@ -1,10 +1,10 @@
-from open_webui.apps.retrieval.vector.dbs.chroma import ChromaClient
-from open_webui.apps.retrieval.vector.dbs.milvus import MilvusClient
-
-
 from open_webui.config import VECTOR_DB
 
 if VECTOR_DB == "milvus":
+    from open_webui.apps.retrieval.vector.dbs.milvus import MilvusClient
+
     VECTOR_DB_CLIENT = MilvusClient()
 else:
+    from open_webui.apps.retrieval.vector.dbs.chroma import ChromaClient
+
     VECTOR_DB_CLIENT = ChromaClient()