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

Multi gpu support (#358)

Browse Source 
* feat: Added multi gpu support for text embedding

* feat: Add support for multi-gpu for special text models

* fix: Fix lazy_load to load the model to child processes when parallel is not none

* feat: Added lazy_load and multi-gpu to colbert

* feat: Add lazy_load and multi gpu to image models

* feat: Support lazy_load and multi-gpu to sparse models (except BM25)

* fix: Fixed BM25 not working

* refactor: Remove redundant GPUParallelProcessor

* refactor: Refactor _embed_*_parallel

* feat: Add cuda argument
refactor: Refactor how worker assign device

* fix: Fix if providers and cuda are None

* fix: Fix providers and cuda are none

* WIP: Multi gpu support review (#361)

* WIP: review

* wip: review

* refactor: refactor images

* refactor: refactor sparse

* refactor: refactor late interaction

* add model loading

* add tests

* fix: uncomment models in tests

* fix: fix variable declaration order

* fix: fix device id assignment

* tests: add multi gpu tests

* fix: fix device id assignment for sparse embeddings

* tests: update multi gpu tests

---------

Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>

* refactor: remove redundant declarations

* fix: rollback redundant changes

* fix: remove num workers device ids dep, fix type hint

* fix: fix post process for sparse models

* fix: remove redundant model loading

* new: add lazy load and new gpu support to cross encoders

* fix: add rerankers to multi gpu tests

* fix: unlock multilingual test

* fix: fix gpu test with cross encoder

---------

Co-authored-by: Andrey Vasnetsov <andrey@vasnetsov.com>
Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>
This commit is contained in:
Hossam Hagag
2024-10-17 00:42:38 +03:00
committed by GitHub
parent 58b5a8ed9a
commit eaecf7d471
29 changed files with 822 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
fastembed/common/onnx_model.py
View File

@@ -50,19 +50,28 @@ class OnnxModel(Generic[T]):
"""
return onnx_input
def load_onnx_model(
def _load_onnx_model(
self,
model_dir: Path,
model_file: str,
threads: Optional[int],
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_id: Optional[int] = None,
) -> None:
model_path = model_dir / model_file
# List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
onnx_providers = (
["CPUExecutionProvider"] if providers is None else list(providers)
)
if providers is not None:
onnx_providers = list(providers)
elif cuda:
if device_id is None:
onnx_providers = ["CUDAExecutionProvider"]
else:
onnx_providers = [("CUDAExecutionProvider", {"device_id": device_id})]
else:
onnx_providers = ["CPUExecutionProvider"]
available_providers = ort.get_available_providers()
requested_provider_names = []
for provider in onnx_providers:
@@ -94,6 +103,9 @@ class OnnxModel(Generic[T]):
RuntimeWarning,
)
def load_onnx_model(self) -> None:
raise NotImplementedError("Subclasses must implement this method")
def onnx_embed(self, *args, **kwargs) -> OnnxOutputContext:
raise NotImplementedError("Subclasses must implement this method")

12
fastembed/image/image_embedding.py
View File

@@ -45,21 +45,23 @@ class ImageEmbedding(ImageEmbeddingBase):
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
**kwargs,
):
super().__init__(model_name, cache_dir, threads, **kwargs)
for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
if any(
model_name.lower() == model["model"].lower()
for model in supported_models
):
if any(model_name.lower() == model["model"].lower() for model in supported_models):
self.model = EMBEDDING_MODEL_TYPE(
model_name,
cache_dir,
threads=threads,
providers=providers,
cuda=cuda,
device_ids=device_ids,
lazy_load=lazy_load,
**kwargs,
)
return

2
fastembed/image/image_embedding_base.py
View File

@@ -30,7 +30,7 @@ class ImageEmbeddingBase(ModelManagement):
Embeds a list of images into a list of embeddings.
Args:
images - The list of image paths to preprocess and embed.
images: The list of image paths to preprocess and embed.
batch_size: Batch size for encoding
parallel:
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.

59
fastembed/image/onnx_embedding.py
View File

@@ -59,6 +59,10 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
):
"""
@@ -68,24 +72,56 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
Defaults to `fastembed_cache` in the system's temp directory.
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.
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.
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.
Raises:
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
"""
super().__init__(model_name, cache_dir, threads, **kwargs)
self.providers = providers
self.lazy_load = lazy_load
model_description = self._get_model_description(model_name)
# List of device ids, that can be used for data parallel processing in workers
self.device_ids = device_ids
self.cuda = cuda
# This device_id will be used if we need to load model in current process
if device_id is not None:
self.device_id = device_id
elif self.device_ids is not None:
self.device_id = self.device_ids[0]
else:
self.device_id = None
self.model_description = self._get_model_description(model_name)
self.cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
model_description, self.cache_dir, local_files_only=self._local_files_only
self.model_description, self.cache_dir, local_files_only=self._local_files_only
)
self.load_onnx_model(
if not self.lazy_load:
self.load_onnx_model()
def load_onnx_model(self) -> None:
"""
Load the onnx model.
"""
self._load_onnx_model(
model_dir=self._model_dir,
model_file=model_description["model_file"],
threads=threads,
providers=providers,
model_file=self.model_description["model_file"],
threads=self.threads,
providers=self.providers,
cuda=self.cuda,
device_id=self.device_id,
)
@classmethod
@@ -120,12 +156,16 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
Returns:
List of embeddings, one per document
"""
yield from self._embed_images(
model_name=self.model_name,
cache_dir=str(self.cache_dir),
images=images,
batch_size=batch_size,
parallel=parallel,
providers=self.providers,
cuda=self.cuda,
device_ids=self.device_ids,
**kwargs,
)
@@ -148,4 +188,9 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
return OnnxImageEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)
return OnnxImageEmbedding(
model_name=model_name,
cache_dir=cache_dir,
threads=1,
**kwargs,
)

48
fastembed/image/onnx_image_model.py
View File

@@ -36,21 +36,28 @@ class OnnxImageModel(OnnxModel[T]):
"""
return onnx_input
def load_onnx_model(
def _load_onnx_model(
self,
model_dir: Path,
model_file: str,
threads: Optional[int],
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_id: Optional[int] = None,
) -> None:
super().load_onnx_model(
super()._load_onnx_model(
model_dir=model_dir,
model_file=model_file,
threads=threads,
providers=providers,
cuda=cuda,
device_id=device_id,
)
self.processor = load_preprocessor(model_dir=model_dir)
def load_onnx_model(self) -> None:
raise NotImplementedError("Subclasses must implement this method")
def _build_onnx_input(self, encoded: np.ndarray) -> Dict[str, np.ndarray]:
return {node.name: encoded for node in self.model.get_inputs()}
@@ -74,33 +81,44 @@ class OnnxImageModel(OnnxModel[T]):
images: ImageInput,
batch_size: int = 256,
parallel: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
**kwargs,
) -> Iterable[T]:
is_small = False
if (
isinstance(images, str)
or isinstance(images, Path)
or (isinstance(images, Image.Image))
):
if isinstance(images, (str, Path, Image.Image)):
images = [images]
is_small = True
if isinstance(images, list):
if len(images) < batch_size:
is_small = True
if parallel == 0:
parallel = os.cpu_count()
if isinstance(images, list) and len(images) < batch_size:
is_small = True
if parallel is None or is_small:
if not hasattr(self, "model") or self.model is None:
self.load_onnx_model()
for batch in iter_batch(images, batch_size):
yield from self._post_process_onnx_output(self.onnx_embed(batch))
else:
if parallel == 0:
parallel = os.cpu_count()
start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
params = {"model_name": model_name, "cache_dir": cache_dir, **kwargs}
params = {
"model_name": model_name,
"cache_dir": cache_dir,
"providers": providers,
**kwargs,
}
pool = ParallelWorkerPool(
parallel, self._get_worker_class(), start_method=start_method
parallel,
self._get_worker_class(),
cuda=cuda,
device_ids=device_ids,
start_method=start_method,
)
for batch in pool.ordered_map(iter_batch(images, batch_size), **params):
yield from self._post_process_onnx_output(batch)

62
fastembed/late_interaction/colbert.py
View File

@@ -117,6 +117,10 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
):
"""
@@ -126,29 +130,60 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
Defaults to `fastembed_cache` in the system's temp directory.
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.
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.
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.
Raises:
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
"""
super().__init__(model_name, cache_dir, threads, **kwargs)
self.providers = providers
self.lazy_load = lazy_load
model_description = self._get_model_description(model_name)
# List of device ids, that can be used for data parallel processing in workers
self.device_ids = device_ids
self.cuda = cuda
# This device_id will be used if we need to load model in current process
if device_id is not None:
self.device_id = device_id
elif self.device_ids is not None:
self.device_id = self.device_ids[0]
else:
self.device_id = None
self.model_description = self._get_model_description(model_name)
self.cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
model_description, self.cache_dir, local_files_only=self._local_files_only
self.model_description, self.cache_dir, local_files_only=self._local_files_only
)
self.mask_token_id = None
self.pad_token_id = None
self.skip_list = set()
self.load_onnx_model(
if not self.lazy_load:
self.load_onnx_model()
def load_onnx_model(self) -> None:
self._load_onnx_model(
model_dir=self._model_dir,
model_file=model_description["model_file"],
threads=threads,
providers=providers,
model_file=self.model_description["model_file"],
threads=self.threads,
providers=self.providers,
cuda=self.cuda,
device_id=self.device_id,
)
self.mask_token_id = self.special_token_to_id["[MASK]"]
self.pad_token_id = self.tokenizer.padding["pad_id"]
self.skip_list = {
self.tokenizer.encode(symbol, add_special_tokens=False).ids[0]
for symbol in string.punctuation
@@ -182,6 +217,9 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
documents=documents,
batch_size=batch_size,
parallel=parallel,
providers=self.providers,
cuda=self.cuda,
device_ids=self.device_ids,
**kwargs,
)
@@ -189,6 +227,9 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
if isinstance(query, str):
query = [query]
if not hasattr(self, "model") or self.model is None:
self.load_onnx_model()
for text in query:
yield from self._post_process_onnx_output(
self.onnx_embed([text], is_doc=False), is_doc=False
@@ -201,4 +242,9 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
class ColbertEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Colbert:
return Colbert(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)
return Colbert(
model_name=model_name,
cache_dir=cache_dir,
threads=1,
**kwargs,
)

22
fastembed/late_interaction/late_interaction_text_embedding.py
View File

@@ -48,18 +48,24 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
**kwargs,
):
super().__init__(model_name, cache_dir, threads, **kwargs)
for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
if any(
model_name.lower() == model["model"].lower()
for model in supported_models
):
if any(model_name.lower() == model["model"].lower() for model in supported_models):
self.model = EMBEDDING_MODEL_TYPE(
model_name, cache_dir, threads, providers=providers, **kwargs
model_name,
cache_dir,
threads=threads,
providers=providers,
cuda=cuda,
device_ids=device_ids,
lazy_load=lazy_load,
**kwargs,
)
return
@@ -92,9 +98,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
"""
yield from self.model.embed(documents, batch_size, parallel, **kwargs)
def query_embed(
self, query: Union[str, Iterable[str]], **kwargs
) -> Iterable[np.ndarray]:
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
"""
Embeds queries

24
fastembed/parallel_processor.py
View File

@@ -8,6 +8,7 @@ from multiprocessing.process import BaseProcess
from multiprocessing.sharedctypes import Synchronized as BaseValue
from queue import Empty
from typing import Any, Dict, Iterable, List, Optional, Tuple, Type
from copy import deepcopy
# Single item should be processed in less than:
processing_timeout = 10 * 60 # seconds
@@ -47,7 +48,9 @@ def _worker(
if kwargs is None:
kwargs = {}
logging.info(f"Reader worker: {worker_id} PID: {os.getpid()}")
logging.info(
f"Reader worker: {worker_id} PID: {os.getpid()} Device: {kwargs.get('device_id', 'CPU')}"
)
try:
worker = worker_class.start(**kwargs)
@@ -85,7 +88,14 @@ def _worker(
class ParallelWorkerPool:
def __init__(self, num_workers: int, worker: Type[Worker], start_method: Optional[str] = None):
def __init__(
self,
num_workers: int,
worker: Type[Worker],
start_method: Optional[str] = None,
device_ids: Optional[List[int]] = None,
cuda: bool = False,
):
self.worker_class = worker
self.num_workers = num_workers
self.input_queue: Optional[Queue] = None
@@ -94,6 +104,8 @@ class ParallelWorkerPool:
self.processes: List[BaseProcess] = []
self.queue_size = self.num_workers * max_internal_batch_size
self.emergency_shutdown = False
self.device_ids = device_ids
self.cuda = cuda
self.num_active_workers: Optional[BaseValue] = None
def start(self, **kwargs: Any) -> None:
@@ -105,6 +117,12 @@ class ParallelWorkerPool:
self.num_active_workers = ctx_value
for worker_id in range(0, self.num_workers):
worker_kwargs = deepcopy(kwargs)
if self.device_ids:
device_id = self.device_ids[worker_id % len(self.device_ids)]
worker_kwargs["device_id"] = device_id
worker_kwargs["cuda"] = self.cuda
assert hasattr(self.ctx, "Process")
process = self.ctx.Process(
target=_worker,
@@ -114,7 +132,7 @@ class ParallelWorkerPool:
self.output_queue,
self.num_active_workers,
worker_id,
kwargs.copy(),
worker_kwargs,
),
)
process.start()

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

@@ -1,5 +1,7 @@
from typing import List, Iterable, Dict, Any, Sequence, Optional
from loguru import logger
from fastembed.common import OnnxProvider
from fastembed.rerank.cross_encoder.onnx_text_model import OnnxCrossEncoderModel
from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
@@ -52,6 +54,10 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
):
"""
@@ -61,30 +67,58 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
Defaults to `fastembed_cache` in the system's temp directory.
threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
providers (Optional[Sequence[OnnxProvider]]): The list of providers to use for the onnxruntime session.
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.
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.
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.
Raises:
ValueError: If the model_name is not in the format <org>/<model> e.g. Xenova/ms-marco-MiniLM-L-6-v2.
"""
super().__init__(
model_name=model_name,
cache_dir=cache_dir,
threads=threads,
providers=providers,
**kwargs,
)
super().__init__(model_name, cache_dir, threads, **kwargs)
self.providers = providers
self.lazy_load = lazy_load
model_description = self._get_model_description(model_name)
# List of device ids, that can be used for data parallel processing in workers
self.device_ids = device_ids
self.cuda = cuda
if self.device_ids is not None and len(self.device_ids) > 1:
logger.warning(
"Parallel execution is currently not supported for cross encoders, "
f"only the first device will be used for inference: {self.device_ids[0]}."
)
# This device_id will be used if we need to load model in current process
if device_id is not None:
self.device_id = device_id
elif self.device_ids is not None:
self.device_id = self.device_ids[0]
else:
self.device_id = None
self.model_description = self._get_model_description(model_name)
self.cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
model_description, self.cache_dir, local_files_only=self._local_files_only
self.model_description, self.cache_dir, local_files_only=self._local_files_only
)
self.load_onnx_model(
if not self.lazy_load:
self.load_onnx_model()
def load_onnx_model(self) -> None:
self._load_onnx_model(
model_dir=self._model_dir,
model_file=model_description["model_file"],
threads=threads,
providers=providers,
model_file=self.model_description["model_file"],
threads=self.threads,
providers=self.providers,
cuda=self.cuda,
device_id=self.device_id,
)
def rerank(

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

@@ -12,18 +12,22 @@ from fastembed.common.utils import iter_batch
class OnnxCrossEncoderModel(OnnxModel):
ONNX_OUTPUT_NAMES: Optional[List[str]] = None
def load_onnx_model(
def _load_onnx_model(
self,
model_dir: Path,
model_file: str,
threads: Optional[int],
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_id: Optional[int] = None,
) -> None:
super().load_onnx_model(
super()._load_onnx_model(
model_dir=model_dir,
model_file=model_file,
threads=threads,
providers=providers,
cuda=cuda,
device_id=device_id,
)
self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
@@ -52,6 +56,8 @@ class OnnxCrossEncoderModel(OnnxModel):
def _rerank_documents(
self, query: str, documents: Iterable[str], batch_size: int, **kwargs
) -> Iterable[float]:
if not hasattr(self, "model") or self.model is None:
self.load_onnx_model()
for batch in iter_batch(documents, batch_size):
yield from self.onnx_embed(query, batch, **kwargs)

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

@@ -43,6 +43,9 @@ class TextCrossEncoder(TextCrossEncoderBase):
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
**kwargs,
):
super().__init__(model_name, cache_dir, threads, **kwargs)
@@ -55,6 +58,9 @@ class TextCrossEncoder(TextCrossEncoderBase):
cache_dir=cache_dir,
threads=threads,
providers=providers,
cuda=cuda,
device_ids=device_ids,
lazy_load=lazy_load,
**kwargs,
)
return

15
fastembed/sparse/bm25.py
View File

@@ -3,15 +3,12 @@ from collections import defaultdict
from multiprocessing import get_all_start_methods
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union
import mmh3
import numpy as np
from snowballstemmer import stemmer as get_stemmer
from fastembed.common.utils import (
define_cache_dir,
iter_batch,
get_all_punctuation,
remove_non_alphanumeric,
)
from fastembed.common.utils import define_cache_dir, iter_batch, get_all_punctuation, remove_non_alphanumeric
from fastembed.parallel_processor import ParallelWorkerPool, Worker
from fastembed.sparse.sparse_embedding_base import (
SparseEmbedding,
@@ -166,13 +163,13 @@ class Bm25(SparseTextEmbeddingBase):
if len(documents) < batch_size:
is_small = True
if parallel == 0:
parallel = os.cpu_count()
if parallel is None or is_small:
for batch in iter_batch(documents, batch_size):
yield from self.raw_embed(batch)
else:
if parallel == 0:
parallel = os.cpu_count()
start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
params = {
"model_name": model_name,

75
fastembed/sparse/bm42.py
View File

@@ -61,6 +61,10 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
alpha: float = 0.5,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
):
"""
@@ -74,39 +78,68 @@ 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.
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.
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.
Raises:
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
"""
super().__init__(model_name, cache_dir, threads, **kwargs)
self.providers = providers
self.lazy_load = lazy_load
model_description = self._get_model_description(model_name)
# List of device ids, that can be used for data parallel processing in workers
self.device_ids = device_ids
self.cuda = cuda
# This device_id will be used if we need to load model in current process
if device_id is not None:
self.device_id = device_id
elif self.device_ids is not None:
self.device_id = self.device_ids[0]
else:
self.device_id = None
self.model_description = self._get_model_description(model_name)
self.cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
model_description, self.cache_dir, local_files_only=self._local_files_only
)
self.load_onnx_model(
model_dir=self._model_dir,
model_file=model_description["model_file"],
threads=threads,
providers=providers,
self.model_description, self.cache_dir, local_files_only=self._local_files_only
)
self.invert_vocab = {}
for token, idx in self.tokenizer.get_vocab().items():
self.invert_vocab[idx] = token
self.special_tokens = set(self.special_token_to_id.keys())
self.special_tokens_ids = set(self.special_token_to_id.values())
self.special_tokens = set()
self.special_tokens_ids = set()
self.punctuation = set(string.punctuation)
self.stopwords = set(self._load_stopwords(self._model_dir))
self.stopwords = set()
self.stemmer = get_stemmer(MODEL_TO_LANGUAGE[model_name])
self.alpha = alpha
if not self.lazy_load:
self.load_onnx_model()
def load_onnx_model(self) -> None:
self._load_onnx_model(
model_dir=self._model_dir,
model_file=self.model_description["model_file"],
threads=self.threads,
providers=self.providers,
cuda=self.cuda,
device_id=self.device_id,
)
for token, idx in self.tokenizer.get_vocab().items():
self.invert_vocab[idx] = token
self.special_tokens = set(self.special_token_to_id.keys())
self.special_tokens_ids = set(self.special_token_to_id.values())
self.stopwords = set(self._load_stopwords(self._model_dir))
def _filter_pair_tokens(self, tokens: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
result = []
for token, value in tokens:
@@ -254,6 +287,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
documents=documents,
batch_size=batch_size,
parallel=parallel,
providers=self.providers,
cuda=self.cuda,
device_ids=self.device_ids,
alpha=self.alpha,
)
@@ -274,6 +310,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
if isinstance(query, str):
query = [query]
if not hasattr(self, "model") or self.model is None:
self.load_onnx_model()
for text in query:
encoded = self.tokenizer.encode(text)
document_tokens_with_ids = enumerate(encoded.tokens)
@@ -290,4 +329,8 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
class Bm42TextEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Bm42:
return Bm42(model_name=model_name, cache_dir=cache_dir, **kwargs)
return Bm42(
model_name=model_name,
cache_dir=cache_dir,
**kwargs,
)

6
fastembed/sparse/sparse_text_embedding.py
View File

@@ -48,6 +48,9 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
**kwargs,
):
super().__init__(model_name, cache_dir, threads, **kwargs)
@@ -68,6 +71,9 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
cache_dir,
threads=threads,
providers=providers,
cuda=cuda,
device_ids=device_ids,
lazy_load=lazy_load,
**kwargs,
)
return

56
fastembed/sparse/splade_pp.py
View File

@@ -64,6 +64,10 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
):
"""
@@ -73,24 +77,53 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
Defaults to `fastembed_cache` in the system's temp directory.
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.
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.
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.
Raises:
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
"""
super().__init__(model_name, cache_dir, threads, **kwargs)
self.providers = providers
self.lazy_load = lazy_load
model_description = self._get_model_description(model_name)
# List of device ids, that can be used for data parallel processing in workers
self.device_ids = device_ids
self.cuda = cuda
# This device_id will be used if we need to load model in current process
if device_id is not None:
self.device_id = device_id
elif self.device_ids is not None:
self.device_id = self.device_ids[0]
else:
self.device_id = None
self.model_description = self._get_model_description(model_name)
self.cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
model_description, self.cache_dir, local_files_only=self._local_files_only
self.model_description, self.cache_dir, local_files_only=self._local_files_only
)
self.load_onnx_model(
if not self.lazy_load:
self.load_onnx_model()
def load_onnx_model(self) -> None:
self._load_onnx_model(
model_dir=self._model_dir,
model_file=model_description["model_file"],
threads=threads,
providers=providers,
model_file=self.model_description["model_file"],
threads=self.threads,
providers=self.providers,
cuda=self.cuda,
device_id=self.device_id,
)
def embed(
@@ -121,6 +154,10 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
documents=documents,
batch_size=batch_size,
parallel=parallel,
providers=self.providers,
cuda=self.cuda,
device_ids=self.device_ids,
**kwargs,
)
@classmethod
@@ -130,4 +167,9 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
class SpladePPEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> SpladePP:
return SpladePP(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)
return SpladePP(
model_name=model_name,
cache_dir=cache_dir,
threads=1,
**kwargs,
)

14
fastembed/text/clip_embedding.py
View File

@@ -34,16 +34,20 @@ class CLIPOnnxEmbedding(OnnxTextEmbedding):
"""
return supported_clip_models
def _post_process_onnx_output(
self, output: OnnxOutputContext
) -> Iterable[np.ndarray]:
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
return output.model_output
class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):
def init_embedding(
self, model_name: str, cache_dir: str, **kwargs
self,
model_name: str,
cache_dir: str,
**kwargs,
) -> OnnxTextEmbedding:
return CLIPOnnxEmbedding(
model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
model_name=model_name,
cache_dir=cache_dir,
threads=1,
**kwargs,
)

