Skip to content

feat: expose ONNX Runtime session config entries through InitOptions #292

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鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
anush008 merged 2 commits into from
Sep 12, 2026
Merged

feat: expose ONNX Runtime session config entries through InitOptions #292

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e6494ad
feat: expose ONNX Runtime session config entries through InitOptions
paulnsorensen Sep 12, 2026
764920b
Update README.md
anush008 Sep 12, 2026
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
9 changes: 6 additions & 3 deletions src/bgem3_embedding/impl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ impl Bgem3Embedding {
show_download_progress,
execution_providers,
intra_threads,
session_config,
} = options;

let model_repo = Bgem3Embedding::retrieve_model(
Expand Down Expand Up @@ -63,7 +64,7 @@ impl Bgem3Embedding {
}
}

let session = init_session_builder(execution_providers, intra_threads)?
let session = init_session_builder(execution_providers, intra_threads, session_config)?
.commit_from_file(model_file_reference)?;

let tokenizer = load_tokenizer_hf_hub(model_repo, max_length)?;
Expand All @@ -79,10 +80,11 @@ impl Bgem3Embedding {
execution_providers,
max_length,
intra_threads,
session_config,
..
} = options;

let session = init_session_builder(execution_providers, intra_threads)?
let session = init_session_builder(execution_providers, intra_threads, session_config)?
.commit_from_memory(&model.onnx_file)?;

let tokenizer = load_tokenizer(model.tokenizer_files, max_length)?;
Expand All @@ -100,10 +102,11 @@ impl Bgem3Embedding {
execution_providers,
max_length,
intra_threads,
session_config,
..
} = options;

let session = init_session_builder(execution_providers, intra_threads)?
let session = init_session_builder(execution_providers, intra_threads, session_config)?
.commit_from_file(model_path.as_ref().join("model.onnx"))?;

let tokenizer = load_tokenizer(tokenizer_files, max_length)?;
Expand Down
16 changes: 16 additions & 0 deletions src/common.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,7 @@ pub fn pull_from_hf(
pub(crate) fn init_session_builder(
execution_providers: Vec<ExecutionProviderDispatch>,
intra_threads: Option<usize>,
session_config: Vec<(String, String)>,
) -> Result<SessionBuilder> {
let threads = match intra_threads {
Some(n) => n,
Expand All @@ -289,6 +290,12 @@ pub(crate) fn init_session_builder(
.with_intra_threads(threads)
.map_err(builder_error)?;

for (key, value) in session_config {
builder = builder
.with_config_entry(&key, &value)
.map_err(builder_error)?;
}

if has_directml {
builder = builder
.with_memory_pattern(false)
Expand Down Expand Up @@ -362,4 +369,13 @@ mod tests {
"error message was: {err}"
);
}
#[test]
fn init_session_builder_applies_config_entry() {
let builder = init_session_builder(
vec![],
Some(1),
vec![("session.disable_prepacking".into(), "1".into())],
);
assert!(builder.is_ok());
}
}
6 changes: 4 additions & 2 deletions src/image_embedding/impl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ impl ImageEmbedding {
cache_dir,
show_download_progress,
intra_threads,
session_config,
} = options;

let model_repo = ImageEmbedding::retrieve_model(
Expand All @@ -61,7 +62,7 @@ impl ImageEmbedding {
source: Box::new(e),
})?;

let session = init_session_builder(execution_providers, intra_threads)?
let session = init_session_builder(execution_providers, intra_threads, session_config)?
.commit_from_file(model_file_reference)?;

Ok(Self::new(preprocessor, session))
Expand All @@ -77,11 +78,12 @@ impl ImageEmbedding {
let ImageInitOptionsUserDefined {
execution_providers,
intra_threads,
session_config,
} = options;

let preprocessor = Compose::from_bytes(model.preprocessor_file)?;

let session = init_session_builder(execution_providers, intra_threads)?
let session = init_session_builder(execution_providers, intra_threads, session_config)?
.commit_from_memory(&model.onnx_file)?;

Ok(Self::new(preprocessor, session))
Expand Down
13 changes: 13 additions & 0 deletions src/image_embedding/init.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,10 @@ pub struct ImageInitOptionsUserDefined {
/// every available CPU core via `std::thread::available_parallelism`.
/// Set this to cap CPU usage (e.g. on laptops) at the cost of throughput.
pub intra_threads: Option<usize>,
/// ONNX Runtime session configuration entries, applied with
/// `SessionBuilder::with_config_entry`. Use this for settings that have
/// no dedicated builder method, such as `mlas.disable_kleidiai`.
pub session_config: Vec<(String, String)>,
}

impl ImageInitOptionsUserDefined {
Expand All @@ -38,6 +42,14 @@ impl ImageInitOptionsUserDefined {
self.intra_threads = Some(intra_threads);
self
}

/// Add an ONNX Runtime session configuration entry, applied with
/// `SessionBuilder::with_config_entry`. Call it once per entry.
/// Example: `.with_session_config("mlas.disable_kleidiai", "1")`.
pub fn with_session_config(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.session_config.push((key.into(), value.into()));
self
}
}

/// Convert ImageInitOptions to ImageInitOptionsUserDefined
Expand All @@ -48,6 +60,7 @@ impl From<ImageInitOptions> for ImageInitOptionsUserDefined {
ImageInitOptionsUserDefined {
execution_providers: options.execution_providers,
intra_threads: options.intra_threads,
session_config: options.session_config,
}
}
}
Expand Down
26 changes: 26 additions & 0 deletions src/init.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,10 @@ pub struct InitOptionsWithLength<M> {
/// every available CPU core via `std::thread::available_parallelism`.
/// Set this to cap CPU usage (e.g. on laptops) at the cost of throughput.
pub intra_threads: Option<usize>,
/// ONNX Runtime session configuration entries, applied with
/// `SessionBuilder::with_config_entry`. Use this for settings that have
/// no dedicated builder method, such as `mlas.disable_kleidiai`.
pub session_config: Vec<(String, String)>,
}

#[derive(Debug, Clone)]
Expand All @@ -31,6 +35,10 @@ pub struct InitOptions<M> {
/// every available CPU core via `std::thread::available_parallelism`.
/// Set this to cap CPU usage (e.g. on laptops) at the cost of throughput.
pub intra_threads: Option<usize>,
/// ONNX Runtime session configuration entries, applied with
/// `SessionBuilder::with_config_entry`. Use this for settings that have
/// no dedicated builder method, such as `mlas.disable_kleidiai`.
pub session_config: Vec<(String, String)>,
}

impl<M: Default + HasMaxLength> Default for InitOptionsWithLength<M> {
Expand All @@ -42,6 +50,7 @@ impl<M: Default + HasMaxLength> Default for InitOptionsWithLength<M> {
show_download_progress: true,
max_length: M::MAX_LENGTH,
intra_threads: None,
session_config: Vec::new(),
}
}
}
Expand All @@ -54,6 +63,7 @@ impl<M: Default> Default for InitOptions<M> {
cache_dir: get_cache_dir().into(),
show_download_progress: true,
intra_threads: None,
session_config: Vec::new(),
}
}
}
Expand Down Expand Up @@ -96,6 +106,14 @@ impl<M: Default + HasMaxLength> InitOptionsWithLength<M> {
self
}

