upstream-archive/fastembed
1
0
Fork 0
mirror of https://github.com/qdrant/fastembed.git synced 2026-07-23 11:20:51 -05:00
Code Issues Packages Projects Releases Wiki Activity

new: use cuda if available (#537)

Browse Source 
* new: use cuda if available

* fix: fix warning msg

* fix: add missing import
This commit is contained in:
George
2025-12-10 20:23:34 +07:00
committed by GitHub
parent b304a2aff0
commit 685fd9b5a1
23 changed files with 98 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
fastembed/common/onnx_model.py
View File

@@ -9,7 +9,7 @@ import onnxruntime as ort
from numpy.typing import NDArray
from tokenizers import Tokenizer
from fastembed.common.types import OnnxProvider, NumpyArray
from fastembed.common.types import OnnxProvider, NumpyArray, Device
from fastembed.parallel_processor import Worker
# Holds type of the embedding result
@@ -60,23 +60,28 @@ class OnnxModel(Generic[T]):
model_file: str,
threads: int | None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_id: int | None = None,
extra_session_options: dict[str, Any] | None = None,
) -> None:
model_path = model_dir / model_file
# List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
available_providers = ort.get_available_providers()
cuda_available = "CUDAExecutionProvider" in available_providers
explicit_cuda = cuda is True or cuda == Device.CUDA
if cuda and providers is not None:
if explicit_cuda and providers is not None:
warnings.warn(
f"`cuda` and `providers` are mutually exclusive parameters, cuda: {cuda}, providers: {providers}",
f"`cuda` and `providers` are mutually exclusive parameters, "
f"cuda: {cuda}, providers: {providers}. If you'd like to use providers, cuda should be one of "
f"[False, Device.CPU, Device.AUTO].",
category=UserWarning,
stacklevel=6,
)
if providers is not None:
onnx_providers = list(providers)
elif cuda:
elif explicit_cuda or (cuda == Device.AUTO and cuda_available):
if device_id is None:
onnx_providers = ["CUDAExecutionProvider"]
else:
@@ -84,7 +89,6 @@ class OnnxModel(Generic[T]):
else:
onnx_providers = ["CPUExecutionProvider"]
available_providers = ort.get_available_providers()
requested_provider_names: list[str] = []
for provider in onnx_providers:
# check providers available

9
fastembed/common/types.py
View File

@@ -1,11 +1,18 @@
from enum import Enum
from pathlib import Path
from typing import Any, TypeAlias
import numpy as np
from numpy.typing import NDArray
from PIL import Image
class Device(str, Enum):
CPU = "cpu"
CUDA = "cuda"
AUTO = "auto"
PathInput: TypeAlias = str | Path
ImageInput: TypeAlias = PathInput | Image.Image

4
fastembed/image/image_embedding.py
View File

@@ -1,7 +1,7 @@
from typing import Any, Iterable, Sequence, Type
from dataclasses import asdict
from fastembed.common.types import NumpyArray
from fastembed.common.types import NumpyArray, Device
from fastembed.common import ImageInput, OnnxProvider
from fastembed.image.image_embedding_base import ImageEmbeddingBase
from fastembed.image.onnx_embedding import OnnxImageEmbedding
@@ -51,7 +51,7 @@ class ImageEmbedding(ImageEmbeddingBase):
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
**kwargs: Any,

12
fastembed/image/onnx_embedding.py
View File

