From cff142379c5ae6f3dc036cf2f626658055162478 Mon Sep 17 00:00:00 2001 From: yangdx Date: Tue, 1 Apr 2025 21:59:05 +0800 Subject: [PATCH] Add FAISS_USE_GPU env var to control which version of FAISS to be installed --- lightrag/kg/faiss_impl.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/lightrag/kg/faiss_impl.py b/lightrag/kg/faiss_impl.py index cf870b3a..c51eb1bf 100644 --- a/lightrag/kg/faiss_impl.py +++ b/lightrag/kg/faiss_impl.py @@ -11,16 +11,20 @@ import pipmaster as pm from lightrag.utils import logger, compute_mdhash_id from lightrag.base import BaseVectorStorage -if not pm.is_installed("faiss"): - pm.install("faiss") - -import faiss # type: ignore from .shared_storage import ( get_storage_lock, get_update_flag, set_all_update_flags, ) +import faiss # type: ignore + +USE_GPU = os.getenv("FAISS_USE_GPU", "0") == "1" +FAISS_PACKAGE = "faiss-gpu" if USE_GPU else "faiss-cpu" + +if not pm.is_installed(FAISS_PACKAGE): + pm.install(FAISS_PACKAGE) + @final @dataclass