diff --git a/src/bgem3_embedding/impl.rs b/src/bgem3_embedding/impl.rs index 09312a0..e0e7f69 100644 --- a/src/bgem3_embedding/impl.rs +++ b/src/bgem3_embedding/impl.rs @@ -79,6 +79,7 @@ impl Bgem3Embedding { execution_providers, max_length, intra_threads, + .. } = options; let session = init_session_builder(execution_providers, intra_threads)? @@ -99,6 +100,7 @@ impl Bgem3Embedding { execution_providers, max_length, intra_threads, + .. } = options; let session = init_session_builder(execution_providers, intra_threads)? diff --git a/src/reranking/impl.rs b/src/reranking/impl.rs index 96dc5b3..abdd81f 100644 --- a/src/reranking/impl.rs +++ b/src/reranking/impl.rs @@ -93,9 +93,24 @@ impl TextRerank { execution_providers, max_length, intra_threads, + disable_cpu_fallback, + dimension_overrides, } = options; let mut session_builder = init_session_builder(execution_providers, intra_threads)?; + let builder_error = |err: ort::Error| { + Error::OrtBuilder(err.to_string()) + }; + if disable_cpu_fallback { + session_builder = session_builder + .with_disable_cpu_fallback() + .map_err(builder_error)?; + } + for (name, size) in dimension_overrides { + session_builder = session_builder + .with_dimension_override(name, size) + .map_err(builder_error)?; + } let session = match &model.onnx_source { OnnxSource::Memory(bytes) => session_builder.commit_from_memory(bytes)?, OnnxSource::File(path) => session_builder.commit_from_file(path)?, diff --git a/src/reranking/init.rs b/src/reranking/init.rs index cdc3d3d..f6d7d9c 100644 --- a/src/reranking/init.rs +++ b/src/reranking/init.rs @@ -33,6 +33,12 @@ pub struct RerankInitOptionsUserDefined { /// 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, + /// Refuse session creation when any graph node would fall back to ORT's + /// default CPU execution provider. + pub disable_cpu_fallback: bool, + /// Override named free dimensions before ORT optimizes and places the + /// user-defined reranker graph. + pub dimension_overrides: Vec<(String, i64)>, } impl Default for RerankInitOptionsUserDefined { @@ -41,6 +47,8 @@ impl Default for RerankInitOptionsUserDefined { execution_providers: Default::default(), max_length: DEFAULT_MAX_LENGTH, intra_threads: None, + disable_cpu_fallback: false, + dimension_overrides: Vec::new(), } } } @@ -72,6 +80,16 @@ impl RerankInitOptionsUserDefined { self.intra_threads = Some(intra_threads); self } + + pub fn with_disable_cpu_fallback(mut self, disable: bool) -> Self { + self.disable_cpu_fallback = disable; + self + } + + pub fn with_dimension_override(mut self, name: impl Into, size: i64) -> Self { + self.dimension_overrides.push((name.into(), size)); + self + } } /// Convert RerankInitOptions to RerankInitOptionsUserDefined @@ -83,6 +101,8 @@ impl From for RerankInitOptionsUserDefined { execution_providers: options.execution_providers, max_length: options.max_length, intra_threads: options.intra_threads, + disable_cpu_fallback: false, + dimension_overrides: Vec::new(), } } } @@ -143,8 +163,12 @@ mod tests { fn userdefined_builders_set_fields() { let o = RerankInitOptionsUserDefined::new() .with_max_length(128) - .with_intra_threads(2); + .with_intra_threads(2) + .with_disable_cpu_fallback(true) + .with_dimension_override("sequence_length", 128); assert_eq!(o.max_length, 128); assert_eq!(o.intra_threads, Some(2)); + assert!(o.disable_cpu_fallback); + assert_eq!(o.dimension_overrides, vec![("sequence_length".into(), 128)]); } } diff --git a/src/text_embedding/impl.rs b/src/text_embedding/impl.rs index 752c9b8..1e40323 100644 --- a/src/text_embedding/impl.rs +++ b/src/text_embedding/impl.rs @@ -92,6 +92,8 @@ impl TextEmbedding { execution_providers, max_length, intra_threads, + disable_cpu_fallback, + dimension_overrides, } = options; let session = { @@ -100,6 +102,17 @@ impl TextEmbedding { }; let mut session_builder = init_session_builder(execution_providers, intra_threads)?; + if disable_cpu_fallback { + session_builder = session_builder + .with_disable_cpu_fallback() + .map_err(builder_error)?; + } + for (name, size) in dimension_overrides { + session_builder = session_builder + .with_dimension_override(name, size) + .map_err(builder_error)?; + } + for external_initializer_file in model.external_initializers { session_builder = session_builder .with_external_initializer_file_in_memory( diff --git a/src/text_embedding/init.rs b/src/text_embedding/init.rs index 36c25b9..7d599a4 100644 --- a/src/text_embedding/init.rs +++ b/src/text_embedding/init.rs @@ -31,6 +31,14 @@ pub struct InitOptionsUserDefined { /// 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, + /// Refuse session creation when any graph node would fall back to the CPU + /// execution provider. This is useful for accelerator conformance tests + /// and for applications where silent partial placement is incorrect. + pub disable_cpu_fallback: bool, + /// Override named free dimensions before ORT optimizes and places the + /// graph. Static-shape execution providers such as QNN can use this with + /// one user-defined model session per admitted shape. + pub dimension_overrides: Vec<(String, i64)>, } impl InitOptionsUserDefined { @@ -60,6 +68,16 @@ impl InitOptionsUserDefined { self.intra_threads = Some(intra_threads); self } + + pub fn with_disable_cpu_fallback(mut self, disable: bool) -> Self { + self.disable_cpu_fallback = disable; + self + } + + pub fn with_dimension_override(mut self, name: impl Into, size: i64) -> Self { + self.dimension_overrides.push((name.into(), size)); + self + } } impl Default for InitOptionsUserDefined { @@ -68,6 +86,8 @@ impl Default for InitOptionsUserDefined { execution_providers: Default::default(), max_length: DEFAULT_MAX_LENGTH, intra_threads: None, + disable_cpu_fallback: false, + dimension_overrides: Vec::new(), } } } @@ -81,6 +101,8 @@ impl From for InitOptionsUserDefined { execution_providers: options.execution_providers, max_length: options.max_length, intra_threads: options.intra_threads, + disable_cpu_fallback: false, + dimension_overrides: Vec::new(), } } } @@ -146,3 +168,28 @@ pub struct TextEmbedding { pub(crate) quantization: QuantizationMode, pub(crate) output_key: Option, } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn user_defined_session_controls_are_opt_in_and_composable() { + let defaults = InitOptionsUserDefined::default(); + assert!(!defaults.disable_cpu_fallback); + assert!(defaults.dimension_overrides.is_empty()); + + let configured = InitOptionsUserDefined::new() + .with_disable_cpu_fallback(true) + .with_dimension_override("batch_size", 1) + .with_dimension_override("sequence_length", 512); + assert!(configured.disable_cpu_fallback); + assert_eq!( + configured.dimension_overrides, + [ + ("batch_size".to_string(), 1), + ("sequence_length".to_string(), 512) + ] + ); + } +}