@@ -1,7 +1,7 @@
from typing import Any, Iterable, Sequence, Type
from fastembed.common.types import NumpyArray
from fastembed.common.types import NumpyArray, Device
from fastembed.common import ImageInput, OnnxProvider
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.utils import define_cache_dir, normalize
@@ -63,10 +63,11 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
def __init__(
self,
model_name: str,
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
device_id: int | None = None,
@@ -82,10 +83,11 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to False.
cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to Device.AUTO.
device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
with `providers`. Defaults to None.
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.

6
fastembed/image/onnx_image_model.py
View File

@@ -8,7 +8,7 @@ import numpy as np
from PIL import Image
from fastembed.image.transform.operators import Compose
from fastembed.common.types import NumpyArray
from fastembed.common.types import NumpyArray, Device
from fastembed.common import ImageInput, OnnxProvider
from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
from fastembed.common.preprocessor_utils import load_preprocessor
@@ -53,7 +53,7 @@ class OnnxImageModel(OnnxModel[T]):
model_file: str,
threads: int | None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_id: int | None = None,
extra_session_options: dict[str, Any] | None = None,
) -> None:
@@ -97,7 +97,7 @@ class OnnxImageModel(OnnxModel[T]):
batch_size: int = 256,
parallel: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
local_files_only: bool = False,
specific_model_path: str | None = None,

11
fastembed/late_interaction/colbert.py
View File

@@ -5,7 +5,7 @@ import numpy as np
from tokenizers import Encoding, Tokenizer
from fastembed.common.preprocessor_utils import load_tokenizer
from fastembed.common.types import NumpyArray
from fastembed.common.types import NumpyArray, Device
from fastembed.common import OnnxProvider
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.utils import define_cache_dir, iter_batch
@@ -143,7 +143,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
device_id: int | None = None,
@@ -159,10 +159,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to False.
cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to Device.AUTO.
device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
with `providers`. Defaults to None.
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.

4
fastembed/late_interaction/late_interaction_text_embedding.py
View File

@@ -2,7 +2,7 @@ from typing import Any, Iterable, Sequence, Type
from dataclasses import asdict
from fastembed.common.model_description import DenseModelDescription
from fastembed.common.types import NumpyArray
from fastembed.common.types import NumpyArray, Device
from fastembed.common import OnnxProvider
from fastembed.late_interaction.colbert import Colbert
from fastembed.late_interaction.jina_colbert import JinaColbert
@@ -54,7 +54,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
**kwargs: Any,

11
fastembed/late_interaction_multimodal/colpali.py
View File

@@ -5,7 +5,7 @@ from tokenizers import Encoding
from fastembed.common import OnnxProvider, ImageInput
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.types import NumpyArray
from fastembed.common.types import NumpyArray, Device
from fastembed.common.utils import define_cache_dir, iter_batch
from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
LateInteractionMultimodalEmbeddingBase,
@@ -49,7 +49,7 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
device_id: int | None = None,
@@ -65,10 +65,11 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to False.
cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to Device.AUTO.
device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
with `providers`. Defaults to None.
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.

4
fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py
View File

@@ -2,7 +2,7 @@ from typing import Any, Iterable, Sequence, Type
from dataclasses import asdict
from fastembed.common import OnnxProvider, ImageInput
from fastembed.common.types import NumpyArray
from fastembed.common.types import NumpyArray, Device
from fastembed.late_interaction_multimodal.colpali import ColPali
from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
@@ -57,7 +57,7 @@ class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase)
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
**kwargs: Any,

8
fastembed/late_interaction_multimodal/onnx_multimodal_model.py
View File

@@ -11,7 +11,7 @@ from tokenizers import Encoding, Tokenizer
from fastembed.common import OnnxProvider, ImageInput
from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
from fastembed.common.preprocessor_utils import load_tokenizer, load_preprocessor
from fastembed.common.types import NumpyArray
from fastembed.common.types import NumpyArray, Device
from fastembed.common.utils import iter_batch
from fastembed.image.transform.operators import Compose
from fastembed.parallel_processor import ParallelWorkerPool
@@ -62,7 +62,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
model_file: str,
threads: int | None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_id: int | None = None,
extra_session_options: dict[str, Any] | None = None,
) -> None:
@@ -120,7 +120,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
batch_size: int = 256,
parallel: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
local_files_only: bool = False,
specific_model_path: str | None = None,
@@ -191,7 +191,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
batch_size: int = 256,
parallel: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
local_files_only: bool = False,
specific_model_path: str | None = None,

3
fastembed/parallel_processor.py
View File

