fastembed-gpu 0.3.6__tar.gz → 0.4.0__tar.gz

{fastembed_gpu-0.3.6 → fastembed_gpu-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fastembed-gpu
-Version: 0.3.6
+Version: 0.4.0
 Summary: Fast, light, accurate library built for retrieval embedding generation
 Home-page: https://github.com/qdrant/fastembed
 License: Apache License
@@ -18,7 +18,7 @@ Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: PyStemmer (>=2.2.0,<3.0.0)
 Requires-Dist: huggingface-hub (>=0.20,<1.0)
 Requires-Dist: loguru (>=0.7.2,<0.8.0)
-Requires-Dist: mmh3 (>=4.0,<5.0)
+Requires-Dist: mmh3 (>=4.1.0,<5.0.0)
 Requires-Dist: numpy (>=1.21,<2) ; python_version < "3.12"
 Requires-Dist: numpy (>=1.26,<2) ; python_version >= "3.12"
 Requires-Dist: onnxruntime-gpu (>=1.17.0,<2.0.0)
@@ -179,7 +179,7 @@ It requires installation of the `fastembed-gpu` package.
 pip install fastembed-gpu
 ```
 
-Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for detailed instructions and CUDA 12.x support.
+Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for detailed instructions, CUDA 12.x support and troubleshooting of the common issues.
 
 ```python
 from fastembed import TextEmbedding

{fastembed_gpu-0.3.6 → fastembed_gpu-0.4.0}/README.md

@@ -147,7 +147,7 @@ It requires installation of the `fastembed-gpu` package.
 pip install fastembed-gpu
 ```
 
-Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for detailed instructions and CUDA 12.x support.
+Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for detailed instructions, CUDA 12.x support and troubleshooting of the common issues.
 
 ```python
 from fastembed import TextEmbedding

{fastembed_gpu-0.3.6 → fastembed_gpu-0.4.0}/fastembed/common/model_management.py

@@ -74,7 +74,7 @@ class ModelManagement:
         if total_size_in_bytes == 0:
             print(f"Warning: Content-length header is missing or zero in the response from {url}.")
 
-        show_progress = total_size_in_bytes and show_progress
+        show_progress = bool(total_size_in_bytes and show_progress)
 
         with tqdm(
             total=total_size_in_bytes,
@@ -95,6 +95,7 @@ class ModelManagement:
         hf_source_repo: str,
         cache_dir: Optional[str] = None,
         extra_patterns: Optional[List[str]] = None,
+        local_files_only: bool = False,
         **kwargs,
     ) -> str:
         """
@@ -104,6 +105,7 @@ class ModelManagement:
             cache_dir (Optional[str]): The path to the cache directory.
             extra_patterns (Optional[List[str]]): extra patterns to allow in the snapshot download, typically
                 includes the required model files.
+            local_files_only (bool, optional): Whether to only use local files. Defaults to False.
         Returns:
             Path: The path to the model directory.
         """
@@ -121,7 +123,8 @@ class ModelManagement:
             repo_id=hf_source_repo,
             allow_patterns=allow_patterns,
             cache_dir=cache_dir,
-            local_files_only=kwargs.get("local_files_only", False),
+            local_files_only=local_files_only,
+            **kwargs,
         )
 
     @classmethod
@@ -162,9 +165,10 @@ class ModelManagement:
         return cache_dir
 
     @classmethod
-    def retrieve_model_gcs(cls, model_name: str, source_url: str, cache_dir: str) -> Path:
+    def retrieve_model_gcs(
+        cls, model_name: str, source_url: str, cache_dir: str, local_files_only: bool = False
+    ) -> Path:
         fast_model_name = f"fast-{model_name.split('/')[-1]}"
-
         cache_tmp_dir = Path(cache_dir) / "tmp"
         model_tmp_dir = cache_tmp_dir / fast_model_name
         model_dir = Path(cache_dir) / fast_model_name
@@ -183,22 +187,32 @@ class ModelManagement:
         if model_tar_gz.exists():
             model_tar_gz.unlink()
 
-        cls.download_file_from_gcs(
-            source_url,
-            output_path=str(model_tar_gz),
-        )
+        if not local_files_only:
+            cls.download_file_from_gcs(
+                source_url,
+                output_path=str(model_tar_gz),
+            )
 
-        cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir))
-        assert model_tmp_dir.exists(), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
+            cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir))
+            assert model_tmp_dir.exists(), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
 
