Skip to content

Commit

Permalink
community: improve FastEmbedEmbeddings support for ONNX execution pro…
Browse files Browse the repository at this point in the history
…vider (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
  • Loading branch information
vemonet authored Feb 6, 2025
1 parent 1bf6202 commit 08b9eaa
Showing 1 changed file with 12 additions and 7 deletions.
19 changes: 12 additions & 7 deletions libs/community/langchain_community/embeddings/fastembed.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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")
Expand All @@ -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

Expand Down

0 comments on commit 08b9eaa

Please sign in to comment.