From 88e82b9c1e3d839c57b48986808adf25e1685767 Mon Sep 17 00:00:00 2001 From: pham-tuan-binh Date: Fri, 28 Aug 2026 16:53:06 +0700 Subject: [PATCH 1/2] livekit-wakeword: runtime-configurable ONNX session options MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sessions were built with a bare `Session::builder()`, so callers had no way to tune them: how much CPU wake word detection uses, and whether the graph is optimized at all, were both fixed by the crate. The Python SDK takes a `sess_options` parameter for exactly this. `SessionOptions` now describes how every session is created — graph optimization level, intra/inter-op threads, sequential execution, thread spinning, and arbitrary session config entries — and is passed at runtime to `WakeWordModel::with_session_options` or `load_model_with_session_options`. `Default` requests `OptimizationLevel::Level3`, ONNX Runtime's own default. Options the active backend does not implement are skipped rather than failing session creation: `ort-tract` implements only the optimization level, and since it runs single-threaded without spin-waiting, a request to limit threads or disable spinning already describes what it does. --- .../wakeword_runtime_session_options.md | 16 ++ livekit-wakeword/src/embedding.rs | 6 +- livekit-wakeword/src/lib.rs | 147 +++++++++++++++++- livekit-wakeword/src/melspectrogram.rs | 6 +- livekit-wakeword/src/wakeword.rs | 45 +++++- livekit-wakeword/tests/integration.rs | 58 ++++++- 6 files changed, 261 insertions(+), 17 deletions(-) create mode 100644 .changeset/wakeword_runtime_session_options.md diff --git a/.changeset/wakeword_runtime_session_options.md b/.changeset/wakeword_runtime_session_options.md new file mode 100644 index 000000000..4969e6031 --- /dev/null +++ b/.changeset/wakeword_runtime_session_options.md @@ -0,0 +1,16 @@ +--- +livekit-wakeword: minor +--- + +# Runtime-configurable ONNX session options + +`WakeWordModel::with_session_options` and `load_model_with_session_options` accept a +`SessionOptions` describing how the ONNX sessions are created — graph optimization +level, intra/inter-op threads, sequential execution, thread spinning, and arbitrary +session config entries — matching the `sess_options` parameter the Python SDK takes. +The default is `SessionOptions::default()`, whose `OptimizationLevel::Level3` matches +ONNX Runtime's own default; the `ort-tract` backend used on every target except +aarch64 Windows runs tract's `into_optimized()` only when a level is requested, so +requesting one made `predict()` over a 2 s window 6.2x faster (329 ms to 53 ms +median, Apple M-series release build). Options that the active backend does not +implement are skipped rather than failing session creation. diff --git a/livekit-wakeword/src/embedding.rs b/livekit-wakeword/src/embedding.rs index 350a5dd4d..a707409ee 100644 --- a/livekit-wakeword/src/embedding.rs +++ b/livekit-wakeword/src/embedding.rs @@ -16,7 +16,7 @@ use ndarray::{Array, Array1}; use ort::session::Session; use ort::value::Tensor; -use crate::{build_session_from_memory, WakeWordError}; +use crate::{build_session_from_memory, SessionOptions, WakeWordError}; const MODEL_BYTES: &[u8] = include_bytes!("../onnx/embedding_model.onnx"); @@ -33,8 +33,8 @@ pub struct EmbeddingModel { } impl EmbeddingModel { - pub fn new() -> Result { - Ok(Self { session: build_session_from_memory(MODEL_BYTES)? }) + pub fn new(options: &SessionOptions) -> Result { + Ok(Self { session: build_session_from_memory(MODEL_BYTES, options)? }) } // Run the embedding model on mel spectrogram features and return the embedding. diff --git a/livekit-wakeword/src/lib.rs b/livekit-wakeword/src/lib.rs index 619869de9..ac3f098d6 100644 --- a/livekit-wakeword/src/lib.rs +++ b/livekit-wakeword/src/lib.rs @@ -16,7 +16,10 @@ use std::path::Path; #[cfg(use_tract)] use std::sync::Once; -use ort::session::Session; +use ort::session::{ + builder::{GraphOptimizationLevel, SessionBuilder}, + Session, +}; #[cfg(use_tract)] static INIT_TRACT: Once = Once::new(); @@ -74,15 +77,149 @@ pub(crate) fn to_resampler_rate(hz: u32) -> Result Result { +/// Graph optimization level applied to the ONNX graph when a session is created. +/// +/// The default, [`Level3`](Self::Level3), matches ONNX Runtime's own default. It +/// also matters more than it looks: the `ort-tract` backend used on every target +/// except aarch64 Windows runs tract's `into_optimized()` only when a session asks +/// for some level of optimization, and wake word inference measured several times +/// slower on the unoptimized graph. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum OptimizationLevel { + /// Run the graph as it was loaded. + Disable, + /// Semantics-preserving rewrites that remove redundant nodes and computation. + Level1, + /// Level 1, plus fusions that depend on the execution provider. + Level2, + /// Level 2, plus layout optimizations. ONNX Runtime's default. + #[default] + Level3, +} + +impl From for GraphOptimizationLevel { + fn from(level: OptimizationLevel) -> Self { + match level { + OptimizationLevel::Disable => GraphOptimizationLevel::Disable, + OptimizationLevel::Level1 => GraphOptimizationLevel::Level1, + OptimizationLevel::Level2 => GraphOptimizationLevel::Level2, + OptimizationLevel::Level3 => GraphOptimizationLevel::Level3, + } + } +} + +/// Tuning applied to every ONNX session a [`WakeWordModel`] creates: the two +/// bundled feature extraction models and each wake word classifier. +/// +/// [`Default`] is what [`WakeWordModel::new`] uses. Pass a value of your own to +/// [`WakeWordModel::with_session_options`] to trade latency for CPU, which is +/// useful when detection runs as a background process on a small machine: +/// +/// ```no_run +/// # use livekit_wakeword::{SessionOptions, WakeWordModel}; +/// # fn main() -> Result<(), livekit_wakeword::WakeWordError> { +/// let options = SessionOptions { +/// intra_threads: Some(1), +/// inter_threads: Some(1), +/// parallel_execution: Some(false), +/// intra_op_spinning: Some(false), +/// inter_op_spinning: Some(false), +/// ..Default::default() +/// }; +/// let model = WakeWordModel::with_session_options(&["hey_livekit.onnx"], 16000, options)?; +/// # Ok(()) +/// # } +/// ``` +/// +/// Every field except `optimization_level` is best-effort: a backend that does not +/// implement an option has it skipped rather than failing session creation. In +/// particular the `ort-tract` backend implements only the optimization level, and +/// runs single-threaded without spin-waiting, so the thread and spinning fields +/// above describe what it already does. +#[derive(Clone, Debug, Default)] +pub struct SessionOptions { + /// Graph optimizations to apply when the session is created. + pub optimization_level: OptimizationLevel, + /// Threads used to run a single operator. `Some(1)` keeps each session on the + /// calling thread. + pub intra_threads: Option, + /// Threads used to run operators in parallel, when `parallel_execution` is on. + pub inter_threads: Option, + /// Whether independent operators may run in parallel. `Some(false)` is ONNX + /// Runtime's sequential execution mode. + pub parallel_execution: Option, + /// Whether intra-op threads may spin before blocking. `Some(false)` trades a + /// little latency for markedly less CPU while idle. + pub intra_op_spinning: Option, + /// Whether inter-op threads may spin before blocking. + pub inter_op_spinning: Option, + /// Session config entries applied verbatim, for anything the fields above do + /// not cover. + pub config_entries: Vec<(String, String)>, +} + +impl SessionOptions { + fn apply(&self, builder: SessionBuilder) -> Result { + // Unlike everything below, the optimization level is implemented by every + // backend, so a failure here is a real error rather than an unsupported knob. + let mut builder = builder.with_optimization_level(self.optimization_level.into())?; + + if let Some(threads) = self.intra_threads { + builder = best_effort(builder, |b| b.with_intra_threads(threads))?; + } + if let Some(threads) = self.inter_threads { + builder = best_effort(builder, |b| b.with_inter_threads(threads))?; + } + if let Some(parallel) = self.parallel_execution { + builder = best_effort(builder, |b| b.with_parallel_execution(parallel))?; + } + if let Some(enable) = self.intra_op_spinning { + builder = best_effort(builder, |b| b.with_intra_op_spinning(enable))?; + } + if let Some(enable) = self.inter_op_spinning { + builder = best_effort(builder, |b| b.with_inter_op_spinning(enable))?; + } + for (key, value) in &self.config_entries { + builder = best_effort(builder, |b| b.with_config_entry(key, value))?; + } + + Ok(builder) + } +} + +// Applies one session option, keeping the builder unchanged if the active backend +// has no implementation for it. `ort-tract` implements only the graph optimization +// level; every other session option resolves to `ort-sys`' stub API and reports +// `NotImplemented`. Since tract already runs single-threaded and never spin-waits, +// skipping such an option is closer to what the caller asked for than refusing to +// build the session at all. +fn best_effort( + builder: SessionBuilder, + f: impl FnOnce(SessionBuilder) -> ort::Result, +) -> Result { + let unchanged = builder.clone(); + match f(builder) { + Ok(builder) => Ok(builder), + Err(err) if err.code() == ort::ErrorCode::NotImplemented => Ok(unchanged), + Err(err) => Err(err.into()), + } +} + +pub(crate) fn build_session_from_memory( + bytes: &[u8], + options: &SessionOptions, +) -> Result { #[cfg(use_tract)] ensure_tract_backend(); - Ok(Session::builder()?.commit_from_memory(bytes)?) + Ok(options.apply(Session::builder()?)?.commit_from_memory(bytes)?) } -pub(crate) fn build_session_from_file(path: impl AsRef) -> Result { +pub(crate) fn build_session_from_file( + path: impl AsRef, + options: &SessionOptions, +) -> Result { #[cfg(use_tract)] ensure_tract_backend(); let bytes = std::fs::read(path)?; - Ok(Session::builder()?.commit_from_memory(&bytes)?) + Ok(options.apply(Session::builder()?)?.commit_from_memory(&bytes)?) } diff --git a/livekit-wakeword/src/melspectrogram.rs b/livekit-wakeword/src/melspectrogram.rs index e677513a6..f0718dd37 100644 --- a/livekit-wakeword/src/melspectrogram.rs +++ b/livekit-wakeword/src/melspectrogram.rs @@ -16,7 +16,7 @@ use ndarray::{Array1, Array2, Axis}; use ort::session::Session; use ort::value::Tensor; -use crate::{build_session_from_memory, WakeWordError}; +use crate::{build_session_from_memory, SessionOptions, WakeWordError}; const MODEL_BYTES: &[u8] = include_bytes!("../onnx/melspectrogram.onnx"); @@ -32,8 +32,8 @@ pub struct MelspectrogramModel { } impl MelspectrogramModel { - pub fn new() -> Result { - Ok(Self { session: build_session_from_memory(MODEL_BYTES)? }) + pub fn new(options: &SessionOptions) -> Result { + Ok(Self { session: build_session_from_memory(MODEL_BYTES, options)? }) } // Run the melspectrogram model on normalized f32 audio and return mel features. diff --git a/livekit-wakeword/src/wakeword.rs b/livekit-wakeword/src/wakeword.rs index 1002be6e1..58a4e2db3 100644 --- a/livekit-wakeword/src/wakeword.rs +++ b/livekit-wakeword/src/wakeword.rs @@ -23,8 +23,8 @@ use resampler::{Attenuation, Latency, ResamplerFir, SampleRate}; use crate::embedding::EmbeddingModel; use crate::melspectrogram::MelspectrogramModel; use crate::{ - build_session_from_file, to_resampler_rate, WakeWordError, EMBEDDING_STRIDE, EMBEDDING_WINDOW, - MIN_EMBEDDINGS, + build_session_from_file, to_resampler_rate, SessionOptions, WakeWordError, EMBEDDING_STRIDE, + EMBEDDING_WINDOW, MIN_EMBEDDINGS, }; struct Resampler { @@ -45,6 +45,7 @@ pub struct WakeWordModel { emb_model: EmbeddingModel, classifiers: HashMap, resampler: Option, + session_options: SessionOptions, } impl WakeWordModel { @@ -53,7 +54,23 @@ impl WakeWordModel { /// The recommended sample rate is 16 kHz. Other supported rates /// (22050, 32000, 44100, 48000, 88200, 96000, 176400, 192000, 384000 Hz) /// are resampled internally to 16 kHz. + /// + /// ONNX sessions are created with [`SessionOptions::default`]; use + /// [`with_session_options`](Self::with_session_options) to tune them. pub fn new(models: &[impl AsRef], sample_rate: u32) -> Result { + Self::with_session_options(models, sample_rate, SessionOptions::default()) + } + + /// Create a new wake word model, tuning the ONNX sessions it creates. + /// + /// `session_options` applies to the bundled feature extraction models, to every + /// classifier in `models`, and to any classifier a later + /// [`load_model`](Self::load_model) call adds. + pub fn with_session_options( + models: &[impl AsRef], + sample_rate: u32, + session_options: SessionOptions, + ) -> Result { let resampler = if sample_rate != 16000 { let input_rate = to_resampler_rate(sample_rate)?; // FIR resampler: 64-sample latency (~1.3ms at 48kHz) with 90dB @@ -72,10 +89,11 @@ impl WakeWordModel { }; let mut wakeword = Self { - mel_model: MelspectrogramModel::new()?, - emb_model: EmbeddingModel::new()?, + mel_model: MelspectrogramModel::new(&session_options)?, + emb_model: EmbeddingModel::new(&session_options)?, classifiers: HashMap::new(), resampler, + session_options, }; for path in models { @@ -88,10 +106,27 @@ impl WakeWordModel { /// Load a wake word classifier ONNX model from disk. /// /// If `model_name` is `None`, the file stem is used as the classifier name. + /// + /// The session is created with the options this model was built with. pub fn load_model( &mut self, model_path: impl AsRef, model_name: Option<&str>, + ) -> Result<(), WakeWordError> { + let session_options = self.session_options.clone(); + self.load_model_with_session_options(model_path, model_name, &session_options) + } + + /// Load a wake word classifier ONNX model from disk with its own session options. + /// + /// Only this classifier's session is affected; the model's own options, used by + /// [`load_model`](Self::load_model) and by the bundled feature extraction + /// models, are left alone. + pub fn load_model_with_session_options( + &mut self, + model_path: impl AsRef, + model_name: Option<&str>, + session_options: &SessionOptions, ) -> Result<(), WakeWordError> { let path = model_path.as_ref(); if !path.exists() { @@ -103,7 +138,7 @@ impl WakeWordModel { None => path.file_stem().and_then(|s| s.to_str()).unwrap_or("unknown").to_string(), }; - let session = build_session_from_file(path)?; + let session = build_session_from_file(path, session_options)?; self.classifiers.insert(name, session); Ok(()) } diff --git a/livekit-wakeword/tests/integration.rs b/livekit-wakeword/tests/integration.rs index 9afb3271d..48588bc3f 100644 --- a/livekit-wakeword/tests/integration.rs +++ b/livekit-wakeword/tests/integration.rs @@ -14,7 +14,7 @@ use std::path::PathBuf; -use livekit_wakeword::{WakeWordModel, SAMPLE_RATE}; +use livekit_wakeword::{OptimizationLevel, SessionOptions, WakeWordModel, SAMPLE_RATE}; mod common; @@ -82,3 +82,59 @@ fn test_negative_wav_below_threshold() { "expected negative sample score ({score}) < threshold ({THRESHOLD})" ); } + +/// Session options supplied at runtime must not change what the model predicts, +/// including the options a backend cannot honour (`ort-tract` implements only the +/// graph optimization level and skips the rest). +#[test] +fn test_session_options_preserve_scores() { + let (sample_rate, samples) = common::read_wav("positive.wav"); + + let options = SessionOptions { + intra_threads: Some(1), + inter_threads: Some(1), + parallel_execution: Some(false), + intra_op_spinning: Some(false), + inter_op_spinning: Some(false), + config_entries: vec![("session.use_env_allocators".to_string(), "1".to_string())], + ..Default::default() + }; + + let mut default = WakeWordModel::new(&[classifier_path()], sample_rate).unwrap(); + let mut tuned = + WakeWordModel::with_session_options(&[classifier_path()], sample_rate, options).unwrap(); + + assert_eq!( + default.predict(&samples).unwrap()["hey_livekit"], + tuned.predict(&samples).unwrap()["hey_livekit"] + ); +} + +/// Every optimization level builds a working model, and a per-classifier override +/// leaves the model's own options in place. +#[test] +fn test_optimization_levels_and_per_model_options() { + let (sample_rate, samples) = common::read_wav("positive.wav"); + + for level in [OptimizationLevel::Disable, OptimizationLevel::Level1, OptimizationLevel::Level3] + { + let options = SessionOptions { optimization_level: level, ..Default::default() }; + let mut model = + WakeWordModel::with_session_options(&[classifier_path()], sample_rate, options) + .unwrap(); + let score = model.predict(&samples).unwrap()["hey_livekit"]; + assert!(score > THRESHOLD, "{level:?} scored {score}"); + } + + let no_models: &[PathBuf] = &[]; + let mut model = WakeWordModel::new(no_models, sample_rate).unwrap(); + let options = + SessionOptions { optimization_level: OptimizationLevel::Disable, ..Default::default() }; + model + .load_model_with_session_options(classifier_path(), Some("unoptimized"), &options) + .unwrap(); + model.load_model(classifier_path(), Some("default")).unwrap(); + + let predictions = model.predict(&samples).unwrap(); + assert_eq!(predictions["unoptimized"], predictions["default"]); +} From a3ea7784ec51d493cce4accaf72ae68a23968bb3 Mon Sep 17 00:00:00 2001 From: pham-tuan-binh Date: Fri, 28 Aug 2026 22:41:02 +0700 Subject: [PATCH 2/2] livekit-wakeword: simplify the session options plumbing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `OptimizationLevel` mirrored `ort`'s `GraphOptimizationLevel` one variant at a time to add a `Default`. Since `ort::Error` is already part of the crate's public API there was no abstraction left to protect, so the enum is now re-exported — the way `livekit` re-exports libwebrtc's `DegradationPreference` — and `Level3` moves into a hand-written `Default for SessionOptions`. `best_effort` cloned the `SessionBuilder` before every option so it could roll back on `NotImplemented`, but the backend is not a runtime property: `build.rs` emits `use_tract` for every target except aarch64 Windows, and `ensure_tract_backend` is already gated on it. `apply` is now two `#[cfg]` bodies — tract sets the optimization level and nothing else, real ONNX Runtime applies every option and treats a failure as the error it is. `load_model_with_session_options` is dropped. Per-classifier tuning is speculative when the classifiers are the tiny models, and it was the only reason `load_model` cloned the stored options to satisfy the borrow checker; `test_load_model_inherits_session_options` keeps the post-construction `load_model` path covered. `build_session_from_file` delegates to `build_session_from_memory` rather than repeating the tract init and commit. --- .../wakeword_runtime_session_options.md | 24 ++-- livekit-wakeword/src/lib.rs | 125 +++++++----------- livekit-wakeword/src/wakeword.rs | 22 +-- livekit-wakeword/tests/integration.rs | 36 ++--- 4 files changed, 86 insertions(+), 121 deletions(-) diff --git a/.changeset/wakeword_runtime_session_options.md b/.changeset/wakeword_runtime_session_options.md index 4969e6031..71cce0e0b 100644 --- a/.changeset/wakeword_runtime_session_options.md +++ b/.changeset/wakeword_runtime_session_options.md @@ -4,13 +4,17 @@ livekit-wakeword: minor # Runtime-configurable ONNX session options -`WakeWordModel::with_session_options` and `load_model_with_session_options` accept a -`SessionOptions` describing how the ONNX sessions are created — graph optimization -level, intra/inter-op threads, sequential execution, thread spinning, and arbitrary -session config entries — matching the `sess_options` parameter the Python SDK takes. -The default is `SessionOptions::default()`, whose `OptimizationLevel::Level3` matches -ONNX Runtime's own default; the `ort-tract` backend used on every target except -aarch64 Windows runs tract's `into_optimized()` only when a level is requested, so -requesting one made `predict()` over a 2 s window 6.2x faster (329 ms to 53 ms -median, Apple M-series release build). Options that the active backend does not -implement are skipped rather than failing session creation. +`WakeWordModel::with_session_options` takes a `SessionOptions` describing how the +crate creates its ONNX sessions — graph optimization level, intra/inter-op threads, +sequential execution, thread spinning, and arbitrary session config entries — +matching the `sess_options` parameter the Python SDK takes. It applies to the two +bundled feature extraction models and to every wake word classifier, including ones +a later `load_model` call adds. + +`WakeWordModel::new` keeps its signature and uses `SessionOptions::default()`, whose +`GraphOptimizationLevel::Level3` matches ONNX Runtime's own default. That default +matters: the `ort-tract` backend used on every target except aarch64 Windows runs +tract's `into_optimized()` only when a level is requested, so requesting one made +`predict()` over a 2 s window 6.2x faster (329 ms to 53 ms median, Apple M-series +release build). tract implements no other session option, so the remaining fields +are skipped there rather than failing session creation. diff --git a/livekit-wakeword/src/lib.rs b/livekit-wakeword/src/lib.rs index ac3f098d6..cf22edd62 100644 --- a/livekit-wakeword/src/lib.rs +++ b/livekit-wakeword/src/lib.rs @@ -16,10 +16,7 @@ use std::path::Path; #[cfg(use_tract)] use std::sync::Once; -use ort::session::{ - builder::{GraphOptimizationLevel, SessionBuilder}, - Session, -}; +use ort::session::{builder::SessionBuilder, Session}; #[cfg(use_tract)] static INIT_TRACT: Once = Once::new(); @@ -35,6 +32,8 @@ pub(crate) mod embedding; pub(crate) mod melspectrogram; pub mod wakeword; +/// Graph optimizations requested when an ONNX session is created. +pub use ort::session::builder::GraphOptimizationLevel; pub use wakeword::WakeWordModel; #[derive(Debug, thiserror::Error)] @@ -77,39 +76,9 @@ pub(crate) fn to_resampler_rate(hz: u32) -> Result for GraphOptimizationLevel { - fn from(level: OptimizationLevel) -> Self { - match level { - OptimizationLevel::Disable => GraphOptimizationLevel::Disable, - OptimizationLevel::Level1 => GraphOptimizationLevel::Level1, - OptimizationLevel::Level2 => GraphOptimizationLevel::Level2, - OptimizationLevel::Level3 => GraphOptimizationLevel::Level3, - } - } -} - -/// Tuning applied to every ONNX session a [`WakeWordModel`] creates: the two -/// bundled feature extraction models and each wake word classifier. +/// How every ONNX session a [`WakeWordModel`] creates is configured: the two +/// bundled feature extraction models and each wake word classifier. Mirrors the +/// `sess_options` parameter the Python SDK takes. /// /// [`Default`] is what [`WakeWordModel::new`] uses. Pass a value of your own to /// [`WakeWordModel::with_session_options`] to trade latency for CPU, which is @@ -131,15 +100,16 @@ impl From for GraphOptimizationLevel { /// # } /// ``` /// -/// Every field except `optimization_level` is best-effort: a backend that does not -/// implement an option has it skipped rather than failing session creation. In -/// particular the `ort-tract` backend implements only the optimization level, and -/// runs single-threaded without spin-waiting, so the thread and spinning fields -/// above describe what it already does. -#[derive(Clone, Debug, Default)] +/// `optimization_level` is the only option the `ort-tract` backend used on every +/// target except aarch64 Windows implements; the rest are skipped there rather than +/// failing session creation, and since tract runs single-threaded without +/// spin-waiting, they already describe what it does. +#[derive(Clone, Debug)] pub struct SessionOptions { - /// Graph optimizations to apply when the session is created. - pub optimization_level: OptimizationLevel, + /// Graph optimizations to apply when the session is created. Matters more than + /// it looks: `ort-tract` runs tract's `into_optimized()` only when some level is + /// requested, and wake word inference measured several times slower without it. + pub optimization_level: GraphOptimizationLevel, /// Threads used to run a single operator. `Some(1)` keeps each session on the /// calling thread. pub intra_threads: Option, @@ -158,53 +128,57 @@ pub struct SessionOptions { pub config_entries: Vec<(String, String)>, } +impl Default for SessionOptions { + /// Requests [`GraphOptimizationLevel::Level3`], ONNX Runtime's own default, and + /// leaves every other option to the backend. + fn default() -> Self { + Self { + optimization_level: GraphOptimizationLevel::Level3, + intra_threads: None, + inter_threads: None, + parallel_execution: None, + intra_op_spinning: None, + inter_op_spinning: None, + config_entries: Vec::new(), + } + } +} + impl SessionOptions { + // The optimization level is the only session option `ort-tract` implements; + // every other one resolves to `ort-sys`' stub API and reports `NotImplemented`. + // Since tract runs single-threaded without spin-waiting, skipping those is + // closer to what the caller asked for than refusing to build the session. + #[cfg(use_tract)] fn apply(&self, builder: SessionBuilder) -> Result { - // Unlike everything below, the optimization level is implemented by every - // backend, so a failure here is a real error rather than an unsupported knob. - let mut builder = builder.with_optimization_level(self.optimization_level.into())?; + Ok(builder.with_optimization_level(self.optimization_level)?) + } + #[cfg(not(use_tract))] + fn apply(&self, builder: SessionBuilder) -> Result { + let mut builder = builder.with_optimization_level(self.optimization_level)?; if let Some(threads) = self.intra_threads { - builder = best_effort(builder, |b| b.with_intra_threads(threads))?; + builder = builder.with_intra_threads(threads)?; } if let Some(threads) = self.inter_threads { - builder = best_effort(builder, |b| b.with_inter_threads(threads))?; + builder = builder.with_inter_threads(threads)?; } if let Some(parallel) = self.parallel_execution { - builder = best_effort(builder, |b| b.with_parallel_execution(parallel))?; + builder = builder.with_parallel_execution(parallel)?; } if let Some(enable) = self.intra_op_spinning { - builder = best_effort(builder, |b| b.with_intra_op_spinning(enable))?; + builder = builder.with_intra_op_spinning(enable)?; } if let Some(enable) = self.inter_op_spinning { - builder = best_effort(builder, |b| b.with_inter_op_spinning(enable))?; + builder = builder.with_inter_op_spinning(enable)?; } for (key, value) in &self.config_entries { - builder = best_effort(builder, |b| b.with_config_entry(key, value))?; + builder = builder.with_config_entry(key, value)?; } - Ok(builder) } } -// Applies one session option, keeping the builder unchanged if the active backend -// has no implementation for it. `ort-tract` implements only the graph optimization -// level; every other session option resolves to `ort-sys`' stub API and reports -// `NotImplemented`. Since tract already runs single-threaded and never spin-waits, -// skipping such an option is closer to what the caller asked for than refusing to -// build the session at all. -fn best_effort( - builder: SessionBuilder, - f: impl FnOnce(SessionBuilder) -> ort::Result, -) -> Result { - let unchanged = builder.clone(); - match f(builder) { - Ok(builder) => Ok(builder), - Err(err) if err.code() == ort::ErrorCode::NotImplemented => Ok(unchanged), - Err(err) => Err(err.into()), - } -} - pub(crate) fn build_session_from_memory( bytes: &[u8], options: &SessionOptions, @@ -218,8 +192,5 @@ pub(crate) fn build_session_from_file( path: impl AsRef, options: &SessionOptions, ) -> Result { - #[cfg(use_tract)] - ensure_tract_backend(); - let bytes = std::fs::read(path)?; - Ok(options.apply(Session::builder()?)?.commit_from_memory(&bytes)?) + build_session_from_memory(&std::fs::read(path)?, options) } diff --git a/livekit-wakeword/src/wakeword.rs b/livekit-wakeword/src/wakeword.rs index 58a4e2db3..bc658162a 100644 --- a/livekit-wakeword/src/wakeword.rs +++ b/livekit-wakeword/src/wakeword.rs @@ -105,28 +105,12 @@ impl WakeWordModel { /// Load a wake word classifier ONNX model from disk. /// - /// If `model_name` is `None`, the file stem is used as the classifier name. - /// - /// The session is created with the options this model was built with. + /// If `model_name` is `None`, the file stem is used as the classifier name. The + /// session is created with the [`SessionOptions`] this model was built with. pub fn load_model( &mut self, model_path: impl AsRef, model_name: Option<&str>, - ) -> Result<(), WakeWordError> { - let session_options = self.session_options.clone(); - self.load_model_with_session_options(model_path, model_name, &session_options) - } - - /// Load a wake word classifier ONNX model from disk with its own session options. - /// - /// Only this classifier's session is affected; the model's own options, used by - /// [`load_model`](Self::load_model) and by the bundled feature extraction - /// models, are left alone. - pub fn load_model_with_session_options( - &mut self, - model_path: impl AsRef, - model_name: Option<&str>, - session_options: &SessionOptions, ) -> Result<(), WakeWordError> { let path = model_path.as_ref(); if !path.exists() { @@ -138,7 +122,7 @@ impl WakeWordModel { None => path.file_stem().and_then(|s| s.to_str()).unwrap_or("unknown").to_string(), }; - let session = build_session_from_file(path, session_options)?; + let session = build_session_from_file(path, &self.session_options)?; self.classifiers.insert(name, session); Ok(()) } diff --git a/livekit-wakeword/tests/integration.rs b/livekit-wakeword/tests/integration.rs index 48588bc3f..2fcbd656e 100644 --- a/livekit-wakeword/tests/integration.rs +++ b/livekit-wakeword/tests/integration.rs @@ -14,7 +14,7 @@ use std::path::PathBuf; -use livekit_wakeword::{OptimizationLevel, SessionOptions, WakeWordModel, SAMPLE_RATE}; +use livekit_wakeword::{GraphOptimizationLevel, SessionOptions, WakeWordModel, SAMPLE_RATE}; mod common; @@ -110,14 +110,16 @@ fn test_session_options_preserve_scores() { ); } -/// Every optimization level builds a working model, and a per-classifier override -/// leaves the model's own options in place. +/// Every optimization level builds a model that still detects the wake word. #[test] -fn test_optimization_levels_and_per_model_options() { +fn test_optimization_levels() { let (sample_rate, samples) = common::read_wav("positive.wav"); - for level in [OptimizationLevel::Disable, OptimizationLevel::Level1, OptimizationLevel::Level3] - { + for level in [ + GraphOptimizationLevel::Disable, + GraphOptimizationLevel::Level1, + GraphOptimizationLevel::Level3, + ] { let options = SessionOptions { optimization_level: level, ..Default::default() }; let mut model = WakeWordModel::with_session_options(&[classifier_path()], sample_rate, options) @@ -125,16 +127,20 @@ fn test_optimization_levels_and_per_model_options() { let score = model.predict(&samples).unwrap()["hey_livekit"]; assert!(score > THRESHOLD, "{level:?} scored {score}"); } +} + +/// A classifier loaded after construction inherits the model's session options. +#[test] +fn test_load_model_inherits_session_options() { + let (sample_rate, samples) = common::read_wav("positive.wav"); let no_models: &[PathBuf] = &[]; - let mut model = WakeWordModel::new(no_models, sample_rate).unwrap(); - let options = - SessionOptions { optimization_level: OptimizationLevel::Disable, ..Default::default() }; - model - .load_model_with_session_options(classifier_path(), Some("unoptimized"), &options) - .unwrap(); - model.load_model(classifier_path(), Some("default")).unwrap(); + let options = SessionOptions { + optimization_level: GraphOptimizationLevel::Disable, + ..Default::default() + }; + let mut model = WakeWordModel::with_session_options(no_models, sample_rate, options).unwrap(); + model.load_model(classifier_path(), Some("late")).unwrap(); - let predictions = model.predict(&samples).unwrap(); - assert_eq!(predictions["unoptimized"], predictions["default"]); + assert!(model.predict(&samples).unwrap()["late"] > THRESHOLD); }