/// Add an ONNX Runtime session configuration entry, applied with
/// `SessionBuilder::with_config_entry`. Call it once per entry.
/// Example: `.with_session_config("mlas.disable_kleidiai", "1")`.
pub fn with_session_config(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.session_config.push((key.into(), value.into()));
self
}

/// Set whether to show download progress
pub fn with_show_download_progress(mut self, show_download_progress: bool) -> Self {
self.show_download_progress = show_download_progress;
Expand Down Expand Up @@ -135,6 +153,14 @@ impl<M: Default> InitOptions<M> {
self
}

/// Add an ONNX Runtime session configuration entry, applied with
/// `SessionBuilder::with_config_entry`. Call it once per entry.
/// Example: `.with_session_config("mlas.disable_kleidiai", "1")`.
pub fn with_session_config(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.session_config.push((key.into(), value.into()));
self
}

/// Set whether to show download progress
pub fn with_show_download_progress(mut self, show_download_progress: bool) -> Self {
self.show_download_progress = show_download_progress;
Expand Down
7 changes: 5 additions & 2 deletions src/reranking/impl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ impl TextRerank {
cache_dir,
show_download_progress,
intra_threads,
session_config,
} = options;

let model_repo = pull_from_hf(model_name.to_string(), cache_dir, show_download_progress)?;
Expand All @@ -75,7 +76,7 @@ impl TextRerank {
})?;
}

let session = init_session_builder(execution_providers, intra_threads)?
let session = init_session_builder(execution_providers, intra_threads, session_config)?
.commit_from_file(model_file_reference)?;

let tokenizer = load_tokenizer_hf_hub(model_repo, max_length)?;
Expand All @@ -95,9 +96,11 @@ impl TextRerank {
intra_threads,
disable_cpu_fallback,
dimension_overrides,
session_config,
} = options;