@@ -10,6 +10,7 @@ from multiprocessing.sharedctypes import Synchronized as BaseValue
from queue import Empty
from typing import Any, Iterable, Type
from fastembed.common.types import Device
# Single item should be processed in less than:
processing_timeout = 10 * 60 # seconds
@@ -95,7 +96,7 @@ class ParallelWorkerPool:
worker: Type[Worker],
start_method: str | None = None,
device_ids: list[int] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
):
self.worker_class = worker
self.num_workers = num_workers

3
fastembed/rerank/cross_encoder/custom_text_cross_encoder.py
View File

@@ -2,6 +2,7 @@ from typing import Sequence, Any
from fastembed.common import OnnxProvider
from fastembed.common.model_description import BaseModelDescription
from fastembed.common.types import Device
from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
@@ -14,7 +15,7 @@ class CustomTextCrossEncoder(OnnxTextCrossEncoder):
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
device_id: int | None = None,

10
fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py
View File

@@ -4,6 +4,7 @@ from loguru import logger
from fastembed.common import OnnxProvider
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.types import Device
from fastembed.common.utils import define_cache_dir
from fastembed.rerank.cross_encoder.onnx_text_model import (
OnnxCrossEncoderModel,
@@ -80,7 +81,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
device_id: int | None = None,
@@ -96,10 +97,11 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to False.
cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to Device.AUTO.
device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
with `providers`. Defaults to None.
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.

6
fastembed/rerank/cross_encoder/onnx_text_model.py
View File

@@ -12,7 +12,7 @@ from fastembed.common.onnx_model import (
OnnxOutputContext,
OnnxProvider,
)
from fastembed.common.types import NumpyArray
from fastembed.common.types import NumpyArray, Device
from fastembed.common.preprocessor_utils import load_tokenizer
from fastembed.common.utils import iter_batch
from fastembed.parallel_processor import ParallelWorkerPool
@@ -31,7 +31,7 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
model_file: str,
threads: int | None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_id: int | None = None,
extra_session_options: dict[str, Any] | None = None,
) -> None:
@@ -94,7 +94,7 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
batch_size: int,
parallel: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
local_files_only: bool = False,
specific_model_path: str | None = None,

3
fastembed/rerank/cross_encoder/text_cross_encoder.py
View File

@@ -2,6 +2,7 @@ from typing import Any, Iterable, Sequence, Type
from dataclasses import asdict
from fastembed.common import OnnxProvider
from fastembed.common.types import Device
from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
@@ -56,7 +57,7 @@ class TextCrossEncoder(TextCrossEncoderBase):
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
**kwargs: Any,

10
fastembed/sparse/bm42.py
View File