10
fastembed/text/e5_onnx_embedding.py
View File

@@ -57,8 +57,14 @@ class E5OnnxEmbedding(OnnxTextEmbedding):
class E5OnnxEmbeddingWorker(OnnxTextEmbeddingWorker):
def init_embedding(
self, model_name: str, cache_dir: str, **kwargs
self,
model_name: str,
cache_dir: str,
**kwargs,
) -> E5OnnxEmbedding:
return E5OnnxEmbedding(
model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
model_name=model_name,
cache_dir=cache_dir,
threads=1,
**kwargs,
)

60
fastembed/text/onnx_embedding.py
View File

@@ -172,6 +172,10 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
):
"""
@@ -181,25 +185,43 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
Defaults to `fastembed_cache` in the system's temp directory.
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.
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.
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.
Raises:
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
"""
super().__init__(model_name, cache_dir, threads, **kwargs)
self.providers = providers
self.lazy_load = lazy_load
model_description = self._get_model_description(model_name)
# List of device ids, that can be used for data parallel processing in workers
self.device_ids = device_ids
self.cuda = cuda
# This device_id will be used if we need to load model in current process
if device_id is not None:
self.device_id = device_id
elif self.device_ids is not None:
self.device_id = self.device_ids[0]
else:
self.device_id = None
self.model_description = self._get_model_description(model_name)
self.cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
model_description, self.cache_dir, local_files_only=self._local_files_only
self.model_description, self.cache_dir, local_files_only=self._local_files_only
)
self.load_onnx_model(
model_dir=self._model_dir,
model_file=model_description["model_file"],
threads=threads,
providers=providers,
)
if not self.lazy_load:
self.load_onnx_model()
def embed(
self,
@@ -229,6 +251,9 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
documents=documents,
batch_size=batch_size,
parallel=parallel,
providers=self.providers,
cuda=self.cuda,
device_ids=self.device_ids,
**kwargs,
)
@@ -248,6 +273,16 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
embeddings = output.model_output
return normalize(embeddings[:, 0]).astype(np.float32)
def load_onnx_model(self) -> None:
self._load_onnx_model(
model_dir=self._model_dir,
model_file=self.model_description["model_file"],
threads=self.threads,
providers=self.providers,
cuda=self.cuda,
device_id=self.device_id,
)
class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(
@@ -256,4 +291,9 @@ class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
cache_dir: str,
**kwargs,
) -> OnnxTextEmbedding:
return OnnxTextEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)
return OnnxTextEmbedding(
model_name=model_name,
cache_dir=cache_dir,
threads=1,
**kwargs,
)