let mut session_builder = init_session_builder(execution_providers, intra_threads)?;
let mut session_builder =
init_session_builder(execution_providers, intra_threads, session_config)?;
let builder_error = |err: ort::Error<ort::session::builder::SessionBuilder>| {
Error::OrtBuilder(err.to_string())
};
Expand Down
31 changes: 31 additions & 0 deletions src/reranking/init.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,10 @@ pub struct RerankInitOptionsUserDefined {
/// Override named free dimensions before ORT optimizes and places the
/// user-defined reranker graph.
pub dimension_overrides: Vec<(String, i64)>,
/// ONNX Runtime session configuration entries, applied with
/// `SessionBuilder::with_config_entry`. Use this for settings that have
/// no dedicated builder method, such as `mlas.disable_kleidiai`.
pub session_config: Vec<(String, String)>,
}

impl Default for RerankInitOptionsUserDefined {
Expand All @@ -49,6 +53,7 @@ impl Default for RerankInitOptionsUserDefined {
intra_threads: None,
disable_cpu_fallback: false,
dimension_overrides: Vec::new(),
session_config: Vec::new(),
}
}
}
Expand Down Expand Up @@ -90,6 +95,14 @@ impl RerankInitOptionsUserDefined {
self.dimension_overrides.push((name.into(), size));
self
}

/// Add an ONNX Runtime session configuration entry, applied with
/// `SessionBuilder::with_config_entry`. Call it once per entry.
/// Example: `.with_session_config("mlas.disable_kleidiai", "1")`.
pub fn with_session_config(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.session_config.push((key.into(), value.into()));
self
}
}

/// Convert RerankInitOptions to RerankInitOptionsUserDefined
Expand All @@ -103,6 +116,7 @@ impl From<RerankInitOptions> for RerankInitOptionsUserDefined {
intra_threads: options.intra_threads,
disable_cpu_fallback: false,
dimension_overrides: Vec::new(),
session_config: options.session_config,
}
}
}
Expand Down Expand Up @@ -171,4 +185,21 @@ mod tests {
assert!(o.disable_cpu_fallback);
assert_eq!(o.dimension_overrides, vec![("sequence_length".into(), 128)]);
}

#[test]
fn session_config_is_collected_and_carried_by_from() {
let opts = RerankInitOptions::new(RerankerModel::default())
.with_session_config("a", "1")
.with_session_config("b", "2");
assert_eq!(
opts.session_config,
vec![("a".into(), "1".into()), ("b".into(), "2".into())]
);

let user_defined = RerankInitOptionsUserDefined::from(opts);
assert_eq!(
user_defined.session_config,
vec![("a".into(), "1".into()), ("b".into(), "2".into())]
);
}
}
3 changes: 2 additions & 1 deletion src/sparse_text_embedding/impl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ impl SparseTextEmbedding {
show_download_progress,
execution_providers,
intra_threads,
session_config,
} = options;

let model_repo = SparseTextEmbedding::retrieve_model(
Expand Down Expand Up @@ -64,7 +65,7 @@ impl SparseTextEmbedding {
}
}

let session = init_session_builder(execution_providers, intra_threads)?
let session = init_session_builder(execution_providers, intra_threads, session_config)?
.commit_from_file(model_file_reference)?;

let tokenizer = load_tokenizer_hf_hub(model_repo, max_length)?;
Expand Down
7 changes: 5 additions & 2 deletions src/text_embedding/impl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ impl TextEmbedding {
cache_dir,
show_download_progress,
intra_threads,
session_config,
} = options;

let model_repo = TextEmbedding::retrieve_model(
Expand Down Expand Up @@ -68,7 +69,7 @@ impl TextEmbedding {
// prioritise loading pooling config if available, if not (thanks qdrant!), look for it in hardcoded
let post_processing = TextEmbedding::get_default_pooling_method(&model_name);

let session = init_session_builder(execution_providers, intra_threads)?
let session = init_session_builder(execution_providers, intra_threads, session_config)?
.commit_from_file(model_file_reference)?;

let tokenizer = load_tokenizer_hf_hub(model_repo, max_length)?;
Expand All @@ -94,13 +95,15 @@ impl TextEmbedding {
intra_threads,
disable_cpu_fallback,
dimension_overrides,
session_config,
} = options;

let session = {
let builder_error = |err: ort::Error<ort::session::builder::SessionBuilder>| {
Error::OrtBuilder(err.to_string())
};
let mut session_builder = init_session_builder(execution_providers, intra_threads)?;
let mut session_builder =
init_session_builder(execution_providers, intra_threads, session_config)?;

if disable_cpu_fallback {
session_builder = session_builder
Expand Down
Loading