@@ -9,6 +9,7 @@ from py_rust_stemmers import SnowballStemmer
from fastembed.common import OnnxProvider
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.types import Device
from fastembed.common.utils import define_cache_dir
from fastembed.sparse.sparse_embedding_base import (
SparseEmbedding,
@@ -69,7 +70,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
alpha: float = 0.5,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
device_id: int | None = None,
@@ -87,10 +88,11 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
alpha (float, optional): Parameter, that defines the importance of the token weight in the document
versus the importance of the token frequency in the corpus. Defaults to 0.5, based on empirical testing.
It is recommended to only change this parameter based on training data for a specific dataset.
cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to False.
cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to Device.AUTO.
device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
with `providers`. Defaults to None.
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.

10
fastembed/sparse/minicoil.py
View File

@@ -10,6 +10,7 @@ from tokenizers import Tokenizer
from fastembed.common.model_description import SparseModelDescription, ModelSource
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common import OnnxProvider
from fastembed.common.types import Device
from fastembed.common.utils import define_cache_dir
from fastembed.sparse.sparse_embedding_base import (
SparseEmbedding,
@@ -78,7 +79,7 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
k: float = 1.2,
b: float = 0.75,
avg_len: float = 150.0,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
device_id: int | None = None,
@@ -98,10 +99,11 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
b (float, optional): The b parameter in the BM25 formula. Defines the importance of the document length.
Defaults to 0.75.
avg_len (float, optional): The average length of the documents in the corpus. Defaults to 150.0.
cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to False.
cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to Device.AUTO.
device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
with `providers`. Defaults to None.
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.

3
fastembed/sparse/sparse_text_embedding.py
View File

@@ -2,6 +2,7 @@ from typing import Any, Iterable, Sequence, Type
from dataclasses import asdict
from fastembed.common import OnnxProvider
from fastembed.common.types import Device
from fastembed.sparse.bm25 import Bm25
from fastembed.sparse.bm42 import Bm42
from fastembed.sparse.minicoil import MiniCOIL
@@ -56,7 +57,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
**kwargs: Any,

10
fastembed/sparse/splade_pp.py
View File

@@ -3,6 +3,7 @@ from typing import Any, Iterable, Sequence, Type
import numpy as np
from fastembed.common import OnnxProvider
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.types import Device
from fastembed.common.utils import define_cache_dir
from fastembed.sparse.sparse_embedding_base import (
SparseEmbedding,
@@ -73,7 +74,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
device_id: int | None = None,
@@ -89,10 +90,11 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to False.
cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to Device.
device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
with `providers`. Defaults to None.
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.

5
fastembed/text/custom_text_embedding.py
View File

@@ -1,5 +1,4 @@
from typing import Sequence, Any, Iterable
from dataclasses import dataclass
import numpy as np
@@ -11,7 +10,7 @@ from fastembed.common.model_description import (
DenseModelDescription,
)
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.types import NumpyArray
from fastembed.common.types import NumpyArray, Device
from fastembed.common.utils import normalize, mean_pooling
from fastembed.text.onnx_embedding import OnnxTextEmbedding
@@ -32,7 +31,7 @@ class CustomTextEmbedding(OnnxTextEmbedding):
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
device_id: int | None = None,

11
fastembed/text/onnx_embedding.py
View File

@@ -1,6 +1,6 @@
from typing import Any, Iterable, Sequence, Type
from fastembed.common.types import NumpyArray, OnnxProvider
from fastembed.common.types import NumpyArray, OnnxProvider, Device
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.utils import define_cache_dir, normalize
from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
@@ -202,7 +202,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
device_id: int | None = None,
@@ -218,10 +218,11 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to False.
cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
Defaults to Device.AUTO.
device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
with `providers`. Defaults to None.
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.

6
fastembed/text/onnx_text_model.py
View File

@@ -7,7 +7,7 @@ import numpy as np
from numpy.typing import NDArray
from tokenizers import Encoding, Tokenizer
from fastembed.common.types import NumpyArray, OnnxProvider
from fastembed.common.types import NumpyArray, OnnxProvider, Device
from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
from fastembed.common.preprocessor_utils import load_tokenizer
from fastembed.common.utils import iter_batch
@@ -52,7 +52,7 @@ class OnnxTextModel(OnnxModel[T]):
model_file: str,
threads: int | None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_id: int | None = None,
extra_session_options: dict[str, Any] | None = None,
) -> None:
@@ -108,7 +108,7 @@ class OnnxTextModel(OnnxModel[T]):
batch_size: int = 256,
parallel: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
local_files_only: bool = False,
specific_model_path: str | None = None,

4
fastembed/text/text_embedding.py
View File

@@ -2,7 +2,7 @@ import warnings
from typing import Any, Iterable, Sequence, Type
from dataclasses import asdict
from fastembed.common.types import NumpyArray, OnnxProvider
from fastembed.common.types import NumpyArray, OnnxProvider, Device
from fastembed.text.clip_embedding import CLIPOnnxEmbedding
from fastembed.text.custom_text_embedding import CustomTextEmbedding
from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
@@ -82,7 +82,7 @@ class TextEmbedding(TextEmbeddingBase):
cache_dir: str | None = None,
threads: int | None = None,
providers: Sequence[OnnxProvider] | None = None,
cuda: bool = False,
cuda: bool | Device = Device.AUTO,
device_ids: list[int] | None = None,
lazy_load: bool = False,
**kwargs: Any,