40
fastembed/text/onnx_text_model.py
View File

@@ -36,21 +36,28 @@ class OnnxTextModel(OnnxModel[T]):
"""
return onnx_input
def load_onnx_model(
def _load_onnx_model(
self,
model_dir: Path,
model_file: str,
threads: Optional[int],
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_id: Optional[int] = None,
) -> None:
super().load_onnx_model(
super()._load_onnx_model(
model_dir=model_dir,
model_file=model_file,
threads=threads,
providers=providers,
cuda=cuda,
device_id=device_id,
)
self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
def load_onnx_model(self) -> None:
raise NotImplementedError("Subclasses must implement this method")
def tokenize(self, documents: List[str], **kwargs) -> List[Encoding]:
return self.tokenizer.encode_batch(documents)
@@ -89,6 +96,9 @@ class OnnxTextModel(OnnxModel[T]):
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
**kwargs,
) -> Iterable[T]:
is_small = False
@@ -101,19 +111,29 @@ class OnnxTextModel(OnnxModel[T]):
if len(documents) < batch_size:
is_small = True
if parallel == 0:
parallel = os.cpu_count()
if parallel is None or is_small:
if not hasattr(self, "model") or self.model is None:
self.load_onnx_model()
for batch in iter_batch(documents, batch_size):
yield from self._post_process_onnx_output(self.onnx_embed(batch))
else:
start_method = (
"forkserver" if "forkserver" in get_all_start_methods() else "spawn"
)
params = {"model_name": model_name, "cache_dir": cache_dir, **kwargs}
if parallel == 0:
parallel = os.cpu_count()
start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
params = {
"model_name": model_name,
"cache_dir": cache_dir,
"providers": providers,
**kwargs,
}
pool = ParallelWorkerPool(
parallel, self._get_worker_class(), start_method=start_method
parallel,
self._get_worker_class(),
cuda=cuda,
device_ids=device_ids,
start_method=start_method,
)
for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
yield from self._post_process_onnx_output(batch)

23
fastembed/text/pooled_embedding.py
View File

@@ -3,7 +3,6 @@ from typing import Any, Dict, Iterable, List, Type
import numpy as np
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.utils import normalize
from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
from fastembed.text.onnx_text_model import TextEmbeddingWorker
@@ -47,14 +46,10 @@ class PooledEmbedding(OnnxTextEmbedding):
return PooledEmbeddingWorker
@classmethod
def mean_pooling(
cls, model_output: np.ndarray, attention_mask: np.ndarray
) -> np.ndarray:
def mean_pooling(cls, model_output: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
token_embeddings = model_output
input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
input_mask_expanded = np.tile(
input_mask_expanded, (1, 1, token_embeddings.shape[-1])
)
input_mask_expanded = np.tile(input_mask_expanded, (1, 1, token_embeddings.shape[-1]))
input_mask_expanded = input_mask_expanded.astype(float)
sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
sum_mask = np.sum(input_mask_expanded, axis=1)
@@ -70,9 +65,7 @@ class PooledEmbedding(OnnxTextEmbedding):
"""
return supported_pooled_models
def _post_process_onnx_output(
self, output: OnnxOutputContext
) -> Iterable[np.ndarray]:
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
embeddings = output.model_output
attn_mask = output.attention_mask
return self.mean_pooling(embeddings, attn_mask).astype(np.float32)
@@ -80,8 +73,14 @@ class PooledEmbedding(OnnxTextEmbedding):
class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):
def init_embedding(
self, model_name: str, cache_dir: str, **kwargs
self,
model_name: str,
cache_dir: str,
**kwargs,
) -> OnnxTextEmbedding:
return PooledEmbedding(
model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
model_name=model_name,
cache_dir=cache_dir,
threads=1,
**kwargs,
)

