diff --git a/Makefile b/Makefile index 4bc18d399..8c7e7e520 100644 --- a/Makefile +++ b/Makefile @@ -167,21 +167,21 @@ check-rust-fmt: check-rust-msrv: @$(REQUIRE_CARGO) @python3 scripts/check_release_preflight.py msrv --toolchain 1.95.0 - RUSTUP_TOOLCHAIN=1.95.0 cargo check --manifest-path $(RUST_MANIFEST) --workspace --locked + RUSTUP_TOOLCHAIN=1.95.0 python3 scripts/resolve_onnxruntime_capi.py --exec -- cargo check --manifest-path $(RUST_MANIFEST) --workspace --locked check-rust-clippy: @$(REQUIRE_CARGO) - cargo clippy --manifest-path $(RUST_MANIFEST) --workspace --all-targets --locked -- -D warnings + python3 scripts/resolve_onnxruntime_capi.py --exec -- cargo clippy --manifest-path $(RUST_MANIFEST) --workspace --all-targets --locked -- -D warnings check-rust-test: @$(REQUIRE_CARGO) - cargo test --manifest-path $(RUST_MANIFEST) --workspace --locked + python3 scripts/resolve_onnxruntime_capi.py --exec -- cargo test --manifest-path $(RUST_MANIFEST) --workspace --locked check-rust-ios: @$(REQUIRE_CARGO) @$(REQUIRE_RUSTUP) @rustup target list --installed 2>/dev/null | grep -qx "$(IOS_TARGET)" || { echo "Rust target $(IOS_TARGET) is required for the iOS gate; run rustup target add $(IOS_TARGET)" >&2; exit 1; } - cargo check --manifest-path $(RUST_MANIFEST) --workspace --exclude solstone-core --exclude solstone-core-indexer-store --lib --target $(IOS_TARGET) --locked + cargo check --manifest-path $(RUST_MANIFEST) --workspace --exclude solstone-core --exclude solstone-core-indexer-store --exclude solstone-core-speakers-onnx --lib --target $(IOS_TARGET) --locked check-rust-deny: @$(REQUIRE_CARGO) diff --git a/core/Cargo.lock b/core/Cargo.lock index 0ee4c3962..8f669532b 100644 --- a/core/Cargo.lock +++ b/core/Cargo.lock @@ -43,9 +43,9 @@ checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" [[package]] name = "cc" -version = "1.2.67" +version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e17dd265a7d0f31ef544e1b20e03add05d3b45b491b633b10d67145d2acc1a38" +checksum = "5add81bb678e6cb321aff7fa0dc7689ad82b112dbc032cea19f91d6b8e3582b9" dependencies = [ "find-msvc-tools", "shlex", @@ -116,21 +116,21 @@ checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" [[package]] name = "futures-core" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" [[package]] name = "futures-task" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" [[package]] name = "futures-util" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" dependencies = [ "futures-core", "futures-task", @@ -140,9 +140,9 @@ dependencies = [ [[package]] name = "glob" -version = "0.3.3" +version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" [[package]] name = "hashbrown" @@ -240,9 +240,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.186" +version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" [[package]] name = "libsqlite3-sys" @@ -261,6 +261,16 @@ version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" +[[package]] +name = "matrixmultiply" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "md5" version = "0.8.1" @@ -273,6 +283,39 @@ version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -288,6 +331,24 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "ort" +version = "2.0.0-rc.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7de3af33d24a745ffb8fab904b13478438d1cd52868e6f17735ef6e1f8bf133" +dependencies = [ + "ndarray", + "ort-sys", + "smallvec", + "tracing", +] + +[[package]] +name = "ort-sys" +version = "2.0.0-rc.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7b497d21a8b6fbb4b5a544f8fadb77e801a09ae0add9e411d31c6f89e3c1e90" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -324,11 +385,26 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" +[[package]] +name = "portable-atomic" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" + +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + [[package]] name = "proc-macro2" -version = "1.0.106" +version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" dependencies = [ "unicode-ident", ] @@ -346,9 +422,9 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.46" +version = "1.0.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" dependencies = [ "proc-macro2", ] @@ -359,6 +435,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "270e04e5ea61d40841942bb15e451c29ee1618637bcf97fc7ede5dd4a9b1601b" +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "rsqlite-vfs" version = "0.1.1" @@ -392,38 +474,38 @@ checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" [[package]] name = "serde" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" dependencies = [ "serde_core", ] [[package]] name = "serde_core" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.3", ] [[package]] name = "serde_json" -version = "1.0.150" +version = "1.0.151" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" dependencies = [ "indexmap", "itoa", @@ -533,6 +615,22 @@ dependencies = [ "solstone-core-sol-client", ] +[[package]] +name = "solstone-core-speakers" +version = "1.0.17" +dependencies = [ + "serde_json", +] + +[[package]] +name = "solstone-core-speakers-onnx" +version = "1.0.17" +dependencies = [ + "ort", + "serde_json", + "solstone-core-speakers", +] + [[package]] name = "sqlite-wasm-rs" version = "0.5.5" @@ -556,24 +654,35 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "thiserror" -version = "2.0.18" +version = "2.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.18" +version = "2.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.3", ] [[package]] @@ -608,6 +717,25 @@ dependencies = [ "winnow", ] +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-core", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + [[package]] name = "unicase" version = "2.9.0" @@ -704,7 +832,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn", + "syn 2.0.119", "wasm-bindgen-shared", ] @@ -738,7 +866,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -749,7 +877,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] diff --git a/core/Cargo.toml b/core/Cargo.toml index 52e02009e..d56cd7eab 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -8,6 +8,8 @@ members = [ "crates/solstone-core-sol", "crates/solstone-core-sol-client", "crates/solstone-core-sol-client-cli", + "crates/solstone-core-speakers", + "crates/solstone-core-speakers-onnx", ] resolver = "3" @@ -26,12 +28,15 @@ solstone-core-indexer = { path = "crates/solstone-core-indexer" } solstone-core-indexer-store = { path = "crates/solstone-core-indexer-store" } solstone-core-journal = { path = "crates/solstone-core-journal" } solstone-core-sol = { path = "crates/solstone-core-sol" } -solstone-core-sol-client-cli = { path = "crates/solstone-core-sol-client-cli" } solstone-core-sol-client = { path = "crates/solstone-core-sol-client" } +solstone-core-sol-client-cli = { path = "crates/solstone-core-sol-client-cli" } +solstone-core-speakers = { path = "crates/solstone-core-speakers" } +solstone-core-speakers-onnx = { path = "crates/solstone-core-speakers-onnx" } chrono = { version = "0.4.45", default-features = false, features = ["clock"] } chrono-tz = "0.10.4" glob = "0.3.3" md5 = "0.8.1" +ort = { version = "=2.0.0-rc.12", default-features = false } pulldown-cmark = { version = "0.13.4", default-features = false } rapidfuzz = "0.5.0" rusqlite = { version = "0.40.1", features = ["bundled"] } diff --git a/core/crates/solstone-core-speakers-onnx/Cargo.toml b/core/crates/solstone-core-speakers-onnx/Cargo.toml new file mode 100644 index 000000000..753cea04b --- /dev/null +++ b/core/crates/solstone-core-speakers-onnx/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "solstone-core-speakers-onnx" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +publish = false + +[lints] +workspace = true + +[dependencies] +ort = { workspace = true, features = ["api-24", "std"] } +solstone-core-speakers.workspace = true + +[target.'cfg(target_vendor = "apple")'.dependencies] +ort = { workspace = true, features = ["coreml"] } + +[dev-dependencies] +serde_json.workspace = true diff --git a/core/crates/solstone-core-speakers-onnx/src/lib.rs b/core/crates/solstone-core-speakers-onnx/src/lib.rs new file mode 100644 index 000000000..cd43298e3 --- /dev/null +++ b/core/crates/solstone-core-speakers-onnx/src/lib.rs @@ -0,0 +1,395 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright (c) 2026 sol pbc + +use std::error::Error; +use std::fmt; +use std::path::Path; + +use ort::ep::{CPU, CoreML, ExecutionProviderDispatch}; +use ort::session::Session; +use ort::value::{Tensor, TensorElementType, ValueType}; +use solstone_core_speakers::{FeatureMatrix, WESPEAKER_EMBEDDING_SIZE, WESPEAKER_MEL_BINS}; + +const INPUT_NAME: &str = "feats"; +const OUTPUT_NAME: &str = "embs"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PlatformFamily { + Apple, + Other, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PlatformDescriptor { + pub family: PlatformFamily, +} + +impl PlatformDescriptor { + pub fn current() -> Self { + Self { + family: if cfg!(target_vendor = "apple") { + PlatformFamily::Apple + } else { + PlatformFamily::Other + }, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SpeakerExecutionProvider { + CoreMl, + Cpu, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct SpeakerEmbedding { + values: [f32; WESPEAKER_EMBEDDING_SIZE], +} + +impl SpeakerEmbedding { + pub fn values(&self) -> &[f32; WESPEAKER_EMBEDDING_SIZE] { + &self.values + } +} + +#[derive(Debug)] +pub struct WespeakerEmbedder { + session: Session, + input_name: String, + output_name: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SpeakerOnnxError { + EmptyProviderPlan, + ProviderUnavailable { provider: &'static str }, + InvalidFeatureMatrix { frames: usize, bins: usize }, + InvalidModelIo { detail: String }, + MissingOutput { name: String }, + Ort { message: String }, +} + +impl fmt::Display for SpeakerOnnxError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::EmptyProviderPlan => formatter.write_str("speaker ONNX provider plan is empty"), + Self::ProviderUnavailable { provider } => { + write!( + formatter, + "speaker ONNX provider is unavailable: {provider}" + ) + } + Self::InvalidFeatureMatrix { frames, bins } => write!( + formatter, + "speaker ONNX features must have at least one frame and {WESPEAKER_MEL_BINS} bins, got frames={frames} bins={bins}" + ), + Self::InvalidModelIo { detail } => { + write!(formatter, "speaker ONNX model IO mismatch: {detail}") + } + Self::MissingOutput { name } => { + write!(formatter, "speaker ONNX output was not returned: {name}") + } + Self::Ort { message } => write!(formatter, "speaker ONNX Runtime error: {message}"), + } + } +} + +impl Error for SpeakerOnnxError {} + +impl From> for SpeakerOnnxError { + fn from(error: ort::Error) -> Self { + Self::Ort { + message: error.to_string(), + } + } +} + +pub fn default_speaker_execution_providers( + platform: PlatformDescriptor, +) -> Vec { + match platform.family { + PlatformFamily::Apple => vec![ + SpeakerExecutionProvider::CoreMl, + SpeakerExecutionProvider::Cpu, + ], + PlatformFamily::Other => vec![SpeakerExecutionProvider::Cpu], + } +} + +impl WespeakerEmbedder { + pub fn open( + model_path: &Path, + providers: &[SpeakerExecutionProvider], + ) -> Result { + if providers.is_empty() { + return Err(SpeakerOnnxError::EmptyProviderPlan); + } + let dispatches = provider_dispatches(providers)?; + let session = Session::builder()? + .with_execution_providers(dispatches)? + .commit_from_file(model_path)?; + validate_session_io(&session)?; + Ok(Self { + session, + input_name: INPUT_NAME.to_string(), + output_name: OUTPUT_NAME.to_string(), + }) + } + + pub fn embed( + &mut self, + features: &FeatureMatrix, + ) -> Result { + if features.frames() == 0 || features.bins() != WESPEAKER_MEL_BINS { + return Err(SpeakerOnnxError::InvalidFeatureMatrix { + frames: features.frames(), + bins: features.bins(), + }); + } + let input = Tensor::from_array(( + [1_usize, features.frames(), WESPEAKER_MEL_BINS], + features.data().to_vec().into_boxed_slice(), + ))?; + let mut outputs = self + .session + .run(ort::inputs![self.input_name.as_str() => input])?; + let output = + outputs + .remove(&self.output_name) + .ok_or_else(|| SpeakerOnnxError::MissingOutput { + name: self.output_name.clone(), + })?; + let (shape, values) = output.try_extract_tensor::()?; + if shape[..] != [1, WESPEAKER_EMBEDDING_SIZE as i64] { + return Err(SpeakerOnnxError::InvalidModelIo { + detail: format!("output shape {shape} is not [1, {WESPEAKER_EMBEDDING_SIZE}]"), + }); + } + let mut embedding = [0.0; WESPEAKER_EMBEDDING_SIZE]; + embedding.copy_from_slice(values); + Ok(SpeakerEmbedding { values: embedding }) + } +} + +fn provider_dispatches( + providers: &[SpeakerExecutionProvider], +) -> Result, SpeakerOnnxError> { + let mut dispatches = Vec::with_capacity(providers.len()); + for provider in providers { + match provider { + SpeakerExecutionProvider::CoreMl => { + if !cfg!(target_vendor = "apple") { + return Err(SpeakerOnnxError::ProviderUnavailable { provider: "coreml" }); + } + dispatches.push(CoreML::default().build()); + } + SpeakerExecutionProvider::Cpu => { + dispatches.push(CPU::default().build()); + } + } + } + Ok(dispatches) +} + +fn validate_session_io(session: &Session) -> Result<(), SpeakerOnnxError> { + let inputs = session.inputs(); + let outputs = session.outputs(); + if inputs.len() != 1 { + return Err(SpeakerOnnxError::InvalidModelIo { + detail: format!("expected one input, got {}", inputs.len()), + }); + } + if outputs.len() != 1 { + return Err(SpeakerOnnxError::InvalidModelIo { + detail: format!("expected one output, got {}", outputs.len()), + }); + } + expect_tensor( + "input", + inputs[0].name(), + inputs[0].dtype(), + INPUT_NAME, + &[ + ExpectedDim::Any, + ExpectedDim::Any, + ExpectedDim::Exact(WESPEAKER_MEL_BINS as i64), + ], + )?; + expect_tensor( + "output", + outputs[0].name(), + outputs[0].dtype(), + OUTPUT_NAME, + &[ + ExpectedDim::Any, + ExpectedDim::Exact(WESPEAKER_EMBEDDING_SIZE as i64), + ], + )?; + Ok(()) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ExpectedDim { + Any, + Exact(i64), +} + +fn expect_tensor( + label: &str, + name: &str, + value_type: &ValueType, + expected_name: &str, + expected_shape: &[ExpectedDim], +) -> Result<(), SpeakerOnnxError> { + if name != expected_name { + return Err(SpeakerOnnxError::InvalidModelIo { + detail: format!("{label} name {name:?} is not {expected_name:?}"), + }); + } + let ValueType::Tensor { ty, shape, .. } = value_type else { + return Err(SpeakerOnnxError::InvalidModelIo { + detail: format!("{label} {name:?} is not a tensor"), + }); + }; + if *ty != TensorElementType::Float32 { + return Err(SpeakerOnnxError::InvalidModelIo { + detail: format!("{label} {name:?} is {ty}, not float32"), + }); + } + if shape.len() != expected_shape.len() { + return Err(SpeakerOnnxError::InvalidModelIo { + detail: format!("{label} {name:?} shape {shape} has wrong rank"), + }); + } + for (index, (actual, expected)) in shape.iter().zip(expected_shape).enumerate() { + match expected { + ExpectedDim::Any => {} + ExpectedDim::Exact(value) if actual == value => {} + ExpectedDim::Exact(value) => { + return Err(SpeakerOnnxError::InvalidModelIo { + detail: format!("{label} {name:?} dim {index} is {actual}, not {value}"), + }); + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use solstone_core_speakers::{WESPEAKER_SAMPLE_RATE_HZ, compute_wespeaker_filterbank_cmn}; + + const FIXTURE: &str = include_str!("../../../fixtures/speaker_filterbank.json"); + + #[test] + fn provider_plan_selects_coreml_then_cpu_for_synthetic_apple() { + let plan = default_speaker_execution_providers(PlatformDescriptor { + family: PlatformFamily::Apple, + }); + + assert_eq!( + plan, + vec![ + SpeakerExecutionProvider::CoreMl, + SpeakerExecutionProvider::Cpu + ] + ); + } + + #[test] + fn provider_plan_selects_cpu_for_non_apple() { + let plan = default_speaker_execution_providers(PlatformDescriptor { + family: PlatformFamily::Other, + }); + + assert_eq!(plan, vec![SpeakerExecutionProvider::Cpu]); + } + + #[test] + fn coreml_provider_open_is_rejected_on_non_apple_builds() { + if cfg!(target_vendor = "apple") { + return; + } + let error = provider_dispatches(&[SpeakerExecutionProvider::CoreMl]).unwrap_err(); + assert_eq!( + error, + SpeakerOnnxError::ProviderUnavailable { provider: "coreml" } + ); + } + + #[test] + fn committed_wespeaker_model_accepts_fixture_features_and_returns_256_floats() { + let fixture = fixture(); + let audio = decode_waveform(&fixture); + let features = + compute_wespeaker_filterbank_cmn(&audio, WESPEAKER_SAMPLE_RATE_HZ).expect("features"); + let model_path = repo_root().join( + "packages/solstone-journal-models/solstone_journal_models/assets/wespeaker-resnet34-256.onnx", + ); + let mut embedder = WespeakerEmbedder::open(&model_path, &[SpeakerExecutionProvider::Cpu]) + .expect("embedder"); + + let embedding = embedder.embed(&features).expect("embedding"); + + assert_eq!(embedding.values().len(), WESPEAKER_EMBEDDING_SIZE); + assert!(embedding.values().iter().all(|value| value.is_finite())); + } + + #[test] + fn session_builder_has_single_production_site() { + let source = include_str!("lib.rs"); + let needle = concat!("Session", "::builder()?"); + assert_eq!(source.matches(needle).count(), 1); + } + + fn repo_root() -> std::path::PathBuf { + std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../..") + } + + fn fixture() -> serde_json::Value { + serde_json::from_str(FIXTURE).expect("fixture JSON") + } + + fn decode_waveform(fixture: &serde_json::Value) -> Vec { + let encoded = fixture["waveform"]["samples_f32_le_base64"] + .as_str() + .expect("waveform base64"); + let bytes = decode_base64(encoded); + bytes + .chunks_exact(4) + .map(|chunk| f32::from_le_bytes(chunk.try_into().expect("f32 bytes"))) + .collect() + } + + fn decode_base64(input: &str) -> Vec { + let mut out = Vec::with_capacity(input.len() / 4 * 3); + let mut quartet = [0_u8; 4]; + let mut len = 0; + for byte in input.bytes().filter(|byte| !byte.is_ascii_whitespace()) { + quartet[len] = match byte { + b'A'..=b'Z' => byte - b'A', + b'a'..=b'z' => byte - b'a' + 26, + b'0'..=b'9' => byte - b'0' + 52, + b'+' => 62, + b'/' => 63, + b'=' => 64, + _ => panic!("invalid base64 byte: {byte}"), + }; + len += 1; + if len == 4 { + out.push((quartet[0] << 2) | (quartet[1] >> 4)); + if quartet[2] != 64 { + out.push((quartet[1] << 4) | (quartet[2] >> 2)); + } + if quartet[3] != 64 { + out.push((quartet[2] << 6) | quartet[3]); + } + len = 0; + } + } + assert_eq!(len, 0); + out + } +} diff --git a/core/crates/solstone-core-speakers/Cargo.toml b/core/crates/solstone-core-speakers/Cargo.toml new file mode 100644 index 000000000..92c2ab66d --- /dev/null +++ b/core/crates/solstone-core-speakers/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "solstone-core-speakers" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +publish = false + +[lints] +workspace = true + +[dev-dependencies] +serde_json.workspace = true diff --git a/core/crates/solstone-core-speakers/src/lib.rs b/core/crates/solstone-core-speakers/src/lib.rs new file mode 100644 index 000000000..adbc84db1 --- /dev/null +++ b/core/crates/solstone-core-speakers/src/lib.rs @@ -0,0 +1,593 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright (c) 2026 sol pbc + +use std::error::Error; +use std::f32::consts::PI; +use std::fmt; + +pub const WESPEAKER_SAMPLE_RATE_HZ: u32 = 16_000; +pub const WESPEAKER_MEL_BINS: usize = 80; +pub const WESPEAKER_FRAME_LENGTH_SAMPLES: usize = 400; +pub const WESPEAKER_FRAME_SHIFT_SAMPLES: usize = 160; +pub const WESPEAKER_FFT_SIZE: usize = 512; +pub const WESPEAKER_EMBEDDING_SIZE: usize = 256; + +const LOW_FREQ_HZ: f32 = 20.0; +const HIGH_FREQ_HZ: f32 = WESPEAKER_SAMPLE_RATE_HZ as f32 / 2.0; +const PREEMPH_COEFF: f32 = 0.97; +const AUDIO_SCALE: f32 = 32768.0; +const ROW_NORM_EPSILON: f32 = 1e-9; + +#[derive(Debug, Clone, PartialEq)] +pub struct FeatureMatrix { + frames: usize, + bins: usize, + data: Vec, +} + +impl FeatureMatrix { + pub fn from_row_major( + frames: usize, + bins: usize, + data: Vec, + ) -> Result { + let expected = frames + .checked_mul(bins) + .ok_or(SpeakerFeatureError::ShapeOverflow { frames, bins })?; + if data.len() != expected { + return Err(SpeakerFeatureError::ShapeMismatch { + frames, + bins, + len: data.len(), + }); + } + Ok(Self { frames, bins, data }) + } + + pub fn frames(&self) -> usize { + self.frames + } + + pub fn bins(&self) -> usize { + self.bins + } + + pub fn data(&self) -> &[f32] { + &self.data + } + + pub fn row(&self, index: usize) -> Option<&[f32]> { + if index >= self.frames { + return None; + } + let start = index * self.bins; + Some(&self.data[start..start + self.bins]) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SpeakerFeatureError { + UnsupportedSampleRate { + expected: u32, + actual: u32, + }, + NonFiniteAudioSample { + index: usize, + }, + ShapeMismatch { + frames: usize, + bins: usize, + len: usize, + }, + ShapeOverflow { + frames: usize, + bins: usize, + }, +} + +impl fmt::Display for SpeakerFeatureError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::UnsupportedSampleRate { expected, actual } => write!( + formatter, + "unsupported sample rate: expected {expected}, got {actual}" + ), + Self::NonFiniteAudioSample { index } => { + write!(formatter, "audio sample at index {index} is not finite") + } + Self::ShapeMismatch { frames, bins, len } => write!( + formatter, + "row-major matrix length mismatch: frames={frames} bins={bins} len={len}" + ), + Self::ShapeOverflow { frames, bins } => { + write!( + formatter, + "row-major matrix shape overflow: frames={frames} bins={bins}" + ) + } + } + } +} + +impl Error for SpeakerFeatureError {} + +pub fn compute_wespeaker_filterbank_cmn( + audio: &[f32], + sample_rate_hz: u32, +) -> Result { + if sample_rate_hz != WESPEAKER_SAMPLE_RATE_HZ { + return Err(SpeakerFeatureError::UnsupportedSampleRate { + expected: WESPEAKER_SAMPLE_RATE_HZ, + actual: sample_rate_hz, + }); + } + if let Some((index, _sample)) = audio + .iter() + .enumerate() + .find(|(_index, sample)| !sample.is_finite()) + { + return Err(SpeakerFeatureError::NonFiniteAudioSample { index }); + } + + let frames = num_snipped_frames(audio.len()); + let mut data = vec![0.0; frames * WESPEAKER_MEL_BINS]; + if frames == 0 { + return FeatureMatrix::from_row_major(0, WESPEAKER_MEL_BINS, data); + } + + let window = povey_window(); + let mel_bank = MelBank::production(); + for frame in 0..frames { + let mut padded = padded_frame(audio, frame, &window); + let power = power_spectrum_512(&mut padded); + let output_start = frame * WESPEAKER_MEL_BINS; + mel_bank.compute_log_energies( + &power, + &mut data[output_start..output_start + WESPEAKER_MEL_BINS], + ); + } + subtract_column_mean(frames, WESPEAKER_MEL_BINS, &mut data); + FeatureMatrix::from_row_major(frames, WESPEAKER_MEL_BINS, data) +} + +pub fn row_l2_normalize(features: &FeatureMatrix) -> FeatureMatrix { + let mut data = features.data.clone(); + for row in data.chunks_mut(features.bins) { + let norm = row.iter().map(|value| value * value).sum::().sqrt(); + let denom = if norm > ROW_NORM_EPSILON { norm } else { 1.0 }; + for value in row { + *value /= denom; + } + } + FeatureMatrix { + frames: features.frames, + bins: features.bins, + data, + } +} + +fn num_snipped_frames(samples: usize) -> usize { + if samples < WESPEAKER_FRAME_LENGTH_SAMPLES { + 0 + } else { + 1 + (samples - WESPEAKER_FRAME_LENGTH_SAMPLES) / WESPEAKER_FRAME_SHIFT_SAMPLES + } +} + +fn povey_window() -> [f32; WESPEAKER_FRAME_LENGTH_SAMPLES] { + let mut window = [0.0; WESPEAKER_FRAME_LENGTH_SAMPLES]; + let scale = 2.0_f64 * std::f64::consts::PI / (WESPEAKER_FRAME_LENGTH_SAMPLES as f64 - 1.0); + for (index, value) in window.iter_mut().enumerate() { + *value = (0.5_f64 - 0.5_f64 * (scale * index as f64).cos()).powf(0.85) as f32; + } + window +} + +fn padded_frame( + audio: &[f32], + frame: usize, + window: &[f32; WESPEAKER_FRAME_LENGTH_SAMPLES], +) -> [f32; WESPEAKER_FFT_SIZE] { + let mut padded = [0.0; WESPEAKER_FFT_SIZE]; + let start = frame * WESPEAKER_FRAME_SHIFT_SAMPLES; + for index in 0..WESPEAKER_FRAME_LENGTH_SAMPLES { + padded[index] = audio[start + index] * AUDIO_SCALE; + } + remove_dc_offset(&mut padded[..WESPEAKER_FRAME_LENGTH_SAMPLES]); + preemphasize(&mut padded[..WESPEAKER_FRAME_LENGTH_SAMPLES]); + apply_window(&mut padded[..WESPEAKER_FRAME_LENGTH_SAMPLES], window); + padded +} + +fn remove_dc_offset(frame: &mut [f32]) { + let mean = frame.iter().sum::() / frame.len() as f32; + for sample in frame { + *sample -= mean; + } +} + +fn preemphasize(frame: &mut [f32]) { + for index in (1..frame.len()).rev() { + frame[index] -= PREEMPH_COEFF * frame[index - 1]; + } + frame[0] -= PREEMPH_COEFF * frame[0]; +} + +fn apply_window(frame: &mut [f32], window: &[f32; WESPEAKER_FRAME_LENGTH_SAMPLES]) { + for (sample, weight) in frame.iter_mut().zip(window) { + *sample *= *weight; + } +} + +fn power_spectrum_512(input: &mut [f32; WESPEAKER_FFT_SIZE]) -> [f32; WESPEAKER_FFT_SIZE / 2 + 1] { + let mut real = *input; + let mut imag = [0.0; WESPEAKER_FFT_SIZE]; + fft_512(&mut real, &mut imag); + let mut power = [0.0; WESPEAKER_FFT_SIZE / 2 + 1]; + for index in 0..power.len() { + power[index] = real[index] * real[index] + imag[index] * imag[index]; + } + power +} + +fn fft_512(real: &mut [f32; WESPEAKER_FFT_SIZE], imag: &mut [f32; WESPEAKER_FFT_SIZE]) { + let mut j = 0; + for i in 1..WESPEAKER_FFT_SIZE { + let mut bit = WESPEAKER_FFT_SIZE >> 1; + while j & bit != 0 { + j ^= bit; + bit >>= 1; + } + j ^= bit; + if i < j { + real.swap(i, j); + imag.swap(i, j); + } + } + + let mut len = 2; + while len <= WESPEAKER_FFT_SIZE { + let angle = -2.0 * PI / len as f32; + let wlen_real = angle.cos(); + let wlen_imag = angle.sin(); + for start in (0..WESPEAKER_FFT_SIZE).step_by(len) { + let mut w_real = 1.0; + let mut w_imag = 0.0; + for offset in 0..(len / 2) { + let even = start + offset; + let odd = even + len / 2; + let odd_real = real[odd] * w_real - imag[odd] * w_imag; + let odd_imag = real[odd] * w_imag + imag[odd] * w_real; + real[odd] = real[even] - odd_real; + imag[odd] = imag[even] - odd_imag; + real[even] += odd_real; + imag[even] += odd_imag; + let next_real = w_real * wlen_real - w_imag * wlen_imag; + w_imag = w_real * wlen_imag + w_imag * wlen_real; + w_real = next_real; + } + } + len *= 2; + } +} + +fn subtract_column_mean(frames: usize, bins: usize, data: &mut [f32]) { + for bin in 0..bins { + let mut sum = 0.0; + for frame in 0..frames { + sum += data[frame * bins + bin]; + } + let mean = sum / frames as f32; + for frame in 0..frames { + data[frame * bins + bin] -= mean; + } + } +} + +#[derive(Debug, Clone)] +struct MelBin { + offset: usize, + weights: Vec, +} + +#[derive(Debug, Clone)] +struct MelBank { + bins: Vec, +} + +impl MelBank { + fn production() -> Self { + let mel_low = mel_scale(LOW_FREQ_HZ); + let mel_high = mel_scale(HIGH_FREQ_HZ); + let delta = (mel_high - mel_low) / (WESPEAKER_MEL_BINS as f32 + 1.0); + let fft_bin_width = WESPEAKER_SAMPLE_RATE_HZ as f32 / WESPEAKER_FFT_SIZE as f32; + let mut bins = Vec::with_capacity(WESPEAKER_MEL_BINS); + for bin in 0..WESPEAKER_MEL_BINS { + let left = mel_low + bin as f32 * delta; + let center = mel_low + (bin as f32 + 1.0) * delta; + let right = mel_low + (bin as f32 + 2.0) * delta; + let mut dense = [0.0; WESPEAKER_FFT_SIZE / 2]; + let mut first = None; + let mut last = 0; + for (fft_bin, weight) in dense.iter_mut().enumerate() { + let freq = fft_bin_width * fft_bin as f32; + let mel = mel_scale(freq); + if mel > left && mel < right { + *weight = if mel <= center { + (mel - left) / (center - left) + } else { + (right - mel) / (right - center) + }; + first.get_or_insert(fft_bin); + last = fft_bin; + } + } + let first = first.expect("production mel bin should have weights"); + bins.push(MelBin { + offset: first, + weights: dense[first..=last].to_vec(), + }); + } + MelBank { bins } + } + + fn compute_log_energies(&self, power: &[f32; WESPEAKER_FFT_SIZE / 2 + 1], output: &mut [f32]) { + for (bin, out) in self.bins.iter().zip(output.iter_mut()) { + let mut energy = 0.0; + for (index, weight) in bin.weights.iter().enumerate() { + energy += weight * power[bin.offset + index]; + } + *out = energy.max(f32::EPSILON).ln(); + } + } + + #[cfg(test)] + fn weight_columns(&self) -> usize { + WESPEAKER_FFT_SIZE / 2 + } +} + +fn mel_scale(freq: f32) -> f32 { + 1127.0 * (1.0 + freq / 700.0).ln() +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::Value; + + const FIXTURE: &str = include_str!("../../../fixtures/speaker_filterbank.json"); + const TOLERANCE: f32 = 1e-2; + + #[test] + fn nyquist_power_bin_is_computed_but_not_consumed_by_mel_bank() { + let bank = MelBank::production(); + assert_eq!(bank.weight_columns(), WESPEAKER_FFT_SIZE / 2); + + let mut power = [0.0; WESPEAKER_FFT_SIZE / 2 + 1]; + let mut baseline = [0.0; WESPEAKER_MEL_BINS]; + bank.compute_log_energies(&power, &mut baseline); + + power[WESPEAKER_FFT_SIZE / 2] = 1.0e12; + let mut perturbed = [0.0; WESPEAKER_MEL_BINS]; + bank.compute_log_energies(&power, &mut perturbed); + + assert_eq!(baseline, perturbed); + } + + #[test] + fn povey_window_spans_400_samples_and_padded_tail_stays_zero() { + let window = povey_window(); + assert_eq!(window.len(), WESPEAKER_FRAME_LENGTH_SAMPLES); + + let audio = vec![0.25; WESPEAKER_FRAME_LENGTH_SAMPLES]; + let padded = padded_frame(&audio, 0, &window); + assert!( + padded[WESPEAKER_FRAME_LENGTH_SAMPLES..] + .iter() + .all(|sample| *sample == 0.0) + ); + } + + #[test] + fn filterbank_cmn_and_row_l2_match_fixture_regions_separately() { + let fixture = fixture(); + let audio = decode_waveform(&fixture); + let cmn = + compute_wespeaker_filterbank_cmn(&audio, WESPEAKER_SAMPLE_RATE_HZ).expect("features"); + let row_l2 = row_l2_normalize(&cmn); + let expected_cmn = fixture_matrix(&fixture, "filterbank_cmn"); + let expected_row_l2 = fixture_matrix(&fixture, "row_l2_normalized"); + let near_rows = fixture_range(&fixture, "near_silent_rows"); + let broad_rows = fixture_range(&fixture, "broadband_rows"); + eprintln!( + "filterbank_cmn_max_abs_diff={}", + max_abs_diff(cmn.data(), &expected_cmn) + ); + + assert_matrix_within("filterbank_cmn", cmn.data(), &expected_cmn, TOLERANCE); + assert_matrix_within( + "row_l2_normalized", + row_l2.data(), + &expected_row_l2, + TOLERANCE, + ); + assert_region_within( + "filterbank_cmn near_silent", + cmn.data(), + &expected_cmn, + cmn.bins(), + near_rows.clone(), + TOLERANCE, + ); + assert_region_within( + "filterbank_cmn broadband", + cmn.data(), + &expected_cmn, + cmn.bins(), + broad_rows.clone(), + TOLERANCE, + ); + assert_region_within( + "row_l2_normalized near_silent", + row_l2.data(), + &expected_row_l2, + row_l2.bins(), + near_rows, + TOLERANCE, + ); + assert_region_within( + "row_l2_normalized broadband", + row_l2.data(), + &expected_row_l2, + row_l2.bins(), + broad_rows, + TOLERANCE, + ); + } + + #[test] + fn fixture_comparison_fails_when_value_exceeds_tolerance() { + let fixture = fixture(); + let audio = decode_waveform(&fixture); + let cmn = + compute_wespeaker_filterbank_cmn(&audio, WESPEAKER_SAMPLE_RATE_HZ).expect("features"); + let mut expected = fixture_matrix(&fixture, "filterbank_cmn"); + expected[0] = cmn.data()[0] + TOLERANCE * 2.0; + + let result = matrix_comparison_error("filterbank_cmn", cmn.data(), &expected, TOLERANCE); + assert!(result.is_some()); + } + + #[test] + fn mel_bank_construction_has_single_production_site() { + let source = include_str!("lib.rs"); + let needle = concat!("MelBank", " { bins"); + assert_eq!(source.matches(needle).count(), 1); + } + + fn fixture() -> Value { + serde_json::from_str(FIXTURE).expect("fixture JSON") + } + + fn decode_waveform(fixture: &Value) -> Vec { + let encoded = fixture["waveform"]["samples_f32_le_base64"] + .as_str() + .expect("waveform base64"); + let bytes = decode_base64(encoded); + bytes + .chunks_exact(4) + .map(|chunk| f32::from_le_bytes(chunk.try_into().expect("f32 bytes"))) + .collect() + } + + fn fixture_matrix(fixture: &Value, name: &str) -> Vec { + let rows = fixture["matrices"][name]["rows"] + .as_array() + .expect("matrix rows"); + rows.iter() + .flat_map(|row| { + row.as_str() + .expect("decimal row") + .split_whitespace() + .map(|value| value.parse::().expect("decimal value")) + }) + .collect() + } + + fn fixture_range(fixture: &Value, name: &str) -> std::ops::Range { + let values = fixture["waveform"][name].as_array().expect("row range"); + let start = values[0].as_u64().expect("range start") as usize; + let end = values[1].as_u64().expect("range end") as usize; + start..end + } + + fn assert_matrix_within(name: &str, actual: &[f32], expected: &[f32], tolerance: f32) { + if let Some(error) = matrix_comparison_error(name, actual, expected, tolerance) { + panic!("{error}"); + } + } + + fn assert_region_within( + name: &str, + actual: &[f32], + expected: &[f32], + bins: usize, + rows: std::ops::Range, + tolerance: f32, + ) { + let start = rows.start * bins; + let end = rows.end * bins; + assert_matrix_within(name, &actual[start..end], &expected[start..end], tolerance); + } + + fn matrix_comparison_error( + name: &str, + actual: &[f32], + expected: &[f32], + tolerance: f32, + ) -> Option { + if actual.len() != expected.len() { + return Some(format!( + "{name} length mismatch: actual={} expected={}", + actual.len(), + expected.len() + )); + } + let (max_abs, max_index) = max_abs_diff_with_index(actual, expected); + if max_abs > tolerance { + Some(format!( + "{name} max_abs_diff={max_abs} at flat_index={max_index} tolerance={tolerance}" + )) + } else { + None + } + } + + fn max_abs_diff(actual: &[f32], expected: &[f32]) -> f32 { + max_abs_diff_with_index(actual, expected).0 + } + + fn max_abs_diff_with_index(actual: &[f32], expected: &[f32]) -> (f32, usize) { + let mut max_abs = 0.0_f32; + let mut max_index = 0; + for (index, (left, right)) in actual.iter().zip(expected).enumerate() { + let diff = (left - right).abs(); + if diff > max_abs { + max_abs = diff; + max_index = index; + } + } + (max_abs, max_index) + } + + fn decode_base64(input: &str) -> Vec { + let mut out = Vec::with_capacity(input.len() / 4 * 3); + let mut quartet = [0_u8; 4]; + let mut len = 0; + for byte in input.bytes().filter(|byte| !byte.is_ascii_whitespace()) { + quartet[len] = match byte { + b'A'..=b'Z' => byte - b'A', + b'a'..=b'z' => byte - b'a' + 26, + b'0'..=b'9' => byte - b'0' + 52, + b'+' => 62, + b'/' => 63, + b'=' => 64, + _ => panic!("invalid base64 byte: {byte}"), + }; + len += 1; + if len == 4 { + out.push((quartet[0] << 2) | (quartet[1] >> 4)); + if quartet[2] != 64 { + out.push((quartet[1] << 4) | (quartet[2] >> 2)); + } + if quartet[3] != 64 { + out.push((quartet[2] << 6) | quartet[3]); + } + len = 0; + } + } + assert_eq!(len, 0); + out + } +} diff --git a/docs/PORTING.md b/docs/PORTING.md index 379becb95..e85f559f5 100644 --- a/docs/PORTING.md +++ b/docs/PORTING.md @@ -26,6 +26,10 @@ adapter, not for the indexer logic. The eventual iOS path is to link the system `libsqlite3` that iOS ships instead of bundling SQLite, then return the store crate to the iOS gate. +`solstone-core-speakers` stays in the iOS canary because it is pure Rust DSP. +`solstone-core-speakers-onnx` is excluded: ONNX Runtime is a host native-runtime +adapter, not mobile-ready subsystem logic. + ## Native Dependency Release Proof A Rust conversion that adds or bumps a dependency with C/C++ build steps or @@ -36,6 +40,25 @@ required toolchain, target, and linker behavior in checked-in repository release paths, not in a local shell profile. If a dependency cannot satisfy a supported target, document the blocker and stop the conversion before merging it. +The ONNX Runtime speaker wrapper is a deliberate exception because it is outside +the `solstone-core` shipping closure. No shipping bin depends on +`solstone-core-speakers-onnx`, and `scripts/core_compile_inputs.py` walks only the +`solstone-core` closure. That unreachability is load-bearing, not a convenience: +reachability from the shipping bin is an open architecture question, currently +believed unsatisfiable as-is. The Linux release lanes are +`x86_64-unknown-linux-musl` and `aarch64-unknown-linux-musl` via zig, i.e. static +musl. The prebuilt ONNX Runtime wheels are glibc-only native libraries; the Linux +1.25.0 wheel requires GLIBC symbols up to `GLIBC_2.27` and links `libstdc++` and +`libgcc_s`. A static-musl binary cannot link or `dlopen` that glibc shared +library, so the three-target native dependency proof would fail rather than +merely remain pending. + +The shipping shape for ONNX-bearing subsystems is a separate founder-owned +decision. The likely direction is a separate glibc-tagged helper artifact, +matching how `parakeet-server` and `llama-server` already arrive as separate +native binaries. Do not design toward linking `solstone-core-speakers-onnx` into +the `solstone-core` bin. + | Evidence | Repository command | Class | Notes | |----------|--------------------|-------|-------| | Rust formatting | `make check-rust-fmt` | GNU-host check | Host source-format evidence only. | @@ -44,10 +67,33 @@ target, document the blocker and stop the conversion before merging it. | Rust tests | `make check-rust-test` | GNU-host check | Runs workspace Rust tests on the GNU host. | | Rust dependency policy | `make check-rust-deny` | GNU-host check | Locked, offline bans/licenses/sources policy over the supported cargo-deny graph. | | Rust advisories | `make audit` | GNU-host check | Verifies a signed advisory mirror packet, materializes its bundle locally, then performs a locked offline advisory check without refreshing or mutating the operator inputs. | -| iOS canary | `make check-rust-ios` | iOS cross-target canary | Cross-target drift evidence for eligible library crates; explicitly excludes `solstone-core-indexer-store` because the native SQLite store is not yet in the iOS gate. | +| iOS canary | `make check-rust-ios` | iOS cross-target canary | Cross-target drift evidence for eligible library crates; explicitly excludes `solstone-core-indexer-store` because the native SQLite store is not yet in the iOS gate, and `solstone-core-speakers-onnx` because ONNX Runtime is host-only native linkage. | | Core sdist compile inputs | `make check-core-sdist-compile-inputs` | Packaging-source check | Verifies shipping Rust compile-time inputs are discovered and covered by the normalized `solstone-core` sdist injection set. | | Release candidate rail | `scripts/release.sh --candidate` / `scripts/release.sh --recover ` | Local readiness evidence | DESTRUCTIVE: `--candidate` is fresh construction; before policy or build work it deletes prior raw build/dist outputs and that version's stale payload/evidence. It binds candidate payload, ledger, and per-target install/smoke proofs, then reports canonical local readiness JSON. `--recover` is retained-byte-only, read-only validation; it preserves retained payload, ledger, and proofs and never rebuilds or refreshes. Proofs cover local candidate bytes and native smoke only; publication is temporarily locked out of this rail. | +### Rust ONNX Runtime Provisioning + +`solstone-core-speakers-onnx` links dynamically to the ONNX Runtime C API from +the journal Python environment. It does not download or vendor ONNX Runtime and +does not read paths from inside the crate. Host Rust workspace compile checks +that build the ONNX crate (`make check-rust-msrv`, `make check-rust-clippy`, and +`make check-rust-test`) run Cargo through `scripts/resolve_onnxruntime_capi.py`, +which imports the Python module `onnxruntime` using `.venv/bin/python`; this +works whether the installed distribution is `onnxruntime` or `onnxruntime-gpu`. + +The resolver stages symlinks, never copies, under +`core/target/onnxruntime-link//lib/`. Linux stages +`libonnxruntime.so`, `libonnxruntime.so.1`, and the full-version shared object. +macOS stages `libonnxruntime.dylib` and the full-version dylib from the wheel. +It then executes Cargo with `ORT_LIB_PATH=`, +`ORT_PREFER_DYNAMIC_LINK=true`, and the matching runtime loader path +(`LD_LIBRARY_PATH` on Linux, `DYLD_LIBRARY_PATH` on macOS) prepended. + +Direct `cargo test --manifest-path core/Cargo.toml --workspace` without the +resolver is unsupported for the ONNX crate because Cargo dependency build +scripts run before dependent crate build scripts. A crate `build.rs` cannot +retroactively provide `ORT_LIB_PATH` to `ort-sys`. + ### Signed Advisory Mirror Audit `make audit` requires four operator-provided inputs: `AUDIT_ADVISORY_BUNDLE` diff --git a/scripts/resolve_onnxruntime_capi.py b/scripts/resolve_onnxruntime_capi.py new file mode 100644 index 000000000..5455efad7 --- /dev/null +++ b/scripts/resolve_onnxruntime_capi.py @@ -0,0 +1,272 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright (c) 2026 sol pbc + +"""Prepare ONNX Runtime C API linkage for Rust speaker tests.""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Mapping, Sequence + +ROOT = Path(__file__).resolve().parent.parent +DEFAULT_PYTHON = ROOT / ".venv" / "bin" / "python" +MIN_ONNXRUNTIME_VERSION = (1, 24, 0) +ABSENT_MESSAGE = ( + "ONNX Runtime C API unavailable: .venv/bin/python could not import the " + "onnxruntime module or find its capi library. Run 'make install' from the " + "repo root; it installs the journal runtime dependency, using " + "onnxruntime-gpu on CUDA hosts and onnxruntime otherwise." +) +TOO_OLD_TEMPLATE = ( + "ONNX Runtime C API too old: found {version} at {capi}; " + "solstone-core-speakers-onnx requires C API 24, provided by ONNX Runtime " + ">=1.24.0. Run 'make install' from the repo root to refresh the journal " + "runtime dependency." +) + +Runner = Callable[..., subprocess.CompletedProcess[str]] + + +@dataclass(frozen=True) +class OnnxRuntimeProbe: + version: str + module_file: Path + capi_dir: Path + platform_tag: str + system: str + + +@dataclass(frozen=True) +class PreparedOnnxRuntime: + probe: OnnxRuntimeProbe + link_dir: Path + runtime_library: Path + staged_links: tuple[Path, ...] + + +class OnnxRuntimeResolverError(RuntimeError): + """ONNX Runtime C API resolution failed.""" + + +def _version_tuple(version: str) -> tuple[int, int, int]: + match = re.match(r"^(\d+)\.(\d+)\.(\d+)", version) + if match is None: + raise OnnxRuntimeResolverError( + TOO_OLD_TEMPLATE.format(version=version, capi="") + ) + return tuple(int(part) for part in match.groups()) + + +def _probe_code() -> str: + return ( + "import json, pathlib, platform, sysconfig\n" + "import onnxruntime\n" + "module = pathlib.Path(onnxruntime.__file__).resolve()\n" + "capi = module.parent / 'capi'\n" + "print(json.dumps({\n" + " 'version': onnxruntime.__version__,\n" + " 'module_file': str(module),\n" + " 'capi_dir': str(capi),\n" + " 'platform_tag': sysconfig.get_platform(),\n" + " 'system': platform.system(),\n" + "}, sort_keys=True))\n" + ) + + +def probe_onnxruntime( + python: Path, + *, + runner: Runner = subprocess.run, +) -> OnnxRuntimeProbe: + if not python.exists(): + raise OnnxRuntimeResolverError(ABSENT_MESSAGE) + result = runner( + [str(python), "-c", _probe_code()], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + raise OnnxRuntimeResolverError(ABSENT_MESSAGE) + try: + payload = json.loads(result.stdout) + except json.JSONDecodeError as exc: + raise OnnxRuntimeResolverError(ABSENT_MESSAGE) from exc + return OnnxRuntimeProbe( + version=str(payload["version"]), + module_file=Path(payload["module_file"]), + capi_dir=Path(payload["capi_dir"]), + platform_tag=str(payload["platform_tag"]), + system=str(payload["system"]), + ) + + +def _require_supported_version(probe: OnnxRuntimeProbe) -> None: + version_tuple = _version_tuple(probe.version) + if version_tuple < MIN_ONNXRUNTIME_VERSION: + raise OnnxRuntimeResolverError( + TOO_OLD_TEMPLATE.format(version=probe.version, capi=probe.capi_dir) + ) + + +def _find_runtime_library(probe: OnnxRuntimeProbe) -> Path: + capi = probe.capi_dir + if not capi.is_dir(): + raise OnnxRuntimeResolverError(ABSENT_MESSAGE) + if probe.system == "Linux": + exact = capi / f"libonnxruntime.so.{probe.version}" + candidates = ( + [exact] if exact.exists() else sorted(capi.glob("libonnxruntime.so.*")) + ) + elif probe.system == "Darwin": + exact = capi / f"libonnxruntime.{probe.version}.dylib" + candidates = ( + [exact] if exact.exists() else sorted(capi.glob("libonnxruntime.*.dylib")) + ) + else: + raise OnnxRuntimeResolverError( + f"ONNX Runtime C API unsupported platform for Rust tests: {probe.system}" + ) + if not candidates: + raise OnnxRuntimeResolverError(ABSENT_MESSAGE) + return candidates[0].resolve() + + +def _platform_component(platform_tag: str) -> str: + cleaned = re.sub(r"[^A-Za-z0-9_.-]+", "-", platform_tag).strip("-") + return cleaned or "unknown" + + +def _replace_symlink(link: Path, target: Path) -> None: + if link.is_symlink() or link.exists(): + if not link.is_symlink(): + raise OnnxRuntimeResolverError( + f"ONNX Runtime link target exists and is not a symlink: {link}" + ) + link.unlink() + link.symlink_to(target) + + +def prepare_onnxruntime( + root: Path, + python: Path, + *, + runner: Runner = subprocess.run, +) -> PreparedOnnxRuntime: + probe = probe_onnxruntime(python, runner=runner) + _require_supported_version(probe) + runtime_library = _find_runtime_library(probe) + link_dir = ( + root + / "core" + / "target" + / "onnxruntime-link" + / _platform_component(probe.platform_tag) + / "lib" + ) + link_dir.mkdir(parents=True, exist_ok=True) + if probe.system == "Linux": + link_names = ( + "libonnxruntime.so", + "libonnxruntime.so.1", + runtime_library.name, + ) + elif probe.system == "Darwin": + link_names = ("libonnxruntime.dylib", runtime_library.name) + else: + raise OnnxRuntimeResolverError( + f"ONNX Runtime C API unsupported platform for Rust tests: {probe.system}" + ) + links = tuple(link_dir / name for name in link_names) + for link in links: + _replace_symlink(link, runtime_library) + return PreparedOnnxRuntime( + probe=probe, + link_dir=link_dir, + runtime_library=runtime_library, + staged_links=links, + ) + + +def env_for_cargo( + prepared: PreparedOnnxRuntime, + base_env: Mapping[str, str] | None = None, +) -> dict[str, str]: + env = dict(os.environ if base_env is None else base_env) + link_dir = str(prepared.link_dir) + env["ORT_LIB_PATH"] = link_dir + env["ORT_PREFER_DYNAMIC_LINK"] = "true" + if prepared.probe.system == "Linux": + key = "LD_LIBRARY_PATH" + elif prepared.probe.system == "Darwin": + key = "DYLD_LIBRARY_PATH" + else: + key = "" + if key: + existing = env.get(key) + env[key] = link_dir if not existing else f"{link_dir}{os.pathsep}{existing}" + return env + + +def _json_payload(prepared: PreparedOnnxRuntime) -> str: + return json.dumps( + { + "capi_dir": str(prepared.probe.capi_dir), + "link_dir": str(prepared.link_dir), + "runtime_library": str(prepared.runtime_library), + "staged_links": [str(path) for path in prepared.staged_links], + "system": prepared.probe.system, + "version": prepared.probe.version, + }, + sort_keys=True, + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Prepare ONNX Runtime C API linkage for Rust speaker tests." + ) + parser.add_argument("--root", type=Path, default=ROOT) + parser.add_argument("--python", type=Path, default=DEFAULT_PYTHON) + parser.add_argument("--json", action="store_true", help="Print resolver details.") + parser.add_argument("--exec", action="store_true") + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + raw_argv = list(sys.argv[1:] if argv is None else argv) + exec_command: list[str] | None = None + if "--exec" in raw_argv: + index = raw_argv.index("--exec") + exec_command = raw_argv[index + 1 :] + raw_argv = raw_argv[: index + 1] + if exec_command and exec_command[0] == "--": + exec_command = exec_command[1:] + args = build_parser().parse_args(raw_argv) + try: + prepared = prepare_onnxruntime(args.root.resolve(), args.python) + except OnnxRuntimeResolverError as exc: + print(str(exc), file=sys.stderr) + return 2 + if args.json: + print(_json_payload(prepared)) + if args.exec: + command = exec_command or [] + if not command: + print("--exec requires a command", file=sys.stderr) + return 2 + result = subprocess.run(command, cwd=args.root, env=env_for_cargo(prepared)) + return result.returncode + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_onnxruntime_capi_resolver.py b/tests/test_onnxruntime_capi_resolver.py new file mode 100644 index 000000000..6e0e0eed0 --- /dev/null +++ b/tests/test_onnxruntime_capi_resolver.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright (c) 2026 sol pbc + +from __future__ import annotations + +import json +import os +import subprocess +from pathlib import Path + +import pytest + +import scripts.resolve_onnxruntime_capi as resolver + + +def _python(tmp_path: Path) -> Path: + path = tmp_path / "python" + path.write_text("#!/bin/sh\n", encoding="utf-8") + path.chmod(0o755) + return path + + +def _runner(payload: dict[str, str]): + def run(argv, **kwargs) -> subprocess.CompletedProcess[str]: + return subprocess.CompletedProcess(argv, 0, json.dumps(payload), "") + + return run + + +def _payload( + capi: Path, + *, + version: str = "1.25.0", + platform_tag: str = "linux-x86_64", + system: str = "Linux", +) -> dict[str, str]: + return { + "version": version, + "module_file": str(capi.parent / "__init__.py"), + "capi_dir": str(capi), + "platform_tag": platform_tag, + "system": system, + } + + +def test_missing_venv_python_reports_actionable_cuda_safe_message( + tmp_path: Path, +) -> None: + with pytest.raises(resolver.OnnxRuntimeResolverError) as excinfo: + resolver.prepare_onnxruntime(tmp_path, tmp_path / "missing-python") + + assert str(excinfo.value) == resolver.ABSENT_MESSAGE + + +def test_missing_capi_reports_actionable_cuda_safe_message(tmp_path: Path) -> None: + python = _python(tmp_path) + capi = tmp_path / "venv" / "onnxruntime" / "capi" + + with pytest.raises(resolver.OnnxRuntimeResolverError) as excinfo: + resolver.prepare_onnxruntime( + tmp_path, + python, + runner=_runner(_payload(capi)), + ) + + assert str(excinfo.value) == resolver.ABSENT_MESSAGE + + +def test_too_old_runtime_reports_required_api_level(tmp_path: Path) -> None: + python = _python(tmp_path) + capi = tmp_path / "venv" / "onnxruntime" / "capi" + capi.mkdir(parents=True) + (capi / "libonnxruntime.so.1.23.0").write_bytes(b"runtime") + + with pytest.raises(resolver.OnnxRuntimeResolverError) as excinfo: + resolver.prepare_onnxruntime( + tmp_path, + python, + runner=_runner(_payload(capi, version="1.23.0")), + ) + + assert str(excinfo.value) == resolver.TOO_OLD_TEMPLATE.format( + version="1.23.0", + capi=capi, + ) + + +def test_linux_stages_linker_and_soname_symlinks(tmp_path: Path) -> None: + python = _python(tmp_path) + capi = tmp_path / "venv" / "onnxruntime" / "capi" + capi.mkdir(parents=True) + runtime = capi / "libonnxruntime.so.1.25.0" + runtime.write_bytes(b"runtime") + + prepared = resolver.prepare_onnxruntime( + tmp_path, + python, + runner=_runner(_payload(capi)), + ) + + link_dir = ( + tmp_path / "core" / "target" / "onnxruntime-link" / "linux-x86_64" / "lib" + ) + assert prepared.link_dir == link_dir + assert {path.name for path in prepared.staged_links} == { + "libonnxruntime.so", + "libonnxruntime.so.1", + "libonnxruntime.so.1.25.0", + } + for link in prepared.staged_links: + assert link.is_symlink() + assert link.resolve() == runtime + + env = resolver.env_for_cargo(prepared, {"LD_LIBRARY_PATH": "/existing"}) + assert env["ORT_LIB_PATH"] == str(link_dir) + assert env["ORT_PREFER_DYNAMIC_LINK"] == "true" + assert env["LD_LIBRARY_PATH"] == f"{link_dir}{os.pathsep}/existing" + + +def test_macos_stages_linker_and_install_name_symlinks(tmp_path: Path) -> None: + python = _python(tmp_path) + capi = tmp_path / "venv" / "onnxruntime" / "capi" + capi.mkdir(parents=True) + runtime = capi / "libonnxruntime.1.25.0.dylib" + runtime.write_bytes(b"runtime") + + prepared = resolver.prepare_onnxruntime( + tmp_path, + python, + runner=_runner( + _payload( + capi, + platform_tag="macosx-14.0-arm64", + system="Darwin", + ) + ), + ) + + link_dir = ( + tmp_path / "core" / "target" / "onnxruntime-link" / "macosx-14.0-arm64" / "lib" + ) + assert prepared.link_dir == link_dir + assert {path.name for path in prepared.staged_links} == { + "libonnxruntime.dylib", + "libonnxruntime.1.25.0.dylib", + } + for link in prepared.staged_links: + assert link.is_symlink() + assert link.resolve() == runtime + + env = resolver.env_for_cargo(prepared, {"DYLD_LIBRARY_PATH": "/existing"}) + assert env["ORT_LIB_PATH"] == str(link_dir) + assert env["ORT_PREFER_DYNAMIC_LINK"] == "true" + assert env["DYLD_LIBRARY_PATH"] == f"{link_dir}{os.pathsep}/existing"