Skip to content

new: expose some onnx session options #578

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
joein merged 4 commits into from
Nov 25, 2025
Merged

new: expose some onnx session options #578

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0e17bea
new: expose some onnx session options
joein Nov 24, 2025
f458562
fix: fix extra session options is None case
joein Nov 24, 2025
298912a
fix: fix missing params
joein Nov 24, 2025
2ef4d17
new: add tests
joein Nov 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 38 additions & 0 deletions fastembed/common/onnx_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ class OnnxOutputContext:


class OnnxModel(Generic[T]):
EXPOSED_SESSION_OPTIONS = ("enable_cpu_mem_arena",)

@tbung tbung Nov 25, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we don't expect this to be expanded at runtime it could be an enum like

from enum import StrEnum, auto

class ExposedSessionOptions(StrEnum):
    enable_cpu_mem_arena = auto()

or similar, would make type hinting nicer and still allows in checks and iteration.

Nvm, StrEnum was added with 3.11.

@joein joein Nov 25, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess it can be emulated with ExposedSessionOptions(str, Enum), but feels like a bit too much, users won't see this enum anyway

@tbung tbung Nov 25, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, it's not worth it to emulate it that way, it was a low hanging fruit with StrEnum built in but definitely not worth reimplementing it (also just (str, Enum) won't give you auto iirc so it'd be even more work). I just misremembered when StrEnum was added.

@tbung tbung Nov 25, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess we could type it as a tuple[Literal[...]] but again, not worth the effort for something that's only internal.


@classmethod
def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]:
raise NotImplementedError("Subclasses must implement this method")
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
)
Expand All @@ -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"]
Comment thread
coderabbitai[bot] marked this conversation as resolved.

def load_onnx_model(self) -> None:
raise NotImplementedError("Subclasses must implement this method")

Expand Down
3 changes: 3 additions & 0 deletions fastembed/image/onnx_embedding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
)

Expand Down
6 changes: 6 additions & 0 deletions fastembed/image/onnx_image_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(),
Expand Down
3 changes: 3 additions & 0 deletions fastembed/late_interaction/colbert.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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,
)

Expand Down
4 changes: 4 additions & 0 deletions fastembed/late_interaction_multimodal/colpali.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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,
)

Expand Down Expand Up @@ -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,
)

Expand Down
10 changes: 10 additions & 0 deletions fastembed/late_interaction_multimodal/onnx_multimodal_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -153,6 +156,9 @@ def _embed_documents(
**kwargs,
}

if extra_session_options is not None:
params.update(extra_session_options)
Comment thread
coderabbitai[bot] marked this conversation as resolved.

pool = ParallelWorkerPool(
num_workers=parallel or 1,
worker=self._get_text_worker_class(),
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -220,6 +227,9 @@ def _embed_images(
**kwargs,
}

if extra_session_options is not None:
params.update(extra_session_options)
Comment on lines +230 to +231

@coderabbitai coderabbitai Bot Nov 24, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Critical bug: session options are incorrectly flattened into params.

Same issue as in _embed_documents: the dictionary structure must be preserved when passing to workers.

Apply this diff:

-            if extra_session_options is not None:
-                params.update(extra_session_options)
+            if extra_session_options is not None:
+                params["extra_session_options"] = extra_session_options
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if extra_session_options is not None:
params.update(extra_session_options)
if extra_session_options is not None:
params["extra_session_options"] = extra_session_options
🤖 Prompt for AI Agents 
In fastembed/late_interaction_multimodal/onnx_multimodal_model.py around lines
230-231, the code currently flattens extra_session_options into params with
params.update(extra_session_options); instead preserve the dictionary structure
by adding the options as a nested entry (e.g., params['extra_session_options']
or the same key used in _embed_documents) instead of updating params, so the
worker receives a single dict field containing the session options rather than
merging its keys into the top-level params.


pool = ParallelWorkerPool(
num_workers=parallel or 1,
worker=self._get_image_worker_class(),
Expand Down
3 changes: 3 additions & 0 deletions fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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,
)

Expand Down
6 changes: 6 additions & 0 deletions fastembed/rerank/cross_encoder/onnx_text_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(),
Expand Down
3 changes: 3 additions & 0 deletions fastembed/sparse/bm42.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions fastembed/sparse/minicoil.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Comment on lines +120 to +121

@coderabbitai coderabbitai Bot Nov 24, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Extra session options initialized but not propagated.

The _extra_session_options are initialized here but never passed to _load_onnx_model (line 151-158) or _embed_documents (lines 211-227, 235-249). This means the session options won't be applied for MiniCOIL models.

Apply this diff to propagate the session options:

     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,
+            extra_session_options=self._extra_session_options,
         )

Also add to the embed method (around line 211):

         yield from self._embed_documents(
             model_name=self.model_name,
             cache_dir=str(self.cache_dir),
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
             k=self.k,
             b=self.b,
             avg_len=self.avg_len,
             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,
         )

And similarly for query_embed (around line 235).

Committable suggestion skipped: line range outside the PR's diff.

self.k = k
self.b = b
self.avg_len = avg_len
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
)

Expand Down
3 changes: 3 additions & 0 deletions fastembed/sparse/splade_pp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,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
Expand Down Expand Up @@ -133,6 +134,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 embed(
Expand Down Expand Up @@ -168,6 +170,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,
)

Expand Down
4 changes: 3 additions & 1 deletion fastembed/text/onnx_embedding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,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
self.cuda = cuda
Expand Down Expand Up @@ -291,6 +291,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,
)

Expand Down Expand Up @@ -327,6 +328,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,
)


Expand Down
Loading
Loading