Add a check to import_utils.py to allow for use of faiss_gpu installation (#37997)

Adding check to import_utils.py for faiss_gpu
This commit is contained in:
Fiona Waters 2025-05-07 14:27:41 +01:00 committed by GitHub
parent fb1e3a4daa
commit c8607a17cb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -131,7 +131,7 @@ _cv2_available = importlib.util.find_spec("cv2") is not None
_yt_dlp_available = importlib.util.find_spec("yt_dlp") is not None
_datasets_available = _is_package_available("datasets")
_detectron2_available = _is_package_available("detectron2")
# We need to check both `faiss` and `faiss-cpu`.
# We need to check `faiss`, `faiss-cpu` and `faiss-gpu`.
_faiss_available = importlib.util.find_spec("faiss") is not None
try:
_faiss_version = importlib.metadata.version("faiss")
@ -141,7 +141,11 @@ except importlib.metadata.PackageNotFoundError:
_faiss_version = importlib.metadata.version("faiss-cpu")
logger.debug(f"Successfully imported faiss version {_faiss_version}")
except importlib.metadata.PackageNotFoundError:
_faiss_available = False
try:
_faiss_version = importlib.metadata.version("faiss-gpu")
logger.debug(f"Successfully imported faiss version {_faiss_version}")
except importlib.metadata.PackageNotFoundError:
_faiss_available = False
_ftfy_available = _is_package_available("ftfy")
_g2p_en_available = _is_package_available("g2p_en")
_hadamard_available = _is_package_available("fast_hadamard_transform")