-        model_tar_gz.unlink()
-        # Rename from tmp to final name is atomic
-        model_tmp_dir.rename(model_dir)
+            model_tar_gz.unlink()
+            # Rename from tmp to final name is atomic
+            model_tmp_dir.rename(model_dir)
+        else:
+            logger.error(
+                f"Could not find the model tar.gz file at {model_dir} and local_files_only=True."
+            )
+            raise ValueError(
+                f"Could not find the model tar.gz file at {model_dir} and local_files_only=True."
+            )
 
         return model_dir
 
     @classmethod
-    def download_model(cls, model: Dict[str, Any], cache_dir: Path, retries=3, **kwargs) -> Path:
+    def download_model(
+        cls, model: Dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs
+    ) -> Path:
         """
         Downloads a model from HuggingFace Hub or Google Cloud Storage.
 
@@ -223,7 +237,8 @@ class ModelManagement:
         Returns:
             Path: The path to the downloaded model directory.
         """
-
+        local_files_only = kwargs.get("local_files_only", False)
+        retries = 1 if local_files_only else retries
         hf_source = model.get("sources", {}).get("hf")
         url_source = model.get("sources", {}).get("url")
 
@@ -241,24 +256,34 @@ class ModelManagement:
                             hf_source,
                             cache_dir=str(cache_dir),
                             extra_patterns=extra_patterns,
-                            local_files_only=kwargs.get("local_files_only", False),
+                            **kwargs,
                         )
                     )
                 except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
-                    logger.error(
-                        f"Could not download model from HuggingFace: {e} "
-                        "Falling back to other sources."
-                    )
-            if url_source:
+                    if not local_files_only:
+                        logger.error(
+                            f"Could not download model from HuggingFace: {e} "
+                            "Falling back to other sources."
+                        )
+            if url_source or local_files_only:
                 try:
-                    return cls.retrieve_model_gcs(model["model"], url_source, str(cache_dir))
+                    return cls.retrieve_model_gcs(
+                        model["model"],
+                        url_source,
+                        str(cache_dir),
+                        local_files_only=local_files_only,
+                    )
                 except Exception:
-                    logger.error(f"Could not download model from url: {url_source}")
-
-            logger.error(
-                f"Could not download model from either source, sleeping for {sleep} seconds, {retries} retries left."
-            )
+                    if not local_files_only:
+                        logger.error(f"Could not download model from url: {url_source}")
+
+            if local_files_only:
+                logger.error("Could not find model in cache_dir")
+            else:
+                logger.error(
+                    f"Could not download model from either source, sleeping for {sleep} seconds, {retries} retries left."
+                )
             time.sleep(sleep)
             sleep *= 3
 
-        raise ValueError(f"Could not download model {model['model']} from any source.")
+        raise ValueError(f"Could not load model {model['model']} from any source.")

{fastembed_gpu-0.3.6 → fastembed_gpu-0.4.0}/fastembed/common/onnx_model.py

@@ -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")
 

{fastembed_gpu-0.3.6 → fastembed_gpu-0.4.0}/fastembed/common/utils.py

@@ -3,8 +3,11 @@ import tempfile
 from itertools import islice
 from pathlib import Path
 from typing import Generator, Iterable, Optional, Union
-
+import unicodedata
+import sys
 import numpy as np
+import re
+from typing import Set
 
 
 def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
@@ -37,7 +40,16 @@ def define_cache_dir(cache_dir: Optional[str] = None) -> Path:
         cache_path = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
     else:
         cache_path = Path(cache_dir)
-
     cache_path.mkdir(parents=True, exist_ok=True)
 
     return cache_path
+
+
+def get_all_punctuation() -> Set[str]:
+    return set(
+        chr(i) for i in range(sys.maxunicode) if unicodedata.category(chr(i)).startswith("P")
+    )
+
+
+def remove_non_alphanumeric(text: str) -> str:
+    return re.sub(r"[^\w\s]", " ", text, flags=re.UNICODE)

{fastembed_gpu-0.3.6 → fastembed_gpu-0.4.0}/fastembed/image/image_embedding.py

@@ -25,6 +25,7 @@ class ImageEmbedding(ImageEmbeddingBase):
                         "model": "Qdrant/clip-ViT-B-32-vision",
                         "dim": 512,
                         "description": "CLIP vision encoder based on ViT-B/32",
