From 318f61161ea53a28e91d32669bb75d08f5b14aa4 Mon Sep 17 00:00:00 2001 From: Jun Siang Cheah Date: Wed, 2 Oct 2024 22:18:42 +0100 Subject: [PATCH] refac: lazy load big dependencies to reduce min memory usage --- backend/open_webui/apps/retrieval/main.py | 5 ++--- backend/open_webui/apps/retrieval/vector/connector.py | 8 ++++---- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/backend/open_webui/apps/retrieval/main.py b/backend/open_webui/apps/retrieval/main.py index 2a79ac90f..e0e4ea341 100644 --- a/backend/open_webui/apps/retrieval/main.py +++ b/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, diff --git a/backend/open_webui/apps/retrieval/vector/connector.py b/backend/open_webui/apps/retrieval/vector/connector.py index 5b203271f..1f33b1721 100644 --- a/backend/open_webui/apps/retrieval/vector/connector.py +++ b/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()