14
fastembed/text/pooled_normalized_embedding.py
View File

@@ -69,9 +69,7 @@ class PooledNormalizedEmbedding(PooledEmbedding):
"""
return supported_pooled_normalized_models
def _post_process_onnx_output(
self, output: OnnxOutputContext
) -> Iterable[np.ndarray]:
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
embeddings = output.model_output
attn_mask = output.attention_mask
return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
@@ -79,8 +77,14 @@ class PooledNormalizedEmbedding(PooledEmbedding):
class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):
def init_embedding(
self, model_name: str, cache_dir: str, **kwargs
self,
model_name: str,
cache_dir: str,
**kwargs,
) -> OnnxTextEmbedding:
return PooledNormalizedEmbedding(
model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
model_name=model_name,
cache_dir=cache_dir,
threads=1,
**kwargs,
)

16
fastembed/text/text_embedding.py
View File

@@ -55,21 +55,23 @@ class TextEmbedding(TextEmbeddingBase):
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[List[int]] = None,
lazy_load: bool = False,
**kwargs,
):
super().__init__(model_name, cache_dir, threads, **kwargs)
for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
if any(
model_name.lower() == model["model"].lower()
for model in supported_models
):
if any(model_name.lower() == model["model"].lower() for model in supported_models):
self.model = EMBEDDING_MODEL_TYPE(
model_name,
cache_dir,
model_name=model_name,
cache_dir=cache_dir,
threads=threads,
providers=providers,
cuda=cuda,
device_ids=device_ids,
lazy_load=lazy_load,
**kwargs,
)
return

15
tests/test_attention_embeddings.py
View File

@@ -142,3 +142,18 @@ def test_special_characters(model_name):
if is_ci:
shutil.rmtree(model.model._model_dir)
@pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions"])
def test_lazy_load(model_name):
model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
assert not hasattr(model.model, "model")
docs = ["hello world", "flag embedding"]
list(model.embed(docs))
assert hasattr(model.model, "model")
model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
list(model.query_embed(docs))
model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
list(model.passage_embed(docs))

12
tests/test_image_onnx_embeddings.py
View File

@@ -103,3 +103,15 @@ def test_parallel_processing(n_dims, model_name):
assert np.allclose(embeddings, embeddings_3, atol=1e-3)
if is_ci:
shutil.rmtree(model.model._model_dir)
@pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
def test_lazy_load(model_name):
model = ImageEmbedding(model_name=model_name, lazy_load=True)
assert not hasattr(model.model, "model")
images = [
TEST_MISC_DIR / "image.jpeg",
str(TEST_MISC_DIR / "small_image.jpeg"),
]
list(model.embed(images))
assert hasattr(model.model, "model")

20
tests/test_late_interaction_embeddings.py
View File

@@ -1,6 +1,7 @@
import os
import shutil
import pytest
import numpy as np
from fastembed.late_interaction.late_interaction_text_embedding import (
@@ -174,3 +175,22 @@ def test_parallel_processing():
if is_ci:
shutil.rmtree(model.model._model_dir)
@pytest.mark.parametrize(
"model_name",
["colbert-ir/colbertv2.0"],
)
def test_lazy_load(model_name):
model = LateInteractionTextEmbedding(model_name=model_name, lazy_load=True)
assert not hasattr(model.model, "model")
docs = ["hello world", "flag embedding"]
list(model.embed(docs))
assert hasattr(model.model, "model")
model = LateInteractionTextEmbedding(model_name=model_name, lazy_load=True)
list(model.query_embed(docs))
model = LateInteractionTextEmbedding(model_name=model_name, lazy_load=True)
list(model.passage_embed(docs))

220
tests/test_multi_gpu.py Normal file
View File

@@ -0,0 +1,220 @@
import pytest
from fastembed import (
TextEmbedding,
SparseTextEmbedding,
LateInteractionTextEmbedding,
ImageEmbedding,
)
from fastembed.rerank.cross_encoder import TextCrossEncoder
from tests.config import TEST_MISC_DIR
CACHE_DIR = "../model_cache"
@pytest.mark.skip(reason="Requires a multi-gpu server")
@pytest.mark.parametrize("device_id", [None, 0, 1])
def test_gpu_via_providers(device_id):
docs = ["hello world", "flag embedding"]
device_id = device_id if device_id is not None else 0
providers = (
["CUDAExecutionProvider"]
if device_id is None
else [("CUDAExecutionProvider", {"device_id": device_id})]
)
embedding_model = TextEmbedding(
"sentence-transformers/all-MiniLM-L6-v2",
providers=providers,
cache_dir=CACHE_DIR,
)
list(embedding_model.embed(docs))
options = embedding_model.model.model.get_provider_options()
assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)
embedding_model = SparseTextEmbedding(
"prithvida/Splade_PP_en_v1",
providers=providers,
cache_dir=CACHE_DIR,
)
list(embedding_model.embed(docs))
options = embedding_model.model.model.get_provider_options()
assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)
embedding_model = SparseTextEmbedding(
"Qdrant/bm42-all-minilm-l6-v2-attentions",
providers=providers,
cache_dir=CACHE_DIR,
)
list(embedding_model.embed(docs))
options = embedding_model.model.model.get_provider_options()
assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)
embedding_model = LateInteractionTextEmbedding(
"colbert-ir/colbertv2.0",
providers=providers,
cache_dir=CACHE_DIR,
)
list(embedding_model.embed(docs))
options = embedding_model.model.model.get_provider_options()
assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)
embedding_model = ImageEmbedding(
model_name="Qdrant/clip-ViT-B-32-vision",
providers=providers,
cache_dir=CACHE_DIR,
)
images = [
TEST_MISC_DIR / "image.jpeg",
str(TEST_MISC_DIR / "small_image.jpeg"),
]
list(embedding_model.embed(images))
options = embedding_model.model.model.get_provider_options()
assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)
model = TextCrossEncoder(
model_name="Xenova/ms-marco-MiniLM-L-6-v2",
providers=providers,
cache_dir=CACHE_DIR,
)
query = "What is the capital of France?"
documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
list(model.rerank(query, documents))
options = embedding_model.model.model.get_provider_options()
assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)
@pytest.mark.skip(reason="Requires a multi-gpu server")
@pytest.mark.parametrize("device_ids", [None, [0], [1], [0, 1]])
def test_gpu_cuda_device_ids(device_ids):
docs = ["hello world", "flag embedding"]
device_id = device_ids[0] if device_ids else 0
embedding_model = TextEmbedding(
"sentence-transformers/all-MiniLM-L6-v2",
cuda=True,
device_ids=device_ids,
cache_dir=CACHE_DIR,
)
list(embedding_model.embed(docs))
options = embedding_model.model.model.get_provider_options()
assert options["CUDAExecutionProvider"]["device_id"] == str(
device_id
), f"Text embedding: {options}"
embedding_model = SparseTextEmbedding(
"prithvida/Splade_PP_en_v1",
cuda=True,
device_ids=device_ids,
cache_dir=CACHE_DIR,
)
list(embedding_model.embed(docs))
options = embedding_model.model.model.get_provider_options()
assert options["CUDAExecutionProvider"]["device_id"] == str(
device_id
), f"Sparse text embedding: {options}"
embedding_model = SparseTextEmbedding(
"Qdrant/bm42-all-minilm-l6-v2-attentions",
cuda=True,
device_ids=device_ids,
cache_dir=CACHE_DIR,
)
list(embedding_model.embed(docs))
options = embedding_model.model.model.get_provider_options()
assert options["CUDAExecutionProvider"]["device_id"] == str(device_id), f"Bm42: {options}"
embedding_model = LateInteractionTextEmbedding(
"colbert-ir/colbertv2.0",
cuda=True,
device_ids=device_ids,
cache_dir=CACHE_DIR,
)
list(embedding_model.embed(docs))
options = embedding_model.model.model.get_provider_options()
assert options["CUDAExecutionProvider"]["device_id"] == str(
device_id
), f"Late interaction text embedding: {options}"
embedding_model = ImageEmbedding(
model_name="Qdrant/clip-ViT-B-32-vision",
cuda=True,
device_ids=device_ids,
cache_dir=CACHE_DIR,
)
images = [
TEST_MISC_DIR / "image.jpeg",
str(TEST_MISC_DIR / "small_image.jpeg"),
]
list(embedding_model.embed(images))
options = embedding_model.model.model.get_provider_options()
assert options["CUDAExecutionProvider"]["device_id"] == str(
device_id
), f"Image embedding: {options}"
if device_ids is None or len(device_ids) == 1:
model = TextCrossEncoder(
model_name="Xenova/ms-marco-MiniLM-L-6-v2",
cuda=True,
device_ids=device_ids,
cache_dir=CACHE_DIR,
)
query = "What is the capital of France?"
documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
list(model.rerank(query, documents))
options = embedding_model.model.model.get_provider_options()
assert options["CUDAExecutionProvider"]["device_id"] == str(
device_id
), f"Text cross encoder: {options}"
@pytest.mark.skip(reason="Requires a multi-gpu server")
@pytest.mark.parametrize(
"device_ids,parallel", [(None, None), (None, 2), ([1], None), ([1], 1), ([1], 2), ([0, 1], 2)]
)
def test_multi_gpu_parallel_inference(device_ids, parallel):
docs = ["hello world", "flag embedding"] * 100
batch_size = 5
embedding_model = TextEmbedding(
"sentence-transformers/all-MiniLM-L6-v2",
cuda=True,
device_ids=device_ids,
cache_dir=CACHE_DIR,
lazy_load=True,
)
list(embedding_model.embed(docs, batch_size=batch_size, parallel=parallel))
embedding_model = SparseTextEmbedding(
"prithvida/Splade_PP_en_v1",
cuda=True,
device_ids=device_ids,
cache_dir=CACHE_DIR,
)
list(embedding_model.embed(docs, batch_size=batch_size, parallel=parallel))
embedding_model = SparseTextEmbedding(
"Qdrant/bm42-all-minilm-l6-v2-attentions",
cuda=True,
device_ids=device_ids,
cache_dir=CACHE_DIR,
)
list(embedding_model.embed(docs, batch_size=batch_size, parallel=parallel))
embedding_model = LateInteractionTextEmbedding(
"colbert-ir/colbertv2.0",
cuda=True,
device_ids=device_ids,
cache_dir=CACHE_DIR,
)
list(embedding_model.embed(docs, batch_size=batch_size, parallel=parallel))
embedding_model = ImageEmbedding(
model_name="Qdrant/clip-ViT-B-32-vision",
cuda=True,
device_ids=device_ids,
cache_dir=CACHE_DIR,
)
images = [
TEST_MISC_DIR / "image.jpeg",
str(TEST_MISC_DIR / "small_image.jpeg"),
] * 100
list(embedding_model.embed(images, batch_size=batch_size, parallel=parallel))

18
tests/test_sparse_embeddings.py
View File

@@ -149,3 +149,21 @@ def test_stem_case_insensitive_stopwords(bm25_instance):
# Assert
expected = ["quick", "brown", "fox", "test", "sentenc"]
assert result == expected, f"Expected {expected}, but got {result}"
@pytest.mark.parametrize(
"model_name",
["prithivida/Splade_PP_en_v1"],
)
def test_lazy_load(model_name):
model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
assert not hasattr(model.model, "model")
docs = ["hello world", "flag embedding"]
list(model.embed(docs))
assert hasattr(model.model, "model")
model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
list(model.query_embed(docs))
model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
list(model.passage_embed(docs))

13
tests/test_text_cross_encoder.py
View File

@@ -56,3 +56,16 @@ def test_batch_rerank(model_name):
), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
if is_ci:
shutil.rmtree(model.model._model_dir)
@pytest.mark.parametrize(
"model_name",
["Xenova/ms-marco-MiniLM-L-6-v2"],
)
def test_lazy_load(model_name):
model = TextCrossEncoder(model_name=model_name, lazy_load=True)
assert not hasattr(model.model, "model")
query = "What is the capital of France?"
documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
list(model.rerank(query, documents))
assert hasattr(model.model, "model")

19
tests/test_text_onnx_embeddings.py
View File

@@ -126,5 +126,24 @@ def test_parallel_processing(n_dims, model_name):
assert embeddings.shape == (200, n_dims)
assert np.allclose(embeddings, embeddings_2, atol=1e-3)
assert np.allclose(embeddings, embeddings_3, atol=1e-3)
if is_ci:
shutil.rmtree(model.model._model_dir)
@pytest.mark.parametrize(
"model_name",
["BAAI/bge-small-en-v1.5"],
)
def test_lazy_load(model_name):
model = TextEmbedding(model_name=model_name, lazy_load=True)
assert not hasattr(model.model, "model")
docs = ["hello world", "flag embedding"]
list(model.embed(docs))
assert hasattr(model.model, "model")
model = TextEmbedding(model_name=model_name, lazy_load=True)
list(model.query_embed(docs))
model = TextEmbedding(model_name=model_name, lazy_load=True)
list(model.passage_embed(docs))