fix: Fix memory access by using ort value

feat: Added io binding
new: Added function to check for cuda enabled
2026-09-22 22:17:49 -05:00 · 2025-03-13 08:13:04 +02:00 · 2025-03-13 08:02:21 +02:00 · 2025-03-13 08:02:21 +02:00 · 2025-03-13 08:02:21 +02:00 · 2025-03-13 08:02:21 +02:00
9 changed files with 62 additions and 283 deletions
@@ -1,4 +1,5 @@
 name: Tests
+run-name: Tests (gpu)

 on:
  pull_request:
@@ -22,8 +23,6 @@ jobs:
          - '3.13.x'
        os:
          - ubuntu-latest
-          - macos-latest
-          - windows-latest

    runs-on: ${{ matrix.os }}

@@ -190,23 +190,6 @@ scores = list(encoder.rerank(query, documents))
 # [-11.48061752319336, 5.472434997558594]
 ```

-Text cross encoders can also be extended with models which are not in the list of supported models.
-
-```python
-from fastembed.rerank.cross_encoder import TextCrossEncoder 
-from fastembed.common.model_description import ModelSource
-
-TextCrossEncoder.add_custom_model(
-    model="Xenova/ms-marco-MiniLM-L-4-v2",
-    model_file="onnx/model.onnx",
-    sources=ModelSource(hf="Xenova/ms-marco-MiniLM-L-4-v2"),
-)
-model = TextCrossEncoder(model_name="Xenova/ms-marco-MiniLM-L-4-v2")
-scores = list(model.rerank_pairs(
-    [("What is AI?", "Artificial intelligence is ..."), ("What is ML?", "Machine learning is ..."),]
-))
-```
-
 ## ⚡️ FastEmbed on a GPU

 FastEmbed supports running on GPU devices.
@@ -5,12 +5,12 @@ import tempfile
 import unicodedata
 from pathlib import Path
 from itertools import islice
-from typing import Iterable, Optional, TypeVar
+from typing import Iterable, Optional, TypeVar, Sequence

 import numpy as np
 from numpy.typing import NDArray

-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, OnnxProvider

 T = TypeVar("T")

@@ -67,3 +67,18 @@ def get_all_punctuation() -> set[str]:

 def remove_non_alphanumeric(text: str) -> str:
    return re.sub(r"[^\w\s]", " ", text, flags=re.UNICODE)
+
+
+def is_cuda_enabled(cuda: bool, providers: Optional[Sequence[OnnxProvider]]) -> bool:
+    """
+    Check if CUDA is enabled based on the `cuda` and `providers` parameters
+    """
+    if cuda:
+        return True
+    if not providers:
+        return False
+    if isinstance(providers, str):
+        return "CUDAExecutionProvider" in providers
+    return isinstance(providers, (list, tuple)) and any(
+        isinstance(p, str) and "CUDAExecutionProvider" in p for p in providers
+    )
@@ -3,8 +3,6 @@ import os
 from collections import defaultdict
 from copy import deepcopy
 from enum import Enum
-from dataclasses import dataclass
-from multiprocessing import shared_memory, Manager, Lock
 from multiprocessing import Queue, get_context
 from multiprocessing.context import BaseContext
 from multiprocessing.process import BaseProcess
@@ -12,11 +10,6 @@ from multiprocessing.sharedctypes import Synchronized as BaseValue
 from queue import Empty
 from typing import Any, Iterable, Optional, Type

-import numpy as np
-from numpy.typing import NDArray
-
-from fastembed.common.types import NumpyArray
-

 # Single item should be processed in less than:
 processing_timeout = 10 * 60  # seconds
@@ -24,13 +17,6 @@ processing_timeout = 10 * 60  # seconds
 max_internal_batch_size = 200


-@dataclass
-class OnnxOutputContext:
-    model_output: NumpyArray
-    attention_mask: Optional[NDArray[np.int64]] = None
-    input_ids: Optional[NDArray[np.int64]] = None
-
-
 class QueueSignals(str, Enum):
    stop = "stop"
    confirm = "confirm"
@@ -46,54 +32,12 @@ class Worker:
        raise NotImplementedError()


-class SharedMemoryPool:
-    def __init__(self, lock: Lock):
-        self._lock = lock
-        self._pool: dict[str, tuple[shared_memory.SharedMemory, int, np.dtype]] = {}
-        self._free_buffers: list[str] = []
-
-    def allocate(self, size: int, dtype: np.dtype) -> tuple[shared_memory.SharedMemory, str]:
-        best_match = None
-        best_size = float("inf")
-        for buf_name in self._free_buffers:
-            shm, buf_size, buf_dtype = self._pool[buf_name]
-            # get best match for needed size
-            if buf_size >= size and buf_dtype == dtype and buf_size < best_size:
-                best_match = buf_name
-                best_size = buf_size
-            if best_match:
-                self._free_buffers.remove(best_match)
-                return self._pool[best_match][0], best_match
-            shm = shared_memory.SharedMemory(create=True, size=size)
-            self._pool[shm.name] = (shm, size, dtype)
-            return shm, shm.name
-        # if no match found, create new buffer
-        shm = shared_memory.SharedMemory(create=True, size=size)
-        self._pool[shm.name] = (shm, size, dtype)
-        return shm, shm.name
-
-    def release(self, name: str) -> None:
-        if name in self._pool and name not in self._free_buffers:
-            self._free_buffers.append(name)
-
-    def cleanup(self) -> None:
-        for shm, _, _ in self._pool.values():
-            shm.close()
-            shm.unlink()
-        self._pool.clear()
-        self._free_buffers.clear()
-
-    def __del__(self):
-        self.cleanup()
-
-
 def _worker(
    worker_class: Type[Worker],
    input_queue: Queue,
    output_queue: Queue,
    num_active_workers: BaseValue,
    worker_id: int,
-    shared_pool: SharedMemoryPool,
    kwargs: Optional[dict[str, Any]] = None,
 ) -> None:
    """
@@ -111,6 +55,7 @@ def _worker(
    try:
        worker = worker_class.start(**kwargs)

+        # Keep going until you get an item that's None.
        def input_queue_iterable() -> Iterable[Any]:
            while True:
                item = input_queue.get()
@@ -119,24 +64,7 @@ def _worker(
                yield item

        for processed_item in worker.process(input_queue_iterable()):
-            idx, output_context = processed_item
-            output_metadata = {}
-            for field in ["model_output", "attention_mask", "input_ids"]:
-                array = getattr(output_context, field, None)
-                if array is not None:
-                    shm, shm_name = shared_pool.allocate(array.nbytes, array.dtype)
-                    shm_array = np.ndarray(array.shape, dtype=array.dtype, buffer=shm.buf)
-                    np.copyto(shm_array, array)
-                    output_metadata[field] = {
-                        "name": shm_name,
-                        "shape": array.shape,
-                        "dtype": array.dtype.str,
-                    }
-                    shm.close()
-            output_queue.put((idx, output_metadata))
-            for field in output_metadata:  # mark release to reuse
-                shared_pool.release(output_metadata[field]["name"])
-
+            output_queue.put(processed_item)
    except Exception as e:  # pylint: disable=broad-except
        logging.exception(e)
        output_queue.put(QueueSignals.error)
@@ -180,8 +108,6 @@ class ParallelWorkerPool:
        self.device_ids = device_ids
        self.cuda = cuda
        self.num_active_workers: Optional[BaseValue] = None
-        self.manager = Manager()
-        self.shared_pool = SharedMemoryPool(self.manager.Lock())

    def start(self, **kwargs: Any) -> None:
        self.input_queue = self.ctx.Queue(self.queue_size)
@@ -207,7 +133,6 @@ class ParallelWorkerPool:
                    self.output_queue,
                    self.num_active_workers,
                    worker_id,
-                    self.shared_pool,
                    worker_kwargs,
                ),
            )
@@ -253,17 +178,7 @@ class ParallelWorkerPool:
                    if out_item == QueueSignals.error:
                        self.join_or_terminate()
                        raise RuntimeError("Thread unexpectedly terminated")
-
-                    idx, output_metadata = out_item
-                    output_arrays = {}
-                    for field, meta in output_metadata.items():
-                        shm = shared_memory.SharedMemory(name=meta["name"])
-                        array = np.ndarray(
-                            meta["shape"], dtype=meta["dtype"], buffer=shm.buf
-                        ).copy()
-                        output_arrays[field] = array
-                        shm.close()
-                    yield (idx, OnnxOutputContext(**output_arrays))
+                    yield out_item
                    read += 1

                self.input_queue.put((idx, item))
@@ -278,18 +193,9 @@ class ParallelWorkerPool:
                if out_item == QueueSignals.error:
                    self.join_or_terminate()
                    raise RuntimeError("Thread unexpectedly terminated")
-
-                idx, output_metadata = out_item
-                output_arrays = {}
-                for field, meta in output_metadata.items():
-                    shm = shared_memory.SharedMemory(name=meta["name"])
-                    array = np.ndarray(meta["shape"], dtype=meta["dtype"], buffer=shm.buf).copy()
-                    output_arrays[field] = array
-                    shm.close()
-                yield (idx, OnnxOutputContext(**output_arrays))
+                yield out_item
                read += 1
        finally:
-            self.shared_pool.cleanup()
            assert self.input_queue is not None, "Input queue is None"
            assert self.output_queue is not None, "Output queue is None"
            self.join()
@@ -325,13 +231,11 @@ class ParallelWorkerPool:
            if process.is_alive():
                process.terminate()
        self.processes.clear()
-        self.shared_pool.cleanup()

    def join(self) -> None:
        for process in self.processes:
            process.join()
        self.processes.clear()
-        self.shared_pool.cleanup()

    def __del__(self) -> None:
        """
@@ -346,4 +250,3 @@ class ParallelWorkerPool:
        for process in self.processes:
            if process.is_alive():
                process.terminate()
-        self.shared_pool.cleanup()
@@ -1,46 +0,0 @@
-from typing import Optional, Sequence, Any
-
-from fastembed.common import OnnxProvider
-from fastembed.common.model_description import BaseModelDescription
-from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
-
-
-class CustomTextCrossEncoder(OnnxTextCrossEncoder):
-    SUPPORTED_MODELS: list[BaseModelDescription] = []
-
-    def __init__(
-        self,
-        model_name: str,
-        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,
-        specific_model_path: Optional[str] = None,
-        **kwargs: Any,
-    ):
-        super().__init__(
-            model_name=model_name,
-            cache_dir=cache_dir,
-            threads=threads,
-            providers=providers,
-            cuda=cuda,
-            device_ids=device_ids,
-            lazy_load=lazy_load,
-            device_id=device_id,
-            specific_model_path=specific_model_path,
-            **kwargs,
-        )
-
-    @classmethod
-    def _list_supported_models(cls) -> list[BaseModelDescription]:
-        return cls.SUPPORTED_MODELS
-
-    @classmethod
-    def add_model(
-        cls,
-        model_description: BaseModelDescription,
-    ) -> None:
-        cls.SUPPORTED_MODELS.append(model_description)
@@ -3,19 +3,13 @@ from dataclasses import asdict

 from fastembed.common import OnnxProvider
 from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
-from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
-
 from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
-from fastembed.common.model_description import (
-    ModelSource,
-    BaseModelDescription,
-)
+from fastembed.common.model_description import BaseModelDescription


 class TextCrossEncoder(TextCrossEncoderBase):
    CROSS_ENCODER_REGISTRY: list[Type[TextCrossEncoderBase]] = [
        OnnxTextCrossEncoder,
-        CustomTextCrossEncoder,
    ]

    @classmethod
@@ -130,34 +124,3 @@ class TextCrossEncoder(TextCrossEncoderBase):
        yield from self.model.rerank_pairs(
            pairs, batch_size=batch_size, parallel=parallel, **kwargs
        )
-
-    @classmethod
-    def add_custom_model(
-        cls,
-        model: str,
-        sources: ModelSource,
-        model_file: str = "onnx/model.onnx",
-        description: str = "",
-        license: str = "",
-        size_in_gb: float = 0.0,
-        additional_files: Optional[list[str]] = None,
-    ) -> None:
-        registered_models = cls._list_supported_models()
-        for registered_model in registered_models:
-            if model == registered_model.model:
-                raise ValueError(
-                    f"Model {model} is already registered in CrossEncoderModel, if you still want to add this model, "
-                    f"please use another model name"
-                )
-
-        CustomTextCrossEncoder.add_model(
-            BaseModelDescription(
-                model=model,
-                sources=sources,
-                model_file=model_file,
-                description=description,
-                license=license,
-                size_in_GB=size_in_gb,
-                additional_files=additional_files or [],
-            )
-        )
@@ -6,11 +6,12 @@ from typing import Any, Iterable, Optional, Sequence, Type, Union
 import numpy as np
 from numpy.typing import NDArray
 from tokenizers import Encoding, Tokenizer
+import onnxruntime as ort

 from fastembed.common.types import NumpyArray, OnnxProvider
 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
+from fastembed.common.utils import iter_batch, is_cuda_enabled
 from fastembed.parallel_processor import ParallelWorkerPool


@@ -62,6 +63,30 @@ class OnnxTextModel(OnnxModel[T]):
    def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
        return self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]

+    def _run_with_io_binding(
+        self, onnx_input: dict[str, NumpyArray], device_id: int = 0
+    ) -> NumpyArray:
+        """Run inference with IO Binding for optimized memory transfer for non-CPU execution providers."""
+        io_binding = self.model.io_binding()
+
+        for name, value in onnx_input.items():
+            ort_value = ort.OrtValue.ortvalue_from_numpy(
+                numpy_obj=value, device_type="cuda", device_id=device_id
+            )
+            io_binding.bind_ortvalue_input(name=name, ortvalue=ort_value)
+
+        output_names = [output.name for output in self.model.get_outputs()]
+        for output_name in output_names:
+            io_binding.bind_output(
+                name=output_name,
+                device_type="cuda",
+                device_id=device_id,
+            )
+
+        self.model.run_with_iobinding(io_binding)
+
+        return io_binding.copy_outputs_to_cpu()
+
    def onnx_embed(
        self,
        documents: list[str],
@@ -82,7 +107,13 @@ class OnnxTextModel(OnnxModel[T]):
            )
        onnx_input = self._preprocess_onnx_input(onnx_input, **kwargs)

-        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
+        providers = kwargs.get("providers", None)
+        cuda = kwargs.get("cuda", False)
+
+        if is_cuda_enabled(cuda, providers):
+            model_output = self._run_with_io_binding(onnx_input)
+        else:
+            model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
        return OnnxOutputContext(
            model_output=model_output[0],
            attention_mask=onnx_input.get("attention_mask", attention_mask),
@@ -115,7 +146,9 @@ class OnnxTextModel(OnnxModel[T]):
            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))
+                yield from self._post_process_onnx_output(
+                    self.onnx_embed(batch, cuda=cuda, providers=providers)
+                )
        else:
            if parallel == 0:
                parallel = os.cpu_count()
@@ -1,5 +1,5 @@
 [tool.poetry]
-name = "fastembed"
+name = "fastembed-gpu"
 version = "0.6.0"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
@@ -18,7 +18,7 @@ numpy = [
    { version = ">=2.1.0", python = ">=3.13" },
    { version = ">=1.21,<2.1.0", python = "<3.10" },
 ]
-onnxruntime = [
+onnxruntime-gpu = [
    { version = ">1.20.0", python = ">=3.13" },
    { version = ">=1.17.0,<1.20.0", python = "<3.10" },
    { version = ">=1.17.0,!=1.20.0", python = ">=3.10,<3.13" },
@@ -3,17 +3,10 @@ import os
 import numpy as np
 import pytest

-from fastembed.common.model_description import (
-    PoolingType,
-    ModelSource,
-    DenseModelDescription,
-    BaseModelDescription,
-)
+from fastembed.common.model_description import PoolingType, ModelSource, DenseModelDescription
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import normalize, mean_pooling
 from fastembed.text.custom_text_embedding import CustomTextEmbedding, PostprocessingConfig
-from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
-from fastembed.rerank.cross_encoder import TextCrossEncoder
 from fastembed.text.text_embedding import TextEmbedding
 from tests.utils import delete_model_cache

@@ -21,10 +14,8 @@ from tests.utils import delete_model_cache
@pytest.fixture(autouse=True)
 def restore_custom_models_fixture():
    CustomTextEmbedding.SUPPORTED_MODELS = []
-    CustomTextCrossEncoder.SUPPORTED_MODELS = []
    yield
    CustomTextEmbedding.SUPPORTED_MODELS = []
-    CustomTextCrossEncoder.SUPPORTED_MODELS = []


 def test_text_custom_model():
@@ -74,43 +65,6 @@ def test_text_custom_model():
        delete_model_cache(model.model._model_dir)


-def test_cross_encoder_custom_model():
-    is_ci = os.getenv("CI")
-    custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2"
-    size_in_gb = 0.08
-    source = ModelSource(hf=custom_model_name)
-    canonical_vector = np.array([-5.7170815, -11.112114], dtype=np.float32)
-
-    TextCrossEncoder.add_custom_model(
-        custom_model_name,
-        model_file="onnx/model.onnx",
-        sources=source,
-        size_in_gb=size_in_gb,
-    )
-
-    assert CustomTextCrossEncoder.SUPPORTED_MODELS[0] == BaseModelDescription(
-        model=custom_model_name,
-        sources=source,
-        model_file="onnx/model.onnx",
-        description="",
-        license="",
-        size_in_GB=size_in_gb,
-    )
-
-    model = TextCrossEncoder(custom_model_name)
-    pairs = [
-        ("What is AI?", "Artificial intelligence is ..."),
-        ("What is ML?", "Machine learning is ..."),
-    ]
-    scores = list(model.rerank_pairs(pairs))
-
-    embeddings = np.stack(scores, axis=0)
-    assert embeddings.shape == (2,)
-    assert np.allclose(embeddings, canonical_vector, atol=1e-3)
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
-
-
 def test_mock_add_custom_models():
    dim = 5
    size_in_gb = 0.1
@@ -202,28 +156,3 @@ def test_do_not_add_existing_model():
            dim=384,
            size_in_gb=0.47,
        )
-
-
-def test_do_not_add_existing_cross_encoder():
-    existing_base_model = "Xenova/ms-marco-MiniLM-L-6-v2"
-    custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2"
-
-    with pytest.raises(ValueError, match=f"Model {existing_base_model} is already registered"):
-        TextCrossEncoder.add_custom_model(
-            existing_base_model,
-            sources=ModelSource(hf=existing_base_model),
-            size_in_gb=0.08,
-        )
-
-    TextCrossEncoder.add_custom_model(
-        custom_model_name,
-        sources=ModelSource(hf=existing_base_model),
-        size_in_gb=0.08,
-    )
-
-    with pytest.raises(ValueError, match=f"Model {custom_model_name} is already registered"):
-        TextCrossEncoder.add_custom_model(
-            custom_model_name,
-            sources=ModelSource(hf=custom_model_name),
-            size_in_gb=0.08,
-        )
Author	SHA1	Message	Date
hh-space-invader	9320e6fb34	fix: Fix memory access by using ort value	2025-03-13 08:13:04 +02:00
hh-space-invader	51e6339e05	feat: Added io binding	2025-03-13 08:02:21 +02:00
hh-space-invader	8cfe93b1ba	new: Added function to check for cuda enabled	2025-03-13 08:02:21 +02:00
Hossam HagagandGeorge Panchuk	33d5ae2f5a	Speedup ci (#489 ) * chore: Trigger CI test * chore: Trigger CI test * chore: Trigger CI test * chore: Trigger CI test * chore: Trigger CI test * chore: Trigger CI test * chore: Trigger CI test * chore: Trigger CI test * chore: Trigger CI test * Trigger CI * Trigger CI * Trigger CI * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * Trigger CI test * new: Added on workflow dispatch * tests: Updated tests * fix: Fix CI * fix: Fix CI * fix: Fix CI * improve: Prevent stop iteration error caused by next * fix: Fix variable might be referenced before assignment * refactor: Revised the way of getting models to test * fix: Fix test in image model * refactor: Call one model * fix: Fix ci * fix: Fix splade model name * tests: Updated tests * chore: Remove cache * tests: Update multi task tests * tests: Update multi task tests * tests: Updated tests * refactor: refactor utils func, add comments, conditions refactor --------- Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>	2025-03-13 08:02:21 +02:00
George	3ee7c3fc66	new: preserve embeddings in a type set by their model (#492 ) * new: preserve embeddings in a type set by their model * fix: remove type coercion * fix: remove redundant type * fix: fix random data type in tests	2025-03-13 08:02:21 +02:00
George	eaa39bd4e3	Deprecate archive struct (#490 ) * bump version to 0.6.0 * new: deprecate gcp archive structure	2025-03-13 08:02:21 +02:00
George	0499ea5a39	new: gpu package (#224 )	2025-02-26 13:58:25 +01:00
George Panchuk	238e8cab8d	add eofl	2025-02-26 13:55:30 +01:00
George Panchuk	ad62c7d85e	sync publih with main	2025-02-26 13:55:30 +01:00
George Panchuk	ee5fa077cb	fix: workflow dispatch can only be triggered from the default branch	2025-02-26 13:55:30 +01:00
George Panchuk	9b1af0ce76	alter workflow	2025-02-26 13:55:30 +01:00
George Panchuk	b3461c0bb3	refactoring: alter workflow names	2025-02-26 13:55:30 +01:00
George Panchuk	9591396245	fix: do not run windows and mac os tests on gpu branch	2025-02-26 13:55:30 +01:00
George Panchuk	ac641d035f	new: gpu package publish workflow	2025-02-26 13:55:30 +01:00