gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add FAISS_USE_GPU env var to control which version of FAISS to be installed

Browse Source
This commit is contained in:
yangdx
2025-04-01 21:59:05 +08:00
parent 0ea55cb8cd
commit cff142379c
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
lightrag/kg/faiss_impl.py
View File

@@ -11,16 +11,20 @@ import pipmaster as pm
from lightrag.utils import logger, compute_mdhash_id from lightrag.utils import logger, compute_mdhash_id
from lightrag.base import BaseVectorStorage from lightrag.base import BaseVectorStorage
if not pm.is_installed("faiss"):
pm.install("faiss")
import faiss # type: ignore
from .shared_storage import ( from .shared_storage import (
get_storage_lock, get_storage_lock,
get_update_flag, get_update_flag,
set_all_update_flags, 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 @final
@dataclass @dataclass