new: expose some onnx session options (#578)

* new: expose some onnx session options

* fix: fix extra session options is None case

* fix: fix missing params

* new: add tests

Author: joein

fastembed/common/onnx_model.py

@@ -24,6 +24,8 @@ class OnnxOutputContext:
 
 
 class OnnxModel(Generic[T]):
+    EXPOSED_SESSION_OPTIONS = ("enable_cpu_mem_arena",)
+
     @classmethod
     def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]:
         raise NotImplementedError("Subclasses must implement this method")
@@ -60,6 +62,7 @@ def _load_onnx_model(
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_id: Optional[int] = None,
+        extra_session_options: Optional[dict[str, Any]] = None,
     ) -> None:
         model_path = model_dir / model_file
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
@@ -99,6 +102,9 @@ def _load_onnx_model(
             so.intra_op_num_threads = threads
             so.inter_op_num_threads = threads
 
+        if extra_session_options is not None:
+            self.add_extra_session_options(so, extra_session_options)
+
         self.model = ort.InferenceSession(
             str(model_path), providers=onnx_providers, sess_options=so
         )
@@ -113,6 +119,38 @@ def _load_onnx_model(
                     RuntimeWarning,
                 )
 
+    @classmethod
+    def _select_exposed_session_options(cls, model_kwargs: dict[str, Any]) -> dict[str, Any]:
+        """A convenience method to select the exposed session options in models
+
+        Args:
+            model_kwargs (dict[str, Any]): The model kwargs.
+
+        Returns:
+            dict[str, Any]: a dict with filtered exposed session options.
+        """
+        return {k: v for k, v in model_kwargs.items() if k in cls.EXPOSED_SESSION_OPTIONS}
+
+    @classmethod
+    def add_extra_session_options(
+        cls, session_options: ort.SessionOptions, extra_options: dict[str, Any]
+    ) -> None:
+        """Add extra session options to the existing options object in-place
+
+        Args:
+            session_options (ort.SessionOptions): The existing session options object.
+            extra_options (dict[str, Any]): The extra session options available in cls.EXPOSED_SESSION_OPTIONS.
+
+        Returns:
+            None
+        """
+        for option in extra_options:
+            assert (
+                option in cls.EXPOSED_SESSION_OPTIONS
+            ), f"{option} is unknown or not exposed (exposed options: {cls.EXPOSED_SESSION_OPTIONS})"
+        if "enable_cpu_mem_arena" in extra_options:
+            session_options.enable_cpu_mem_arena = extra_options["enable_cpu_mem_arena"]
+
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 

fastembed/image/onnx_embedding.py

@@ -98,6 +98,7 @@ def __init__(
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
@@ -134,6 +135,7 @@ def load_onnx_model(self) -> None:
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
 
     @classmethod
@@ -180,6 +182,7 @@ def embed(
             device_ids=self.device_ids,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 

fastembed/image/onnx_image_model.py

@@ -55,6 +55,7 @@ def _load_onnx_model(
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_id: Optional[int] = None,
+        extra_session_options: Optional[dict[str, Any]] = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -63,6 +64,7 @@ def _load_onnx_model(
             providers=providers,
             cuda=cuda,
             device_id=device_id,
+            extra_session_options=extra_session_options,
         )
         self.processor = load_preprocessor(model_dir=model_dir)
 
@@ -99,6 +101,7 @@ def _embed_images(
         device_ids: Optional[list[int]] = None,
         local_files_only: bool = False,
         specific_model_path: Optional[str] = None,
+        extra_session_options: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -130,6 +133,9 @@ def _embed_images(
                 **kwargs,
             }
 
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
                 worker=self._get_worker_class(),

fastembed/late_interaction/colbert.py

@@ -143,6 +143,7 @@ def __init__(
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
@@ -182,6 +183,7 @@ def load_onnx_model(self) -> None:
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
         self.query_tokenizer, _ = load_tokenizer(model_dir=self._model_dir)
 
@@ -235,6 +237,7 @@ def embed(
             device_ids=self.device_ids,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 

fastembed/late_interaction_multimodal/colpali.py

@@ -80,6 +80,7 @@ def __init__(
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
@@ -125,6 +126,7 @@ def load_onnx_model(self) -> None:
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
 
     def _post_process_onnx_image_output(
@@ -238,6 +240,7 @@ def embed_text(
             device_ids=self.device_ids,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 
@@ -273,6 +276,7 @@ def embed_image(
             device_ids=self.device_ids,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 

fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -64,6 +64,7 @@ def _load_onnx_model(
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_id: Optional[int] = None,
+        extra_session_options: Optional[dict[str, Any]] = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -72,6 +73,7 @@ def _load_onnx_model(
             providers=providers,
             cuda=cuda,
             device_id=device_id,
+            extra_session_options=extra_session_options,
         )
         self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
         assert self.tokenizer is not None
@@ -122,6 +124,7 @@ def _embed_documents(
         device_ids: Optional[list[int]] = None,
         local_files_only: bool = False,
         specific_model_path: Optional[str] = None,
+        extra_session_options: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -153,6 +156,9 @@ def _embed_documents(
                 **kwargs,
             }
 
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
                 worker=self._get_text_worker_class(),
@@ -189,6 +195,7 @@ def _embed_images(
         device_ids: Optional[list[int]] = None,
         local_files_only: bool = False,
         specific_model_path: Optional[str] = None,
+        extra_session_options: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -220,6 +227,9 @@ def _embed_images(
                 **kwargs,
             }
 
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
                 worker=self._get_image_worker_class(),

fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -111,6 +111,7 @@ def __init__(
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
@@ -150,6 +151,7 @@ def load_onnx_model(self) -> None:
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
 
     def rerank(
@@ -192,6 +194,7 @@ def rerank_pairs(
             device_ids=self.device_ids,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )
 

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -33,6 +33,7 @@ def _load_onnx_model(
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_id: Optional[int] = None,
+        extra_session_options: Optional[dict[str, Any]] = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -41,6 +42,7 @@ def _load_onnx_model(
             providers=providers,
             cuda=cuda,
             device_id=device_id,
+            extra_session_options=extra_session_options,
         )
         self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
         assert self.tokenizer is not None
@@ -96,6 +98,7 @@ def _rerank_pairs(
         device_ids: Optional[list[int]] = None,
         local_files_only: bool = False,
         specific_model_path: Optional[str] = None,
+        extra_session_options: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         is_small = False
@@ -127,6 +130,9 @@ def _rerank_pairs(
                 **kwargs,
             }
 
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
                 worker=self._get_worker_class(),

fastembed/sparse/bm42.py

@@ -103,6 +103,7 @@ def __init__(
         super().__init__(model_name, cache_dir, threads, **kwargs)
         self.providers = providers
         self.lazy_load = lazy_load
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
 
         # List of device ids, that can be used for data parallel processing in workers
         self.device_ids = device_ids
@@ -146,6 +147,7 @@ def load_onnx_model(self) -> None:
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
 
         for token, idx in self.tokenizer.get_vocab().items():  # type: ignore[union-attr]
@@ -312,6 +314,7 @@ def embed(
             alpha=self.alpha,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
         )
 
     @classmethod

fastembed/sparse/minicoil.py

@@ -117,6 +117,8 @@ def __init__(
         self.device_ids = device_ids
         self.cuda = cuda
         self.device_id = device_id
+        self._extra_session_options = self._select_exposed_session_options(kwargs)
+
         self.k = k
         self.b = b
         self.avg_len = avg_len
@@ -153,6 +155,7 @@ def load_onnx_model(self) -> None:
             providers=self.providers,
             cuda=self.cuda,
             device_id=self.device_id,
+            extra_session_options=self._extra_session_options,
         )
 
         assert self.tokenizer is not None
@@ -221,6 +224,7 @@ def embed(
             is_query=False,
             local_files_only=self._local_files_only,
             specific_model_path=self._specific_model_path,
+            extra_session_options=self._extra_session_options,
             **kwargs,
         )