+                        "license": "mit",
                         "size_in_GB": 0.33,
                         "sources": {
                             "hf": "Qdrant/clip-ViT-B-32-vision",
@@ -45,21 +46,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

{fastembed_gpu-0.3.6 → fastembed_gpu-0.4.0}/fastembed/image/image_embedding_base.py

@@ -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.

{fastembed_gpu-0.3.6 → fastembed_gpu-0.4.0}/fastembed/image/onnx_embedding.py

@@ -13,6 +13,7 @@ supported_onnx_models = [
         "model": "Qdrant/clip-ViT-B-32-vision",
         "dim": 512,
         "description": "Image embeddings, Multimodal (text&image), 2021 year",
+        "license": "mit",
         "size_in_GB": 0.34,
         "sources": {
             "hf": "Qdrant/clip-ViT-B-32-vision",
@@ -23,6 +24,7 @@ supported_onnx_models = [
         "model": "Qdrant/resnet50-onnx",
         "dim": 2048,
         "description": "Image embeddings, Unimodal (image), 2016 year",
+        "license": "apache-2.0",
         "size_in_GB": 0.1,
         "sources": {
             "hf": "Qdrant/resnet50-onnx",
@@ -33,6 +35,7 @@ supported_onnx_models = [
         "model": "Qdrant/Unicom-ViT-B-16",
         "dim": 768,
         "description": "Image embeddings (more detailed than Unicom-ViT-B-32), Multimodal (text&image), 2023 year",
+        "license": "apache-2.0",
         "size_in_GB": 0.82,
         "sources": {
             "hf": "Qdrant/Unicom-ViT-B-16",
@@ -43,6 +46,7 @@ supported_onnx_models = [
         "model": "Qdrant/Unicom-ViT-B-32",
         "dim": 512,
         "description": "Image embeddings, Multimodal (text&image), 2023 year",
+        "license": "apache-2.0",
         "size_in_GB": 0.48,
         "sources": {
             "hf": "Qdrant/Unicom-ViT-B-32",
@@ -59,6 +63,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 +76,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)
-
-        model_description = self._get_model_description(model_name)
+        self.providers = providers
+        self.lazy_load = lazy_load
+
+        # 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)
-        model_dir = self.download_model(
-            model_description, self.cache_dir, local_files_only=self._local_files_only
+        self._model_dir = self.download_model(
+            self.model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
-        self.load_onnx_model(
-            model_dir=model_dir,
-            model_file=model_description["model_file"],
-            threads=threads,
-            providers=providers,
+        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=self.model_description["model_file"],
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
         )
 
     @classmethod
@@ -120,12 +160,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 +192,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,
+        )

{fastembed_gpu-0.3.6 → fastembed_gpu-0.4.0}/fastembed/image/onnx_image_model.py

@@ -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
+                num_workers=parallel or 1,
+                worker=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)

{fastembed_gpu-0.3.6 → fastembed_gpu-0.4.0}/fastembed/late_interaction/colbert.py

@@ -17,6 +17,7 @@ supported_colbert_models = [
         "model": "colbert-ir/colbertv2.0",
         "dim": 128,
         "description": "Late interaction model",
+        "license": "mit",
         "size_in_GB": 0.44,
         "sources": {
             "hf": "colbert-ir/colbertv2.0",
@@ -27,6 +28,7 @@ supported_colbert_models = [
         "model": "answerdotai/answerai-colbert-small-v1",
         "dim": 96,
         "description": "Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, 2024 year",
+        "license": "apache-2.0",
         "size_in_GB": 0.13,
         "sources": {
             "hf": "answerdotai/answerai-colbert-small-v1",
@@ -48,6 +50,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         if not is_doc:
             return output.model_output.astype(np.float32)
 
+        if output.input_ids is None or output.attention_mask is None:
+            raise ValueError(
+                "input_ids and attention_mask must be provided for document post-processing"
+            )
+
         for i, token_sequence in enumerate(output.input_ids):
             for j, token_id in enumerate(token_sequence):
                 if token_id in self.skip_list or token_id == self.pad_token_id:
@@ -117,6 +124,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 +137,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
+
+        # 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
 
-        model_description = self._get_model_description(model_name)
+        self.model_description = self._get_model_description(model_name)
         self.cache_dir = define_cache_dir(cache_dir)
 
-        model_dir = self.download_model(
-            model_description, self.cache_dir, local_files_only=self._local_files_only
+        self._model_dir = self.download_model(
+            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(
-            model_dir=model_dir,
-            model_file=model_description["model_file"],
-            threads=threads,
-            providers=providers,
+        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,
         )
         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,13 +224,19 @@ 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,
         )
 
-    def query_embed(self, query: Union[str, List[str]], **kwargs) -> np.ndarray:
+    def query_embed(self, query: Union[str, List[str]], **kwargs) -> Iterable[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 +249,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,
+        )