diff --git a/core/Cargo.lock b/core/Cargo.lock index 0ffc19d6b..8120b7aff 100644 --- a/core/Cargo.lock +++ b/core/Cargo.lock @@ -332,6 +332,16 @@ dependencies = [ "hashbrown 0.17.1", ] +[[package]] +name = "hdbscan" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "466aa8bfaa22c81efcf67c84e6ad046de13a05c6cc9d41a5094361bff464247a" +dependencies = [ + "kdtree", + "num-traits", +] + [[package]] name = "hkdf" version = "0.12.4" @@ -417,6 +427,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "kdtree" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f0a0e9f770b65bac9aad00f97a67ab5c5319effed07f6da385da3c2115e47ba" +dependencies = [ + "num-traits", + "thiserror 1.0.69", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -1048,6 +1068,7 @@ dependencies = [ name = "solstone-core-speakers" version = "1.0.17" dependencies = [ + "hdbscan", "serde_json", ] diff --git a/core/Cargo.toml b/core/Cargo.toml index 41f377e5e..c55cc438a 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -39,6 +39,7 @@ 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" +hdbscan = { version = "0.12.0", default-features = false, features = ["serial"] } md5 = "0.8.1" nix = { version = "0.30.1", default-features = false, features = ["fs"] } ort = { version = "=2.0.0-rc.12", default-features = false } diff --git a/core/crates/solstone-core-speakers-analyze/src/lib.rs b/core/crates/solstone-core-speakers-analyze/src/lib.rs index aca0f9fa7..c97b58fc8 100644 --- a/core/crates/solstone-core-speakers-analyze/src/lib.rs +++ b/core/crates/solstone-core-speakers-analyze/src/lib.rs @@ -24,7 +24,10 @@ use std::path::Path; use serde_json::{Map, Value, json}; use solstone_core_speakers::diarization::{ DiarizationError, FrameLogProbs, MIN_INTERVAL_S, SentenceTiming, SpeakerInterval, - assign_sentences, cluster_embeddings, find_intervals, + assign_sentences, cluster_embeddings as cluster_diarization_embeddings, find_intervals, +}; +use solstone_core_speakers::discovery::{ + DiscoveryClusteringError, cluster_embeddings as cluster_discovery_embeddings, }; use solstone_core_speakers::{ PYANNOTE_CLASS_COUNT, PYANNOTE_DIARIZE_STRIDE_S, SpeakerEvidence, SpeakerFeatureError, @@ -40,15 +43,21 @@ use solstone_core_speakers_onnx::{ pub const REQUEST_SCHEMA: &str = "solstone-speaker-analyze-request-v1"; pub const RESPONSE_SCHEMA: &str = "solstone-speaker-analyze-response-v1"; pub const ERROR_SCHEMA: &str = "solstone-speaker-analyze-error-v1"; -pub const USAGE: &str = "Usage: solstone-core-speakers-analyze < request.json > response.json"; +pub const DISCOVERY_CLUSTER_REQUEST_SCHEMA: &str = "solstone-speaker-discovery-cluster-request-v1"; +pub const DISCOVERY_CLUSTER_RESPONSE_SCHEMA: &str = + "solstone-speaker-discovery-cluster-response-v1"; +pub const USAGE: &str = "Usage: solstone-core-speakers-analyze < request.json > response.json\n solstone-core-speakers-analyze discovery-cluster < request.json > response.json"; const PAYLOAD_FORMAT: &str = "raw-f32le-row-major-v1"; const DTYPE_F32LE: &str = "float32-le"; const ENCODER: &str = "wespeaker-resnet34-256"; +const DISCOVERY_CLUSTER_COMMAND: &str = "discovery-cluster"; +const DISCOVERY_CLUSTER_ALGORITHM: &str = "hdbscan-eom-euclidean-f64-prim-mst"; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Command { Run, + DiscoveryCluster, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -69,29 +78,77 @@ impl fmt::Display for UsageError { impl Error for UsageError {} pub fn evaluate_args(args: &[OsString]) -> Result { - if let Some(argument) = args.first() { - return Err(UsageError::UnexpectedArgument { + match args { + [] => Ok(Command::Run), + [argument] if argument == DISCOVERY_CLUSTER_COMMAND => Ok(Command::DiscoveryCluster), + [argument, ..] => Err(UsageError::UnexpectedArgument { argument: argument.to_string_lossy().into_owned(), - }); + }), } - Ok(Command::Run) } #[derive(Debug, Clone, PartialEq, Eq)] pub enum AnalyzeError { - MalformedRequest { detail: String }, - UnknownSchema { schema: String }, - AudioUnreadable { path: String, detail: String }, - AudioInvalid { path: String, detail: String }, - UnsupportedSampleRate { expected: u32, actual: u32 }, - AudioNonFinite { path: String, index: usize }, - ModelUnreadable { field: &'static str, path: String }, - ModelInvalid { field: &'static str, detail: String }, - ModelIoMismatch { field: &'static str, detail: String }, - ProviderUnavailable { detail: String }, - OnnxRuntime { detail: String }, - OutputUnwritable { path: String, detail: String }, - Internal { detail: String }, + MalformedRequest { + detail: String, + }, + UnknownSchema { + schema: String, + }, + AudioUnreadable { + path: String, + detail: String, + }, + AudioInvalid { + path: String, + detail: String, + }, + UnsupportedSampleRate { + expected: u32, + actual: u32, + }, + AudioNonFinite { + path: String, + index: usize, + }, + PayloadUnreadable { + path: String, + detail: String, + }, + PayloadInvalid { + path: String, + detail: String, + }, + PayloadNonFinite { + path: String, + row: usize, + col: usize, + }, + ModelUnreadable { + field: &'static str, + path: String, + }, + ModelInvalid { + field: &'static str, + detail: String, + }, + ModelIoMismatch { + field: &'static str, + detail: String, + }, + ProviderUnavailable { + detail: String, + }, + OnnxRuntime { + detail: String, + }, + OutputUnwritable { + path: String, + detail: String, + }, + Internal { + detail: String, + }, } impl AnalyzeError { @@ -103,6 +160,9 @@ impl AnalyzeError { Self::AudioInvalid { .. } => "audio-invalid", Self::UnsupportedSampleRate { .. } => "unsupported-sample-rate", Self::AudioNonFinite { .. } => "audio-non-finite", + Self::PayloadUnreadable { .. } => "payload-unreadable", + Self::PayloadInvalid { .. } => "payload-invalid", + Self::PayloadNonFinite { .. } => "payload-non-finite", Self::ModelUnreadable { .. } => "model-unreadable", Self::ModelInvalid { .. } => "model-invalid", Self::ModelIoMismatch { .. } => "model-io-mismatch", @@ -120,6 +180,9 @@ impl AnalyzeError { | Self::AudioInvalid { .. } | Self::UnsupportedSampleRate { .. } | Self::AudioNonFinite { .. } + | Self::PayloadUnreadable { .. } + | Self::PayloadInvalid { .. } + | Self::PayloadNonFinite { .. } | Self::ModelUnreadable { .. } | Self::ModelInvalid { .. } | Self::ModelIoMismatch { .. } => 69, @@ -148,6 +211,15 @@ impl AnalyzeError { Self::AudioNonFinite { path, index } => { format!("audio path {path:?} contains non-finite sample at index {index}") } + Self::PayloadUnreadable { path, detail } => { + format!("payload path {path:?} is unreadable: {detail}") + } + Self::PayloadInvalid { path, detail } => { + format!("payload path {path:?} is not raw little-endian f32 row-major: {detail}") + } + Self::PayloadNonFinite { path, row, col } => { + format!("payload path {path:?} contains non-finite value at row={row} col={col}") + } Self::ModelUnreadable { field, path } => { format!( "{field} is missing or unreadable at {path:?}; provide a readable ONNX model path" @@ -197,6 +269,17 @@ pub fn error_line_for_analyze_error(error: &AnalyzeError) -> String { } pub fn run_request(input: &str) -> Result { + run_command_request(Command::Run, input) +} + +pub fn run_command_request(command: Command, input: &str) -> Result { + match command { + Command::Run => run_analyze_request(input), + Command::DiscoveryCluster => run_discovery_cluster_request(input), + } +} + +fn run_analyze_request(input: &str) -> Result { let request = parse_request(input)?; analyze_request(&request) } @@ -214,6 +297,16 @@ struct Request { diarization_spans: Vec, } +#[derive(Debug, Clone, PartialEq, Eq)] +struct DiscoveryClusterRequest { + embeddings_f32le_path: String, + rows: usize, + cols: usize, + byte_count: usize, + min_cluster_size: usize, + min_samples: usize, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum StatementAudioBuffer { Full, @@ -398,6 +491,47 @@ fn analyze_request(request: &Request) -> Result { Ok(response) } +fn run_discovery_cluster_request(input: &str) -> Result { + let request = parse_discovery_cluster_request(input)?; + let embeddings = read_embedding_payload_f32le(&request)?; + let labels = cluster_discovery_embeddings( + &embeddings, + request.rows, + request.cols, + request.min_cluster_size, + request.min_samples, + ) + .map_err(|error| map_discovery_clustering_error(&request.embeddings_f32le_path, error))?; + Ok(discovery_cluster_response_value(&request, &labels)) +} + +fn discovery_cluster_response_value( + request: &DiscoveryClusterRequest, + labels: &[Option], +) -> Value { + let noise_count = labels.iter().filter(|label| label.is_none()).count(); + let cluster_count = labels + .iter() + .filter_map(|label| *label) + .collect::>() + .len(); + let labels = labels + .iter() + .map(|label| label.map_or(-1_i64, |label| label as i64)) + .collect::>(); + json!({ + "schema": DISCOVERY_CLUSTER_RESPONSE_SCHEMA, + "labels": labels, + "cluster_count": cluster_count, + "noise_count": noise_count, + "parameters": { + "min_cluster_size": request.min_cluster_size, + "min_samples": request.min_samples, + }, + "algorithm": DISCOVERY_CLUSTER_ALGORITHM, + }) +} + fn statement_audio_buffer_for_request(request: &Request) -> StatementAudioBuffer { if request.reduced_audio_f32le_path.is_some() { StatementAudioBuffer::Reduced @@ -459,6 +593,39 @@ fn parse_request(input: &str) -> Result { }) } +fn parse_discovery_cluster_request(input: &str) -> Result { + let value: Value = + serde_json::from_str(input).map_err(|error| AnalyzeError::MalformedRequest { + detail: format!("request body is not valid JSON: {error}"), + })?; + let object = value + .as_object() + .ok_or_else(|| malformed("request body must be a JSON object"))?; + let schema = required_string(object, "schema")?; + if schema != DISCOVERY_CLUSTER_REQUEST_SCHEMA { + return Err(AnalyzeError::UnknownSchema { + schema: schema.to_string(), + }); + } + + let embeddings_f32le_path = required_string(object, "embeddings_f32le_path")?.to_string(); + let (rows, cols) = required_shape(object, "shape")?; + let byte_count = checked_payload_byte_count(rows, cols)?; + require_literal(object, "payload_format", PAYLOAD_FORMAT)?; + require_literal(object, "dtype", DTYPE_F32LE)?; + let min_cluster_size = required_usize_i64(object, "min_cluster_size")?; + let min_samples = required_usize_i64(object, "min_samples")?; + + Ok(DiscoveryClusterRequest { + embeddings_f32le_path, + rows, + cols, + byte_count, + min_cluster_size, + min_samples, + }) +} + fn required_object<'a>( object: &'a Map, field: &'static str, @@ -502,6 +669,76 @@ fn required_u32(object: &Map, field: &'static str) -> Result, + field: &'static str, +) -> Result { + let value = object + .get(field) + .and_then(Value::as_i64) + .ok_or_else(|| malformed(format!("{field} must be an i64 integer")))?; + if value < 0 { + return Err(malformed(format!("{field} must be non-negative"))); + } + usize::try_from(value).map_err(|_error| malformed(format!("{field} is out of range for usize"))) +} + +fn required_shape( + object: &Map, + field: &'static str, +) -> Result<(usize, usize), AnalyzeError> { + let shape = object + .get(field) + .and_then(Value::as_array) + .ok_or_else(|| malformed(format!("{field} must be [rows, cols]")))?; + if shape.len() != 2 { + return Err(malformed(format!( + "{field} must contain exactly two integers" + ))); + } + let rows = usize_from_i64_value(&shape[0], "shape[0]")?; + let cols = usize_from_i64_value(&shape[1], "shape[1]")?; + if cols == 0 { + return Err(malformed("shape[1] must be at least 1")); + } + Ok((rows, cols)) +} + +fn usize_from_i64_value(value: &Value, field: &'static str) -> Result { + let value = value + .as_i64() + .ok_or_else(|| malformed(format!("{field} must be an i64 integer")))?; + if value < 0 { + return Err(malformed(format!("{field} must be non-negative"))); + } + usize::try_from(value).map_err(|_error| malformed(format!("{field} is out of range for usize"))) +} + +fn checked_payload_byte_count(rows: usize, cols: usize) -> Result { + let values = rows + .checked_mul(cols) + .ok_or_else(|| malformed(format!("shape overflow: rows={rows} cols={cols}")))?; + values + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| { + malformed(format!( + "payload byte count overflow: rows={rows} cols={cols}" + )) + }) +} + +fn require_literal( + object: &Map, + field: &'static str, + expected: &'static str, +) -> Result<(), AnalyzeError> { + let actual = required_string(object, field)?; + if actual != expected { + return Err(malformed(format!("{field} must be {expected:?}"))); + } + Ok(()) +} + fn parse_spans( object: &Map, field_path: &'static str, @@ -648,6 +885,68 @@ fn read_audio_f32le(path: &str) -> Result, AnalyzeError> { Ok(audio) } +fn read_embedding_payload_f32le( + request: &DiscoveryClusterRequest, +) -> Result, AnalyzeError> { + let metadata = fs::metadata(&request.embeddings_f32le_path).map_err(|error| { + AnalyzeError::PayloadUnreadable { + path: request.embeddings_f32le_path.clone(), + detail: error.to_string(), + } + })?; + if !metadata.is_file() { + return Err(AnalyzeError::PayloadUnreadable { + path: request.embeddings_f32le_path.clone(), + detail: "path is not a regular file".to_string(), + }); + } + let actual_len = metadata.len(); + let expected_len = request.byte_count as u64; + if actual_len != expected_len { + return Err(AnalyzeError::PayloadInvalid { + path: request.embeddings_f32le_path.clone(), + detail: format!("byte length {actual_len} does not match expected {expected_len}"), + }); + } + + let mut file = fs::File::open(&request.embeddings_f32le_path).map_err(|error| { + AnalyzeError::PayloadUnreadable { + path: request.embeddings_f32le_path.clone(), + detail: error.to_string(), + } + })?; + let mut bytes = Vec::with_capacity(request.byte_count); + file.read_to_end(&mut bytes) + .map_err(|error| AnalyzeError::PayloadUnreadable { + path: request.embeddings_f32le_path.clone(), + detail: error.to_string(), + })?; + if bytes.len() != request.byte_count { + return Err(AnalyzeError::PayloadInvalid { + path: request.embeddings_f32le_path.clone(), + detail: format!( + "byte length {} changed after metadata check; expected {}", + bytes.len(), + request.byte_count + ), + }); + } + + let mut values = Vec::with_capacity(request.rows * request.cols); + for (index, chunk) in bytes.chunks_exact(4).enumerate() { + let value = f32::from_le_bytes(chunk.try_into().expect("four bytes")); + if !value.is_finite() { + return Err(AnalyzeError::PayloadNonFinite { + path: request.embeddings_f32le_path.clone(), + row: index / request.cols, + col: index % request.cols, + }); + } + values.push(value); + } + Ok(values) +} + fn preflight_model_path(field: &'static str, path: &str) -> Result<(), AnalyzeError> { let metadata = fs::metadata(path).map_err(|_error| AnalyzeError::ModelUnreadable { field, @@ -751,7 +1050,7 @@ fn diarization_with_interval_embeddings( interval_payload_path: Option<&str>, ) -> Result<(Value, Option), AnalyzeError> { let rows = valid_intervals.len(); - let cluster = cluster_embeddings( + let cluster = cluster_diarization_embeddings( interval_embedding_values, rows, WESPEAKER_EMBEDDING_SIZE, @@ -956,6 +1255,35 @@ fn map_diarization_error(error: DiarizationError) -> AnalyzeError { } } +fn map_discovery_clustering_error(path: &str, error: DiscoveryClusteringError) -> AnalyzeError { + match error { + DiscoveryClusteringError::InvalidMinClusterSize { .. } + | DiscoveryClusteringError::InvalidMinSamples { .. } + | DiscoveryClusteringError::MinSamplesExceedsRows { .. } + | DiscoveryClusteringError::NonUnitEmbeddingRow { .. } => AnalyzeError::PayloadInvalid { + path: path.to_string(), + detail: error.to_string(), + }, + DiscoveryClusteringError::NonFiniteCoordinate { row, col } => { + AnalyzeError::PayloadNonFinite { + path: path.to_string(), + row, + col, + } + } + DiscoveryClusteringError::ZeroColumns + | DiscoveryClusteringError::ShapeOverflow { .. } + | DiscoveryClusteringError::ShapeMismatch { .. } + | DiscoveryClusteringError::HdbscanEmptyDataset + | DiscoveryClusteringError::HdbscanWrongDimension { .. } + | DiscoveryClusteringError::HdbscanNonFiniteCoordinate { .. } + | DiscoveryClusteringError::HdbscanOutputLength { .. } + | DiscoveryClusteringError::HdbscanInvalidLabel { .. } => AnalyzeError::Internal { + detail: error.to_string(), + }, + } +} + fn map_open_onnx_error(field: &'static str, error: SpeakerOnnxError) -> AnalyzeError { match error { // Reserved under the current fixed provider plan: @@ -1042,6 +1370,18 @@ mod tests { }) } + fn base_discovery_cluster_request(path: &str, rows: usize, cols: usize) -> Value { + json!({ + "schema": DISCOVERY_CLUSTER_REQUEST_SCHEMA, + "embeddings_f32le_path": path, + "payload_format": PAYLOAD_FORMAT, + "dtype": DTYPE_F32LE, + "shape": [rows as i64, cols as i64], + "min_cluster_size": 3, + "min_samples": 2, + }) + } + fn request_string(value: Value) -> String { serde_json::to_string(&value).expect("request JSON") } @@ -1097,7 +1437,17 @@ mod tests { } #[test] - fn argv_rejects_any_argument_as_usage() { + fn argv_accepts_discovery_cluster_subcommand() { + assert_eq!( + evaluate_args(&[OsString::from(DISCOVERY_CLUSTER_COMMAND)]), + Ok(Command::DiscoveryCluster) + ); + } + + #[test] + fn argv_rejects_unknown_argument_as_usage() { + // Bare invocation is still Command::Run and no existing caller passes + // argv today, so accepting the new discovery-cluster token is additive. let error = evaluate_args(&[OsString::from("--help")]).unwrap_err(); let line = error_line_for_usage(&error); @@ -1116,6 +1466,105 @@ mod tests { assert_eq!(error.exit_code(), 64); } + #[test] + fn discovery_cluster_unknown_schema_is_rejected() { + let dir = TestDir::new(); + let mut request = base_discovery_cluster_request(&dir.path("embeddings.f32"), 0, 2); + request["schema"] = json!("solstone-speaker-discovery-cluster-request-v2"); + + let error = + run_command_request(Command::DiscoveryCluster, &request_string(request)).unwrap_err(); + + assert_eq!(error.reason(), "unknown-schema"); + assert_eq!(error.exit_code(), 64); + } + + #[test] + fn discovery_cluster_missing_payload_path_reports_payload_unreadable() { + let dir = TestDir::new(); + let request = base_discovery_cluster_request(&dir.path("missing.f32"), 6, 2); + + let error = + run_command_request(Command::DiscoveryCluster, &request_string(request)).unwrap_err(); + + assert_eq!(error.reason(), "payload-unreadable"); + assert_eq!(error.exit_code(), 69); + } + + #[test] + fn discovery_cluster_byte_length_mismatch_reports_payload_invalid() { + let dir = TestDir::new(); + let path = dir.path("embeddings.f32"); + write_f32le(&path, &[0.0_f32, 1.0]); + let request = base_discovery_cluster_request(&path, 2, 2); + + let error = + run_command_request(Command::DiscoveryCluster, &request_string(request)).unwrap_err(); + + assert_eq!(error.reason(), "payload-invalid"); + assert_eq!(error.exit_code(), 69); + } + + #[test] + fn discovery_cluster_happy_path_returns_labels_and_counts() { + let dir = TestDir::new(); + let path = dir.path("embeddings.f32"); + write_f32le( + &path, + &unit_rows_2d(&[ + (1.0, 0.0), + (1.0, 0.03), + (1.0, -0.03), + (-1.0, 0.0), + (-1.0, 0.03), + (-1.0, -0.03), + ]), + ); + let request = base_discovery_cluster_request(&path, 6, 2); + + let response = run_command_request(Command::DiscoveryCluster, &request_string(request)) + .expect("response"); + + assert_eq!(response["schema"], DISCOVERY_CLUSTER_RESPONSE_SCHEMA); + assert_eq!(response["labels"], json!([0, 0, 0, 1, 1, 1])); + assert_eq!(response["cluster_count"], 2); + assert_eq!(response["noise_count"], 0); + assert_eq!(response["parameters"]["min_cluster_size"], 3); + assert_eq!(response["parameters"]["min_samples"], 2); + assert_eq!(response["algorithm"], DISCOVERY_CLUSTER_ALGORITHM); + } + + #[test] + fn discovery_cluster_path_does_not_preflight_models() { + let dir = TestDir::new(); + let path = dir.path("embeddings.f32"); + write_f32le( + &path, + &unit_rows_2d(&[ + (1.0, 0.0), + (1.0, 0.03), + (1.0, -0.03), + (-1.0, 0.0), + (-1.0, 0.03), + (-1.0, -0.03), + ]), + ); + let mut request = base_discovery_cluster_request(&path, 6, 2); + request["models"] = json!({ + "pyannote_segmentation_onnx_path": dir.path("absent-pyannote.onnx"), + "wespeaker_onnx_path": dir.path("absent-wespeaker.onnx"), + }); + + // lib.rs:651-668 preflights model paths for the analyze path. These + // intentionally absent model paths prove the cluster path does not route + // through that preflight despite the ONNX crate import at module scope. + let response = run_command_request(Command::DiscoveryCluster, &request_string(request)) + .expect("response"); + + assert_eq!(response["schema"], DISCOVERY_CLUSTER_RESPONSE_SCHEMA); + assert_eq!(response["labels"], json!([0, 0, 0, 1, 1, 1])); + } + #[test] fn payload_paths_may_not_collide_with_each_other() { let mut request = base_request(); @@ -1535,4 +1984,14 @@ mod tests { } fs::write(path, bytes).expect("write audio"); } + + fn unit_rows_2d(points: &[(f32, f32)]) -> Vec { + let mut out = Vec::with_capacity(points.len() * 2); + for (x, y) in points { + let norm = (x * x + y * y).sqrt(); + out.push(*x / norm); + out.push(*y / norm); + } + out + } } diff --git a/core/crates/solstone-core-speakers-analyze/src/main.rs b/core/crates/solstone-core-speakers-analyze/src/main.rs index fae3ab2f1..0f9c57ccd 100644 --- a/core/crates/solstone-core-speakers-analyze/src/main.rs +++ b/core/crates/solstone-core-speakers-analyze/src/main.rs @@ -6,7 +6,8 @@ use std::io::{self, Read}; use std::process; use solstone_core_speakers_analyze::{ - error_json_line, error_line_for_analyze_error, error_line_for_usage, evaluate_args, run_request, + error_json_line, error_line_for_analyze_error, error_line_for_usage, evaluate_args, + run_command_request, }; const EXIT_USAGE: i32 = 64; @@ -15,10 +16,13 @@ const EXIT_TEMPFAIL: i32 = 75; fn main() { let args: Vec<_> = env::args_os().skip(1).collect(); - if let Err(error) = evaluate_args(&args) { - eprintln!("{}", error_line_for_usage(&error)); - process::exit(EXIT_USAGE); - } + let command = match evaluate_args(&args) { + Ok(command) => command, + Err(error) => { + eprintln!("{}", error_line_for_usage(&error)); + process::exit(EXIT_USAGE); + } + }; let mut input = String::new(); if let Err(error) = io::stdin().read_to_string(&mut input) { @@ -29,7 +33,7 @@ fn main() { process::exit(EXIT_TEMPFAIL); } - match run_request(&input) { + match run_command_request(command, &input) { Ok(response) => { println!( "{}", diff --git a/core/crates/solstone-core-speakers/Cargo.toml b/core/crates/solstone-core-speakers/Cargo.toml index 92c2ab66d..11fbf1dfa 100644 --- a/core/crates/solstone-core-speakers/Cargo.toml +++ b/core/crates/solstone-core-speakers/Cargo.toml @@ -6,6 +6,9 @@ rust-version.workspace = true license.workspace = true publish = false +[dependencies] +hdbscan.workspace = true + [lints] workspace = true diff --git a/core/crates/solstone-core-speakers/src/discovery.rs b/core/crates/solstone-core-speakers/src/discovery.rs new file mode 100644 index 000000000..e57f28320 --- /dev/null +++ b/core/crates/solstone-core-speakers/src/discovery.rs @@ -0,0 +1,627 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright (c) 2026 sol pbc + +use std::collections::BTreeMap; +use std::error::Error; +use std::fmt; + +use hdbscan::{DistanceMetric, Hdbscan, HdbscanError, HdbscanHyperParams, NnAlgorithm}; + +const UNIT_NORM_TOLERANCE: f64 = 1.0e-3; + +#[derive(Debug, Clone, PartialEq)] +pub enum DiscoveryClusteringError { + InvalidMinClusterSize { + actual: usize, + }, + InvalidMinSamples { + actual: usize, + }, + ZeroColumns, + ShapeOverflow { + rows: usize, + cols: usize, + }, + ShapeMismatch { + rows: usize, + cols: usize, + len: usize, + }, + MinSamplesExceedsRows { + min_samples: usize, + rows: usize, + }, + NonFiniteCoordinate { + row: usize, + col: usize, + }, + NonUnitEmbeddingRow { + row: usize, + norm: f64, + }, + HdbscanEmptyDataset, + HdbscanWrongDimension { + detail: String, + }, + HdbscanNonFiniteCoordinate { + detail: String, + }, + HdbscanOutputLength { + expected: usize, + actual: usize, + }, + HdbscanInvalidLabel { + label: i32, + }, +} + +impl fmt::Display for DiscoveryClusteringError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::InvalidMinClusterSize { actual } => write!( + formatter, + "invalid min_cluster_size: expected at least 2, got {actual}" + ), + Self::InvalidMinSamples { actual } => write!( + formatter, + "invalid min_samples: expected at least 1, got {actual}" + ), + Self::ZeroColumns => { + write!(formatter, "row-major matrix must have at least one column") + } + Self::ShapeOverflow { rows, cols } => { + write!( + formatter, + "row-major matrix shape overflow: rows={rows} cols={cols}" + ) + } + Self::ShapeMismatch { rows, cols, len } => write!( + formatter, + "row-major matrix length mismatch: rows={rows} cols={cols} len={len}" + ), + Self::MinSamplesExceedsRows { min_samples, rows } => write!( + formatter, + "invalid min_samples: min_samples={min_samples} exceeds rows={rows}" + ), + Self::NonFiniteCoordinate { row, col } => { + write!( + formatter, + "matrix value at row={row} col={col} is not finite" + ) + } + Self::NonUnitEmbeddingRow { row, norm } => write!( + formatter, + "embedding row {row} is not unit length: norm={norm}" + ), + Self::HdbscanEmptyDataset => write!(formatter, "hdbscan rejected an empty dataset"), + Self::HdbscanWrongDimension { detail } => { + write!(formatter, "hdbscan rejected matrix dimensions: {detail}") + } + Self::HdbscanNonFiniteCoordinate { detail } => { + write!( + formatter, + "hdbscan rejected a non-finite coordinate: {detail}" + ) + } + Self::HdbscanOutputLength { expected, actual } => write!( + formatter, + "hdbscan returned {actual} labels for {expected} rows" + ), + Self::HdbscanInvalidLabel { label } => { + write!(formatter, "hdbscan returned invalid label {label}") + } + } + } +} + +impl Error for DiscoveryClusteringError {} + +/// Cluster unknown speaker embeddings with the verified `hdbscan` EOM kernel. +/// +/// The caller supplies all sklearn-tuned parameters explicitly. Noise is returned +/// as `None`; consumers that need sklearn's JSON-facing `-1` sentinel should map +/// it at that boundary. +/// +/// Degenerate inputs are named deliberately: invalid HDBSCAN parameters error, +/// zero rows returns an empty result for direct callers, too few rows returns all +/// noise to preserve the Python discovery early-out, excessive `min_samples` +/// mirrors sklearn's fit-time error, and non-finite coordinates error instead +/// of sklearn's non-finite-row remapping. +pub fn cluster_embeddings( + embeddings: &[f32], + rows: usize, + cols: usize, + min_cluster_size: usize, + min_samples: usize, +) -> Result>, DiscoveryClusteringError> { + if min_cluster_size < 2 { + // sklearn rejects this before fit proceeds + // (.venv/lib/python3.13/site-packages/sklearn/cluster/_hdbscan/hdbscan.py:640-642). + return Err(DiscoveryClusteringError::InvalidMinClusterSize { + actual: min_cluster_size, + }); + } + if min_samples == 0 { + // sklearn rejects zero min_samples alongside the min_cluster_size guard + // (.venv/lib/python3.13/site-packages/sklearn/cluster/_hdbscan/hdbscan.py:642). + return Err(DiscoveryClusteringError::InvalidMinSamples { + actual: min_samples, + }); + } + if cols == 0 { + // A zero-width Euclidean space would make every row indistinguishable, + // so direct callers get a named shape error instead of accidental labels. + return Err(DiscoveryClusteringError::ZeroColumns); + } + validate_shape(embeddings, rows, cols)?; + if rows == 0 { + // Python has no direct kernel call for rows==0. The caller normally + // short-circuits before clustering, so this protects direct callers and + // mirrors diarization.rs:238-248. + return Ok(Vec::new()); + } + if rows < min_cluster_size { + // Keep solstone/apps/speakers/discovery.py:306-308 reachable: too few + // points for a legal cluster is not an error, and this happens before + // min_samples is compared with rows. + return Ok(vec![None; rows]); + } + if min_samples > rows { + // The hdbscan crate validates shape and infinities + // (~/.cargo/registry/src/.../hdbscan-0.12.0/src/validation.rs:16-47) + // but does not reject min_samples > rows before k-NN indexing. + return Err(DiscoveryClusteringError::MinSamplesExceedsRows { min_samples, rows }); + } + validate_finite(embeddings, rows, cols)?; + validate_unit_rows(embeddings, rows, cols)?; + + let data = matrix_as_f64_rows(embeddings, rows, cols); + let hyper_params = HdbscanHyperParams::builder() + .min_cluster_size(min_cluster_size) + .min_samples(min_samples) + .max_cluster_size(usize::MAX) + .allow_single_cluster(false) + .epsilon(0.0) + .dist_metric(DistanceMetric::Euclidean) + // The crate exposes the nearest-neighbour choice. Pin KD-tree instead + // of Auto so rows <= 250 do not take the crate's brute-force branch + // (~/.cargo/registry/src/.../hdbscan-0.12.0/src/core_distances.rs:1-18), + // while sklearn resolves this production metric to KD-tree + // (.venv/lib/python3.13/site-packages/sklearn/cluster/_hdbscan/hdbscan.py:848-857). + .nn_algorithm(NnAlgorithm::KdTree) + .build(); + let clusterer = Hdbscan::new(&data, hyper_params); + let raw_labels = clusterer.cluster().map_err(map_hdbscan_error)?; + option_labels_from_hdbscan(raw_labels, rows) +} + +fn validate_shape( + embeddings: &[f32], + rows: usize, + cols: usize, +) -> Result<(), DiscoveryClusteringError> { + let expected = rows + .checked_mul(cols) + .ok_or(DiscoveryClusteringError::ShapeOverflow { rows, cols })?; + if embeddings.len() != expected { + return Err(DiscoveryClusteringError::ShapeMismatch { + rows, + cols, + len: embeddings.len(), + }); + } + Ok(()) +} + +fn validate_finite( + embeddings: &[f32], + rows: usize, + cols: usize, +) -> Result<(), DiscoveryClusteringError> { + for row in 0..rows { + for col in 0..cols { + if !embeddings[row * cols + col].is_finite() { + // Deliberate divergence: sklearn remaps non-finite rows before + // fitting (.venv/lib/python3.13/site-packages/sklearn/cluster/_hdbscan/hdbscan.py:747-772). + // The Rust kernel names the bad coordinate instead. This also + // prevents the hdbscan crate's NaN path from reaching + // partial_cmp(...).expect("Invalid float") + // (~/.cargo/registry/src/.../hdbscan-0.12.0/src/core_distances/serial.rs:76-78). + return Err(DiscoveryClusteringError::NonFiniteCoordinate { row, col }); + } + } + } + Ok(()) +} + +fn validate_unit_rows( + embeddings: &[f32], + rows: usize, + cols: usize, +) -> Result<(), DiscoveryClusteringError> { + for row in 0..rows { + let mut norm_squared = 0.0; + for col in 0..cols { + let value = f64::from(embeddings[row * cols + col]); + norm_squared += value * value; + } + let norm = norm_squared.sqrt(); + if (norm - 1.0).abs() > UNIT_NORM_TOLERANCE { + // solstone/think/entities/voiceprints.py:31 normalize_embedding, + // reached through solstone/apps/speakers/discovery.py:271, is the + // production normalization point. This assertion catches callers + // that silently stop normalizing; in-kernel normalization would be + // idempotent and invisible to tests on already-normalized inputs. + return Err(DiscoveryClusteringError::NonUnitEmbeddingRow { row, norm }); + } + } + Ok(()) +} + +fn matrix_as_f64_rows(embeddings: &[f32], rows: usize, cols: usize) -> Vec> { + let mut out = Vec::with_capacity(rows); + for row in 0..rows { + let start = row * cols; + // sklearn casts HDBSCAN fit input to float64 before clustering + // (.venv/lib/python3.13/site-packages/sklearn/cluster/_hdbscan/hdbscan.py:739-745). + // The hdbscan crate API takes &[Vec], so this necessarily allocates + // one Vec per row before invoking the dependency. + out.push( + embeddings[start..start + cols] + .iter() + .map(|value| f64::from(*value)) + .collect(), + ); + } + out +} + +fn map_hdbscan_error(error: HdbscanError) -> DiscoveryClusteringError { + match error { + HdbscanError::EmptyDataset => DiscoveryClusteringError::HdbscanEmptyDataset, + HdbscanError::WrongDimension(detail) => { + DiscoveryClusteringError::HdbscanWrongDimension { detail } + } + HdbscanError::NonFiniteCoordinate(detail) => { + DiscoveryClusteringError::HdbscanNonFiniteCoordinate { detail } + } + } +} + +fn option_labels_from_hdbscan( + labels: Vec, + rows: usize, +) -> Result>, DiscoveryClusteringError> { + if labels.len() != rows { + return Err(DiscoveryClusteringError::HdbscanOutputLength { + expected: rows, + actual: labels.len(), + }); + } + + let mut members_by_raw_label: BTreeMap> = BTreeMap::new(); + let mut out = vec![None; rows]; + for (row, label) in labels.into_iter().enumerate() { + match label { + -1 => {} + 0..=i32::MAX => { + members_by_raw_label.entry(label).or_default().push(row); + } + _ => return Err(DiscoveryClusteringError::HdbscanInvalidLabel { label }), + } + } + + let mut clusters: Vec> = members_by_raw_label.into_values().collect(); + clusters.sort_by_key(|members| members[0]); + for (label, members) in clusters.into_iter().enumerate() { + for row in members { + out[row] = Some(label); + } + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::BTreeSet; + use std::time::Instant; + + #[test] + fn empty_matrix_returns_empty_labels() { + let labels = cluster_embeddings(&[], 0, 2, 5, 3).expect("empty matrix clusters"); + + assert!(labels.is_empty()); + } + + #[test] + fn fewer_rows_than_min_cluster_size_returns_noise_before_min_samples_row_check() { + let matrix = vec![0.0_f32, 0.0, 1.0, 1.0, 2.0, 2.0]; + + let labels = + cluster_embeddings(&matrix, 3, 2, 5, 99).expect("too few rows is not an error"); + + assert_eq!(labels, vec![None, None, None]); + } + + #[test] + fn all_noise_result_is_representable() { + let matrix = unit_rows_2d(&[(1.0, 0.0), (0.0, 1.0), (-1.0, 0.0), (0.0, -1.0), (1.0, 1.0)]); + + let labels = cluster_embeddings(&matrix, 5, 2, 5, 2).expect("matrix clusters"); + + assert_eq!(labels, vec![None, None, None, None, None]); + } + + #[test] + fn well_separated_clusters_are_recovered_as_partition() { + let matrix = unit_rows_2d(&[ + (1.0, 0.0), + (1.0, 0.03), + (1.0, -0.03), + (-1.0, 0.0), + (-1.0, 0.03), + (-1.0, -0.03), + ]); + + let labels = cluster_embeddings(&matrix, 6, 2, 3, 2).expect("matrix clusters"); + + assert_eq!( + partition(&labels), + partition_from_sets(&[&[0, 1, 2], &[3, 4, 5]]) + ); + } + + #[test] + fn min_cluster_size_boundary_excludes_root_inherited_exact_size_group() { + let matrix = unit_rows_2d(&[ + (1.0, 0.0), + (1.0, 0.03), + (1.0, -0.03), + (-1.0, 0.0), + (-1.0, 0.03), + ]); + + let labels = cluster_embeddings(&matrix, 5, 2, 3, 2).expect("matrix clusters"); + + assert_eq!(labels, vec![None, None, None, None, None]); + } + + #[test] + fn non_unit_embedding_rows_are_rejected_instead_of_normalized() { + let matrix = vec![ + 0.9_f32, 0.0, 1.0, 0.0, 1.1, 0.0, 9.0, 0.0, 10.0, 0.0, 11.0, 0.0, + ]; + + let error = cluster_embeddings(&matrix, 6, 2, 3, 2).unwrap_err(); + + let DiscoveryClusteringError::NonUnitEmbeddingRow { row, norm } = error else { + panic!("expected NonUnitEmbeddingRow"); + }; + assert_eq!(row, 0); + assert!((norm - 0.9).abs() < 1.0e-6); + } + + #[test] + fn many_well_separated_256_dim_groups_recover_expected_cluster_count() { + let rows = 2_000; + let cols = 256; + let (matrix, expected_clusters) = seeded_unit_normalized_discovery_matrix(rows, cols, 11); + + let labels = cluster_embeddings(&matrix, rows, cols, 5, 3).expect("matrix clusters"); + + assert_eq!(partition(&labels).len(), expected_clusters); + } + + #[test] + fn invalid_min_cluster_size_is_error() { + assert_eq!( + cluster_embeddings(&[], 0, 2, 1, 1), + Err(DiscoveryClusteringError::InvalidMinClusterSize { actual: 1 }) + ); + } + + #[test] + fn invalid_min_samples_is_error() { + assert_eq!( + cluster_embeddings(&[], 0, 2, 2, 0), + Err(DiscoveryClusteringError::InvalidMinSamples { actual: 0 }) + ); + } + + #[test] + fn zero_columns_is_error() { + assert_eq!( + cluster_embeddings(&[], 0, 0, 2, 1), + Err(DiscoveryClusteringError::ZeroColumns) + ); + } + + #[test] + fn min_samples_greater_than_rows_is_error() { + let matrix = unit_rows_2d(&[(1.0, 0.0), (-1.0, 0.0)]); + + assert_eq!( + cluster_embeddings(&matrix, 2, 2, 2, 3), + Err(DiscoveryClusteringError::MinSamplesExceedsRows { + min_samples: 3, + rows: 2 + }) + ); + } + + #[test] + fn non_finite_coordinate_is_error() { + let matrix = vec![1.0_f32, 0.0, f32::NAN, 1.0]; + + assert_eq!( + cluster_embeddings(&matrix, 2, 2, 2, 1), + Err(DiscoveryClusteringError::NonFiniteCoordinate { row: 1, col: 0 }) + ); + } + + #[test] + fn shape_overflow_is_error() { + assert_eq!( + cluster_embeddings(&[], usize::MAX, 2, 2, 1), + Err(DiscoveryClusteringError::ShapeOverflow { + rows: usize::MAX, + cols: 2 + }) + ); + } + + #[test] + fn slice_length_mismatch_is_error() { + let matrix = vec![0.0_f32, 0.0, 1.0]; + + assert_eq!( + cluster_embeddings(&matrix, 2, 2, 2, 1), + Err(DiscoveryClusteringError::ShapeMismatch { + rows: 2, + cols: 2, + len: 3 + }) + ); + } + + #[test] + #[ignore = "release-mode wall-clock benchmark for the 10_000 x 256 discovery scan"] + fn benchmark_10_000_by_256_discovery_clustering() { + let rows = 10_000; + let cols = 256; + let (matrix, _expected_clusters) = + seeded_unit_normalized_discovery_matrix(rows, cols, 0x5eed_5eed_cafe_babe); + + let started = Instant::now(); + let labels = cluster_embeddings(&matrix, rows, cols, 5, 3).expect("benchmark clusters"); + let elapsed = started.elapsed(); + let noise_count = labels.iter().filter(|label| label.is_none()).count(); + let cluster_count: BTreeSet = labels.iter().filter_map(|label| *label).collect(); + + println!( + "discovery_hdbscan_benchmark rows={rows} cols={cols} profile={} elapsed_ms={} clusters={} noise={noise_count}", + if cfg!(debug_assertions) { + "debug" + } else { + "release" + }, + elapsed.as_millis(), + cluster_count.len() + ); + assert_eq!(labels.len(), rows); + } + + fn unit_rows_2d(points: &[(f32, f32)]) -> Vec { + let mut out = Vec::with_capacity(points.len() * 2); + for (x, y) in points { + let norm = (x * x + y * y).sqrt(); + out.push(*x / norm); + out.push(*y / norm); + } + out + } + + fn partition(labels: &[Option]) -> BTreeSet> { + let mut by_label: BTreeMap> = BTreeMap::new(); + for (index, label) in labels.iter().enumerate() { + if let Some(label) = label { + by_label.entry(*label).or_default().insert(index); + } + } + by_label.into_values().collect() + } + + fn partition_from_sets(clusters: &[&[usize]]) -> BTreeSet> { + clusters + .iter() + .map(|cluster| cluster.iter().copied().collect()) + .collect() + } + + fn seeded_unit_normalized_discovery_matrix( + rows: usize, + cols: usize, + seed: u64, + ) -> (Vec, usize) { + let mut rng = Lcg::new(seed); + let marginal = 5; + let group_count = 3.max((rows - marginal) / 60); + let mut sizes = vec![(rows - marginal) / group_count; group_count]; + sizes[0] += (rows - marginal) - sizes.iter().sum::(); + sizes.push(marginal); + + let mut centers = Vec::with_capacity(sizes.len()); + for _ in 0..sizes.len() { + let mut center = Vec::with_capacity(cols); + for _ in 0..cols { + center.push(rng.next_standard_normal()); + } + normalize_f64_row(&mut center); + centers.push(center); + } + + let mut matrix = Vec::with_capacity(rows * cols); + for (group_index, size) in sizes.iter().enumerate() { + let total_spread = if group_index == sizes.len() - 1 { + 0.45 + } else { + 0.12 + }; + let per_dimension_spread = total_spread / (cols as f64).sqrt(); + for _ in 0..*size { + let mut row = Vec::with_capacity(cols); + for center_value in ¢ers[group_index] { + row.push(*center_value + rng.next_standard_normal() * per_dimension_spread); + } + normalize_f64_row(&mut row); + matrix.extend(row.into_iter().map(|value| value as f32)); + } + } + + (matrix, sizes.len()) + } + + fn normalize_f64_row(row: &mut [f64]) { + let norm = row.iter().map(|value| value * value).sum::().sqrt(); + for value in row { + *value /= norm; + } + } + + #[derive(Debug, Clone, Copy)] + struct Lcg { + state: u64, + spare_normal: Option, + } + + impl Lcg { + fn new(seed: u64) -> Self { + Self { + state: seed, + spare_normal: None, + } + } + + fn next_unit_f64(&mut self) -> f64 { + self.state = self + .state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1); + let bits = self.state >> 11; + bits as f64 / (1_u64 << 53) as f64 + } + + fn next_standard_normal(&mut self) -> f64 { + if let Some(value) = self.spare_normal.take() { + return value; + } + let u1 = self.next_unit_f64().max(f64::MIN_POSITIVE); + let u2 = self.next_unit_f64(); + let radius = (-2.0 * u1.ln()).sqrt(); + let angle = std::f64::consts::TAU * u2; + self.spare_normal = Some(radius * angle.sin()); + radius * angle.cos() + } + } +} diff --git a/core/crates/solstone-core-speakers/src/lib.rs b/core/crates/solstone-core-speakers/src/lib.rs index 8a92d43d7..7d4294795 100644 --- a/core/crates/solstone-core-speakers/src/lib.rs +++ b/core/crates/solstone-core-speakers/src/lib.rs @@ -9,6 +9,7 @@ use std::error::Error; use std::fmt; pub mod diarization; +pub mod discovery; pub use filterbank::{ WESPEAKER_EMBEDDING_SIZE, WESPEAKER_FFT_SIZE, WESPEAKER_FRAME_LENGTH_SAMPLES, diff --git a/docs/PORTING.md b/docs/PORTING.md index 28de79abd..a15052bb9 100644 --- a/docs/PORTING.md +++ b/docs/PORTING.md @@ -26,7 +26,9 @@ 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` stays in the iOS canary because its DSP and discovery +clustering graph remains Rust-only; the `hdbscan`/`kdtree` clustering crates add +no C/C++ build steps or native linkage. `solstone-core-speakers-analyze` and `solstone-core-speakers-onnx` are excluded: the analyzer transitively depends on the ONNX Runtime host native-runtime adapter, which is not mobile-ready subsystem logic. @@ -250,6 +252,12 @@ writes its report under the harness work directory unless `--report` is supplied `tests/verify_speaker_differential.py` runs the local speaker-pipeline differential harness and writes/compares versioned `.npz` result bundles for Python-to-port parity checks. + +`tests/verify_speaker_discovery_clustering_differential.py` runs the +unknown-speaker discovery clustering differential, feeding an `.npz` embedding +matrix to sklearn and the native analyzer `discovery-cluster` subcommand while +separating noise-boundary flips from cluster-to-cluster structural moves. + `tests/verify_speaker_verdict.py` consumes those recorded bundles without rerunning speaker models, adding decision-flip replay for clustering, owner-claim, and acoustic-tier outcomes plus DER scoring against diff --git a/tests/test_speaker_discovery_clustering_differential.py b/tests/test_speaker_discovery_clustering_differential.py new file mode 100644 index 000000000..9a5f33445 --- /dev/null +++ b/tests/test_speaker_discovery_clustering_differential.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright (c) 2026 sol pbc + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +from tests import verify_speaker_discovery_clustering_differential as harness + + +def _seeded_discovery_matrix( + rows: int, seed: int, cols: int = 256 +) -> tuple[np.ndarray, int]: + rng = np.random.default_rng(seed) + marginal = 5 + n_groups = max(3, (rows - marginal) // 60) + sizes = [(rows - marginal) // n_groups] * n_groups + sizes[0] += (rows - marginal) - sum(sizes) + sizes.append(marginal) + centers = rng.normal(size=(len(sizes), cols)) + centers /= np.linalg.norm(centers, axis=1, keepdims=True) + + groups = [] + for index, size in enumerate(sizes): + # Spread is total displacement from the unit-norm center. Divide by + # sqrt(cols), or a 256-dim "tight" group becomes a broad overlapping + # blob because every coordinate contributes variance. + total_spread = 0.45 if index == len(sizes) - 1 else 0.12 + spread = total_spread / np.sqrt(cols) + groups.append(centers[index] + rng.normal(scale=spread, size=(size, cols))) + + matrix = np.vstack(groups).astype(np.float32) + matrix /= np.linalg.norm(matrix, axis=1, keepdims=True) + return matrix.astype(np.float32), len(sizes) + + +def test_label_permutation_compares_equal() -> None: + report = harness._compare_clustering( + np.asarray([0, 0, 1, 1, -1], dtype=np.int64), + np.asarray([8, 8, 3, 3, -1], dtype=np.int64), + cols=2, + ) + + assert report["classification"] == harness.SHIP + assert report["partition_equal_up_to_relabelling"] is True + assert report["noise_to_clustered"] == [] + assert report["clustered_to_noise"] == [] + assert report["cluster_to_cluster_moves"] == [] + + +def test_noise_flip_classifies_separately() -> None: + left = np.asarray([0] * 500 + [1] * 500, dtype=np.int64) + right = left.copy() + right[0] = harness.NOISE + + report = harness._compare_clustering(left, right, cols=256) + + assert report["classification"] == harness.INVESTIGATE + assert report["noise_to_clustered"] == [] + assert report["clustered_to_noise"] == [0] + assert report["cluster_to_cluster_moves"] == [] + + +def test_cluster_to_cluster_move_fails_structurally() -> None: + report = harness._compare_clustering( + np.asarray([0, 0, 1, 1], dtype=np.int64), + np.asarray([0, 1, 1, 1], dtype=np.int64), + cols=2, + ) + + assert report["classification"] == harness.TRAILS_THE_WAVE + assert report["cluster_to_cluster_moves"] + + +def test_seeded_discovery_matrix_is_unit_normalized_npz_input(tmp_path: Path) -> None: + matrix, expected_clusters = _seeded_discovery_matrix(200, seed=11) + + assert matrix.shape == (200, 256) + assert matrix.dtype == np.float32 + assert np.allclose(np.linalg.norm(matrix, axis=1), 1.0, atol=1e-6) + + matrix_path = tmp_path / "matrix.npz" + np.savez(matrix_path, embeddings=matrix) + + loaded = harness.load_matrix(matrix_path) + labels = harness.run_sklearn(loaded) + + assert loaded.flags["C_CONTIGUOUS"] + assert harness._cluster_count(labels) == expected_clusters + assert harness._noise_count(labels) == 0 + + +def test_refuses_in_repo_report_destination() -> None: + with pytest.raises(harness.HarnessError): + harness._refuse_repo_destination(harness.ROOT / "report.json") + + +def test_rust_bin_is_required() -> None: + with pytest.raises(SystemExit) as exc: + harness.parse_args(["matrix.npz"]) + + assert exc.value.code == 2 diff --git a/tests/verify_speaker_discovery_clustering_differential.py b/tests/verify_speaker_discovery_clustering_differential.py new file mode 100644 index 000000000..10dfb241e --- /dev/null +++ b/tests/verify_speaker_discovery_clustering_differential.py @@ -0,0 +1,419 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright (c) 2026 sol pbc + +"""Differential harness for unknown-speaker discovery clustering.""" + +from __future__ import annotations + +import argparse +import json +import platform +import subprocess +import sys +import tempfile +from collections.abc import Sequence +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +import numpy as np +import sklearn +from sklearn.cluster import HDBSCAN + +from solstone.apps.speakers.discovery import MIN_CLUSTER_SIZE, MIN_SAMPLES +from solstone.think.utils import get_rev + +ROOT = Path(__file__).resolve().parent.parent +REPORT_SCHEMA = "solstone-speaker-discovery-clustering-differential-report" +SCHEMA_VERSION = 1 + +REQUEST_SCHEMA = "solstone-speaker-discovery-cluster-request-v1" +RESPONSE_SCHEMA = "solstone-speaker-discovery-cluster-response-v1" +PAYLOAD_FORMAT = "raw-f32le-row-major-v1" +DTYPE_F32LE = "float32-le" +RUST_COMMAND = "discovery-cluster" + +NOISE = -1 + +SHIP = "ship" +INVESTIGATE = "investigate" +TRAILS_THE_WAVE = "trails-the-wave" +HARNESS_ERROR = "harness-error" + + +class HarnessError(RuntimeError): + """Raised when the harness cannot produce a trustworthy comparison.""" + + +def _json_default(value: object) -> object: + if isinstance(value, Path): + return str(value) + if isinstance(value, np.generic): + return value.item() + raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable") + + +def _render_report(report: dict[str, Any]) -> str: + return json.dumps(report, default=_json_default, indent=2, sort_keys=True) + "\n" + + +def _refuse_repo_destination(path: Path) -> Path: + resolved = path.expanduser().resolve() + if resolved == ROOT or resolved.is_relative_to(ROOT): + raise HarnessError( + f"speaker differential refuses in-repo destination: {resolved}" + ) + return resolved + + +def _provenance(*, matrix_path: Path | None, rust_bin: Path | None) -> dict[str, Any]: + return { + "generated_at": datetime.now(UTC).isoformat().replace("+00:00", "Z"), + "harness": { + "name": "tests.verify_speaker_discovery_clustering_differential", + "repo_commit": get_rev(), + "schema_version": SCHEMA_VERSION, + }, + "host": { + "platform": platform.platform(), + "python": platform.python_version(), + }, + "versions": { + "numpy": np.__version__, + "sklearn": sklearn.__version__, + }, + "inputs": { + "matrix_path": str(matrix_path) if matrix_path is not None else None, + "rust_bin": str(rust_bin) if rust_bin is not None else None, + }, + } + + +def _base_report( + *, matrix_path: Path | None = None, rust_bin: Path | None = None +) -> dict[str, Any]: + return { + "schema": REPORT_SCHEMA, + "schema_version": SCHEMA_VERSION, + "classification": HARNESS_ERROR, + "failure": None, + "provenance": _provenance(matrix_path=matrix_path, rust_bin=rust_bin), + "parameters": { + "min_cluster_size": MIN_CLUSTER_SIZE, + "min_samples": MIN_SAMPLES, + }, + "rows": None, + "cols": None, + "sklearn": { + "cluster_count": None, + "noise_count": None, + }, + "rust": { + "cluster_count": None, + "noise_count": None, + }, + "partition_equal_up_to_relabelling": None, + "noise_to_clustered": [], + "clustered_to_noise": [], + "cluster_to_cluster_moves": [], + } + + +def load_matrix(path: Path) -> np.ndarray: + if path.suffix != ".npz": + raise HarnessError("matrix input must be a .npz file") + with np.load(path) as payload: + if "embeddings" not in payload: + raise HarnessError("matrix .npz must contain an embeddings array") + matrix = payload["embeddings"] + if matrix.ndim != 2: + raise HarnessError(f"embeddings must be 2-D, got shape {matrix.shape}") + if matrix.dtype != np.float32: + raise HarnessError(f"embeddings dtype must be float32, got {matrix.dtype}") + if not np.isfinite(matrix).all(): + raise HarnessError("embeddings contain non-finite values") + return np.ascontiguousarray(matrix) + + +def run_sklearn(matrix: np.ndarray) -> np.ndarray: + clusterer = HDBSCAN( + min_cluster_size=MIN_CLUSTER_SIZE, + min_samples=MIN_SAMPLES, + metric="euclidean", + cluster_selection_method="eom", + ) + return clusterer.fit_predict(matrix).astype(np.int64, copy=False) + + +def run_rust( + matrix: np.ndarray, rust_bin: Path, *, temp_parent: Path | None = None +) -> np.ndarray: + with tempfile.TemporaryDirectory( + prefix="solstone-speaker-discovery-", dir=temp_parent + ) as temp_dir: + payload_path = Path(temp_dir) / "embeddings.f32" + matrix.astype(" int: + return len({int(label) for label in labels if int(label) != NOISE}) + + +def _noise_count(labels: Sequence[int]) -> int: + return sum(1 for label in labels if int(label) == NOISE) + + +def _partition(labels: Sequence[int]) -> set[frozenset[int]]: + groups: dict[int, set[int]] = {} + for index, label in enumerate(labels): + label = int(label) + if label == NOISE: + continue + groups.setdefault(label, set()).add(index) + return {frozenset(members) for members in groups.values()} + + +def _partition_without(labels: Sequence[int], ignored: set[int]) -> set[frozenset[int]]: + groups = set() + for members in _partition(labels): + remaining = frozenset(index for index in members if index not in ignored) + if remaining: + groups.add(remaining) + return groups + + +def _index_members( + labels: Sequence[int], *, ignored: set[int] | None = None +) -> dict[int, frozenset[int]]: + ignored = ignored or set() + by_label: dict[int, set[int]] = {} + for index, label in enumerate(labels): + if index in ignored: + continue + label = int(label) + if label == NOISE: + continue + by_label.setdefault(label, set()).add(index) + out: dict[int, frozenset[int]] = {} + for members in by_label.values(): + frozen = frozenset(members) + for index in members: + out[index] = frozen + return out + + +def _compare_clustering( + sklearn_labels: Sequence[int], + rust_labels: Sequence[int], + *, + cols: int | None = None, +) -> dict[str, Any]: + if len(sklearn_labels) != len(rust_labels): + raise HarnessError( + f"label length mismatch: sklearn={len(sklearn_labels)} rust={len(rust_labels)}" + ) + rows = len(sklearn_labels) + sklearn_labels = [int(label) for label in sklearn_labels] + rust_labels = [int(label) for label in rust_labels] + noise_to_clustered = [ + index + for index, (left, right) in enumerate( + zip(sklearn_labels, rust_labels, strict=True) + ) + if left == NOISE and right != NOISE + ] + clustered_to_noise = [ + index + for index, (left, right) in enumerate( + zip(sklearn_labels, rust_labels, strict=True) + ) + if left != NOISE and right == NOISE + ] + ignored_noise_flips = set(noise_to_clustered) | set(clustered_to_noise) + sklearn_members = _index_members(sklearn_labels, ignored=ignored_noise_flips) + rust_members = _index_members(rust_labels, ignored=ignored_noise_flips) + cluster_to_cluster_moves = [ + index + for index in range(rows) + if sklearn_labels[index] != NOISE + and rust_labels[index] != NOISE + and index not in ignored_noise_flips + and sklearn_members[index] != rust_members[index] + ] + partition_equal = _partition(sklearn_labels) == _partition(rust_labels) + cluster_counts_match_after_noise_flips = _partition_without( + sklearn_labels, ignored_noise_flips + ) == _partition_without(rust_labels, ignored_noise_flips) + + classification = _classify( + rows=rows, + partition_equal=partition_equal, + sklearn_cluster_count=_cluster_count(sklearn_labels), + rust_cluster_count=_cluster_count(rust_labels), + noise_flip_count=len(ignored_noise_flips), + cluster_to_cluster_move_count=len(cluster_to_cluster_moves), + cluster_counts_match_after_noise_flips=cluster_counts_match_after_noise_flips, + ) + return { + "classification": classification, + "rows": rows, + "cols": cols, + "sklearn": { + "cluster_count": _cluster_count(sklearn_labels), + "noise_count": _noise_count(sklearn_labels), + }, + "rust": { + "cluster_count": _cluster_count(rust_labels), + "noise_count": _noise_count(rust_labels), + }, + "partition_equal_up_to_relabelling": partition_equal, + "noise_to_clustered": noise_to_clustered, + "clustered_to_noise": clustered_to_noise, + "cluster_to_cluster_moves": cluster_to_cluster_moves, + "cluster_counts_match_after_noise_flips": cluster_counts_match_after_noise_flips, + } + + +def _classify( + *, + rows: int, + partition_equal: bool, + sklearn_cluster_count: int, + rust_cluster_count: int, + noise_flip_count: int, + cluster_to_cluster_move_count: int, + cluster_counts_match_after_noise_flips: bool, +) -> str: + if ( + partition_equal + and noise_flip_count == 0 + and cluster_to_cluster_move_count == 0 + and sklearn_cluster_count == rust_cluster_count + ): + return SHIP + + # A noise-boundary flip is a point sitting on the stability margin where + # f64 summation order can legitimately move it, whereas a cluster-to-cluster + # move means the two implementations built different structure. Those are + # different findings with different consequences, and averaging them into + # one percentage would hide the one that matters. + noise_flips_within_threshold = ( + rows > 0 and noise_flip_count * 1000 <= rows and noise_flip_count <= 3 + ) + cluster_count_mismatch_explained = ( + sklearn_cluster_count == rust_cluster_count + or cluster_counts_match_after_noise_flips + ) + if ( + cluster_to_cluster_move_count == 0 + and noise_flip_count > 0 + and noise_flips_within_threshold + and cluster_count_mismatch_explained + ): + return INVESTIGATE + + return TRAILS_THE_WAVE + + +def compare_matrix( + matrix: np.ndarray, rust_bin: Path, *, temp_parent: Path | None = None +) -> dict[str, Any]: + sklearn_labels = run_sklearn(matrix) + rust_labels = run_rust(matrix, rust_bin, temp_parent=temp_parent) + return _compare_clustering(sklearn_labels, rust_labels, cols=int(matrix.shape[1])) + + +def compare_matrix_file(matrix_path: Path, rust_bin: Path) -> dict[str, Any]: + report = _base_report(matrix_path=matrix_path, rust_bin=rust_bin) + try: + matrix = load_matrix(matrix_path) + report.update( + compare_matrix( + matrix, + rust_bin, + temp_parent=_temp_parent_for_matrix(matrix_path), + ) + ) + except Exception as exc: + report["classification"] = HARNESS_ERROR + report["failure"] = {"class": HARNESS_ERROR, "message": str(exc)} + return report + + +def _temp_parent_for_matrix(matrix_path: Path) -> Path | None: + resolved = matrix_path.expanduser().resolve() + if resolved == ROOT or resolved.is_relative_to(ROOT): + return None + return resolved.parent + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("matrix_path", help="Input .npz containing embeddings") + parser.add_argument( + "--rust-bin", required=True, help="Path to Rust analyzer binary" + ) + parser.add_argument( + "--report", help="JSON report destination outside the repository" + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + requested_report_path = Path(args.report).resolve() if args.report else None + report_path: Path | None = None + try: + if requested_report_path is not None: + report_path = _refuse_repo_destination(requested_report_path) + matrix_path = Path(args.matrix_path) + rust_bin = Path(args.rust_bin) + report = compare_matrix_file(matrix_path, rust_bin) + except Exception as exc: + report = _base_report() + report["failure"] = {"class": HARNESS_ERROR, "message": str(exc)} + report["classification"] = HARNESS_ERROR + + rendered = _render_report(report) + if report_path is not None: + report_path.parent.mkdir(parents=True, exist_ok=True) + report_path.write_text(rendered, encoding="utf-8") + sys.stdout.write(rendered) + return 0 if report.get("classification") == SHIP else 1 + + +if __name__ == "__main__": + raise SystemExit(main())