From 08b9eaaa6f267ae1b4863ade5bae7727238d99cc Mon Sep 17 00:00:00 2001
From: Vincent Emonet <vincent.emonet@gmail.com>
Date: Thu, 6 Feb 2025 21:31:23 +0100
Subject: [PATCH] community: improve FastEmbedEmbeddings support for ONNX
 execution provider (e.g. GPU)  (#29645)

I made a change to how was implemented the support for GPU in
`FastEmbedEmbeddings` to be more consistent with the existing
implementation `langchain-qdrant` sparse embeddings implementation

It is directly enabling to provide the list of ONNX execution providers:
https://github.com/langchain-ai/langchain/blob/master/libs/partners/qdrant/langchain_qdrant/fastembed_sparse.py#L15

It is a bit less clear to a user that just wants to enable GPU, but
gives more capabilities to work with other execution providers that are
not the `CUDAExecutionProvider`, and is more future proof

Sorry for the disturbance @ccurme

> Nice to see you just moved to `uv`! It is so much nicer to run
format/lint/test! No need to manually rerun the `poetry install` with
all required extras now
---
 .../embeddings/fastembed.py                   | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/libs/community/langchain_community/embeddings/fastembed.py b/libs/community/langchain_community/embeddings/fastembed.py
index 12499adb1cdfd..d46f9210607e5 100644
--- a/libs/community/langchain_community/embeddings/fastembed.py
+++ b/libs/community/langchain_community/embeddings/fastembed.py
@@ -1,6 +1,6 @@
 import importlib
 import importlib.metadata
-from typing import Any, Dict, List, Literal, Optional, cast
+from typing import Any, Dict, List, Literal, Optional, Sequence, cast
 
 import numpy as np
 from langchain_core.embeddings import Embeddings
@@ -65,11 +65,12 @@ class FastEmbedEmbeddings(BaseModel, Embeddings):
     Defaults to `None`.
     """
 
-    gpu: bool = False
-    """Enable the use of GPU through CUDA. This requires to install `fastembed-gpu`
+    providers: Optional[Sequence[Any]] = None
+    """List of ONNX execution providers. Use `["CUDAExecutionProvider"]` to enable the
+    use of GPU when generating embeddings. This requires to install `fastembed-gpu`
     instead of `fastembed`. See https://qdrant.github.io/fastembed/examples/FastEmbed_GPU
     for more details.
-    Defaults to False.
+    Defaults to `None`.
     """
 
     model: Any = None  # : :meta private:
@@ -83,8 +84,12 @@ def validate_environment(cls, values: Dict) -> Dict:
         max_length = values.get("max_length")
         cache_dir = values.get("cache_dir")
         threads = values.get("threads")
-        gpu = values.get("gpu")
-        pkg_to_install = "fastembed-gpu" if gpu else "fastembed"
+        providers = values.get("providers")
+        pkg_to_install = (
+            "fastembed-gpu"
+            if providers and "CUDAExecutionProvider" in providers
+            else "fastembed"
+        )
 
         try:
             fastembed = importlib.import_module("fastembed")
@@ -106,7 +111,7 @@ def validate_environment(cls, values: Dict) -> Dict:
             max_length=max_length,
             cache_dir=cache_dir,
             threads=threads,
-            providers=["CUDAExecutionProvider"] if gpu else None,
+            providers=providers,
         )
         return values