diff --git a/CHANGELOG.md b/CHANGELOG.md index 0baf87fc..87678eb3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -75,6 +75,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Miri-friendly test target `tests/miri_resegmentation.rs` covering no-overlap, single-overlap, and centroid math paths. +### Added (M6a — Pipeline + Profile API, additive) +- `polyvoice::pipeline_v2` module: `Pipeline::builder()` returning + `PipelineBuilder` with `.profile(Profile::Mobile|Balanced|Custom)`, + `.with_models_from(ModelRegistry)`, `.with_segmenter/embedder/clusterer/resegmenter()`, + and a validated `.build()`. `PipelineConfig`, `ClustererKind`, + `ExecutionProvider`, and `ConfigError` all per spec §5.2/§5.4. +- `Pipeline::run(&samples, SampleRate)` orchestrates M1 segmenter → M2 + embedder (with `apply_overlap_mask` per primary chunk) → M3 clusterer → + M4 resegmenter → legacy `merge_segments` → `DiarizationResult`. +- New Cargo feature `pipeline_v2` (in default features). Requires + `onnx + segmentation + embedder + clusterer + resegmentation`; missing + any of these triggers a `compile_error!` with an actionable message. +- Public re-exports `polyvoice::PipelineV2`, `PipelineBuilder`, + `PipelineConfig`, `ClustererKind`, `ExecutionProvider`, `ConfigError`, + `PipelineV2Error`. Legacy `polyvoice::Pipeline` is unchanged; M6b will + rename `pipeline_v2 → pipeline` and remove the legacy code path. +- Synthetic integration test on Custom profile (`tests/pipeline_v2_synthetic_test.rs`, + 7 tests) + `#[ignore]` E2E test on Balanced profile via `ModelRegistry` + (`tests/pipeline_v2_e2e_test.rs`). + ### Added (M5 — INT8 quantization) - New scripts: `download-voxconverse-dev.sh`, `download-voxceleb1-subset.sh`, `quantize_models.py` + `quantize-models.sh`, `validate_int8.py` + diff --git a/Cargo.toml b/Cargo.toml index 7e6eddb3..28512cc1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ categories = ["multimedia::audio", "science"] crate-type = ["cdylib", "rlib"] [features] -default = ["spectral", "segmentation", "embedder", "clusterer", "resegmentation"] +default = ["spectral", "segmentation", "embedder", "clusterer", "resegmentation", "pipeline_v2"] # ONNX-based embedding extractors (pulls ort). onnx = ["dep:ort"] # C FFI bindings. @@ -52,6 +52,13 @@ clusterer = [] # caller (M6 Pipeline wires the embedder pool into this). resegmentation = [] +# v1.0 Pipeline + Profile builder API (additive in M6a; replaces legacy in M6b). +# Requires the full M1–M5 stack — the builder loads ONNX through ModelRegistry +# (download feature) for Mobile/Balanced profiles. The module gates itself +# with a `compile_error!` if onnx/segmentation/embedder/clusterer/resegmentation +# are missing, so half-wired feature combos cannot ship. +pipeline_v2 = ["download"] + # Model profile bundles. M0 only emits metadata; CLI uses these to pick which # manifest entries to ensure(). M5 will introduce INT8 model variants. profile-mobile = [] diff --git a/docs/superpowers/plans/2026-05-07-m6a-pipeline-v2-plan.md b/docs/superpowers/plans/2026-05-07-m6a-pipeline-v2-plan.md new file mode 100644 index 00000000..8248a5da --- /dev/null +++ b/docs/superpowers/plans/2026-05-07-m6a-pipeline-v2-plan.md @@ -0,0 +1,1773 @@ +# M6a — Pipeline + Profile API (additive) Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax. + +**Goal:** Add a new `polyvoice::pipeline_v2` module that exposes `Pipeline::builder()` API wiring M1–M5 (powerset segmenter + CAM++/ResNet34 embedder + NME-SC/AHC clusterer + OverlapResegmenter + INT8 manifest) into a single end-to-end `run(&samples, SampleRate) -> Result` call. **Additive only** — legacy `polyvoice::Pipeline` is untouched. M6b removes legacy and renames `pipeline_v2 → pipeline`. + +**Architecture:** New directory `src/pipeline_v2/{mod.rs, config.rs, builder.rs, mocks.rs}` gated behind a default-on `pipeline_v2` Cargo feature that requires `onnx + segmentation + embedder + clusterer + resegmentation`. Builder validates Profile vs Custom-component contracts, resolves Mobile/Balanced ONNX through `ModelRegistry` (M0), or accepts caller-supplied trait objects for Custom. `Pipeline::run()` orchestrates segment → mask + embed → cluster → resegment → merge → emit `DiarizationResult`. + +**Tech Stack:** Rust 2024. No new crate dependencies. Reuses M1 `Segmenter` + `RawSegment`, M2 `Embedder` + `EmbedderPool` + `apply_overlap_mask`, M3 `Clusterer`, M4 `Resegmenter` + `compute_centroids` + `extract_overlap_time_ranges`, M5 `ModelRegistry` + `Profile`. `std::thread::available_parallelism` for pool sizing (no `num_cpus` dep). + +--- + +## File structure + +| Path | Action | Responsibility | +|---|---|---| +| `Cargo.toml` | modify | Add `pipeline_v2` feature (default-on) | +| `src/pipeline_v2/mod.rs` | create | `Pipeline` struct, `run()`, `PipelineError`, feature-gate `compile_error!` | +| `src/pipeline_v2/config.rs` | create | `PipelineConfig` + `ClustererKind` + `ExecutionProvider` + Default | +| `src/pipeline_v2/builder.rs` | create | `PipelineBuilder` + `ConfigError` | +| `src/pipeline_v2/mocks.rs` | create | Test-only `MockSegmenter / MockEmbedder / MockClusterer` (`#[cfg(test)]`) | +| `src/lib.rs` | modify | `pub mod pipeline_v2;` gated, re-exports | +| `tests/pipeline_v2_synthetic_test.rs` | create | Builder validation + Custom profile end-to-end on synthetic data | +| `tests/pipeline_v2_e2e_test.rs` | create | `#[ignore]` integration test on real ONNX (Balanced profile) | +| `CHANGELOG.md` | modify | Unreleased M6a section | + +Total roughly 1100 LOC Rust + ~250 LOC tests + ~30 lines of doc. + +--- + +## Task 1: Add `pipeline_v2` Cargo feature + +**Files:** +- Modify: `Cargo.toml` + +- [ ] **Step 1.1: Update default + add feature** + +In `/Users/ekhodzitsky/Documents/personal/polyvoice/Cargo.toml`, find: + +```toml +default = ["spectral", "segmentation", "embedder", "clusterer", "resegmentation"] +``` + +Replace with: + +```toml +default = ["spectral", "segmentation", "embedder", "clusterer", "resegmentation", "pipeline_v2"] +``` + +After the `resegmentation = []` line, append: + +```toml + +# v1.0 Pipeline + Profile builder API (additive in M6a; replaces legacy in M6b). +# Requires the full M1–M5 stack: onnx, segmentation, embedder, clusterer, +# resegmentation. The new module compiles to a `compile_error!` outside that +# feature combo, so we don't accidentally ship a half-wired pipeline. +pipeline_v2 = [] +``` + +- [ ] **Step 1.2: Verify build matrix** + +```bash +cd /Users/ekhodzitsky/Documents/personal/polyvoice +cargo check +cargo check --features pipeline_v2 +cargo check --features pipeline_v2,onnx,segmentation,embedder,clusterer,resegmentation +cargo check --no-default-features +cargo check --no-default-features --features pipeline_v2 # MUST FAIL with compile_error +cargo check --target wasm32-unknown-unknown --no-default-features --features resegmentation --lib +cargo check --all-features +``` + +The `--no-default-features --features pipeline_v2` invocation will be expected to fail once Task 5 ships the `compile_error!` guard. For Task 1 it succeeds (empty module yet). Note this expectation in the commit message — Task 5 will tighten it. + +- [ ] **Step 1.3: Commit** + +```bash +git add Cargo.toml +git commit -m "feat(cargo): add pipeline_v2 feature flag for M6a" +``` + +--- + +## Task 2: `PipelineConfig` + `ClustererKind` + `ExecutionProvider` + +**Files:** +- Create: `src/pipeline_v2/config.rs` +- Create: `src/pipeline_v2/mod.rs` (stub — only `pub mod config;`) +- Modify: `src/lib.rs` (gated `pub mod pipeline_v2;`) + +- [ ] **Step 2.1: Write failing tests first** + +Create `/Users/ekhodzitsky/Documents/personal/polyvoice/src/pipeline_v2/config.rs`: + +```rust +//! M6a — `PipelineConfig`, `ClustererKind`, `ExecutionProvider`. +//! +//! Spec: `docs/superpowers/specs/2026-05-07-m6a-pipeline-v2-design.md` §3. + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::Profile; + + #[test] + fn pipeline_config_default_is_balanced() { + let cfg = PipelineConfig::default(); + assert_eq!(cfg.profile, Profile::Balanced); + assert_eq!(cfg.sample_rate.get(), 16000); + assert!((cfg.seg_window_secs - 10.0).abs() < f32::EPSILON); + assert!((cfg.seg_hop_secs - 0.5).abs() < f32::EPSILON); + assert!(matches!(cfg.clusterer, ClustererKind::NmeSc)); + assert_eq!(cfg.max_speakers, 20); + assert_eq!(cfg.min_cluster_size, 12); + assert!(cfg.resegment_overlap); + assert!((cfg.resegment_min_overlap_secs - 0.1).abs() < f32::EPSILON); + assert!((cfg.min_speech_secs - 0.25).abs() < f32::EPSILON); + assert!((cfg.max_gap_secs - 0.5).abs() < f32::EPSILON); + assert!(cfg.embedder_pool_size >= 1); + assert!(cfg.embedder_pool_size <= 4); + } + + #[test] + fn clusterer_kind_ahc_with_threshold() { + let k = ClustererKind::Ahc { threshold: 0.7 }; + if let ClustererKind::Ahc { threshold } = k { + assert!((threshold - 0.7).abs() < f32::EPSILON); + } else { + panic!("expected Ahc variant"); + } + } + + #[test] + fn execution_provider_auto_returns_some_variant() { + let ep = ExecutionProvider::auto(); + // Just assert it returns *some* variant — actual platform default + // varies (CoreMl on macOS aarch64, XnnPack on linux aarch64, Cpu else). + let _ = ep; + } +} +``` + +- [ ] **Step 2.2: Wire stub mod into lib.rs and pipeline_v2/mod.rs** + +Create `/Users/ekhodzitsky/Documents/personal/polyvoice/src/pipeline_v2/mod.rs`: + +```rust +//! M6a — additive `polyvoice::pipeline_v2` module. +//! +//! Spec: `docs/superpowers/specs/2026-05-07-m6a-pipeline-v2-design.md`. + +pub mod config; +``` + +In `/Users/ekhodzitsky/Documents/personal/polyvoice/src/lib.rs`, after the existing block: + +```rust +#[cfg(all(feature = "resegmentation", feature = "segmentation"))] +pub use resegmentation::extract_overlap_time_ranges; +``` + +append: + +```rust + +#[cfg(all( + feature = "pipeline_v2", + feature = "onnx", + feature = "segmentation", + feature = "embedder", + feature = "clusterer", + feature = "resegmentation", +))] +pub mod pipeline_v2; +``` + +- [ ] **Step 2.3: Confirm compile-failure** + +```bash +cd /Users/ekhodzitsky/Documents/personal/polyvoice +cargo test --all-features --lib pipeline_v2::config 2>&1 | head -20 +``` + +Expected: errors that `PipelineConfig`, `ClustererKind`, `ExecutionProvider` are not defined. + +- [ ] **Step 2.4: Implement config** + +Replace `src/pipeline_v2/config.rs` with: + +```rust +//! M6a — `PipelineConfig`, `ClustererKind`, `ExecutionProvider`. +//! +//! Spec: `docs/superpowers/specs/2026-05-07-m6a-pipeline-v2-design.md` §3. + +use crate::types::{Profile, SampleRate}; + +/// Top-level configuration for the v1.0 Pipeline. Mirrors spec §5.2 verbatim. +#[derive(Clone, Debug)] +pub struct PipelineConfig { + pub profile: Profile, + pub sample_rate: SampleRate, + pub seg_window_secs: f32, + pub seg_hop_secs: f32, + pub clusterer: ClustererKind, + pub max_speakers: u8, + pub min_cluster_size: usize, + pub resegment_overlap: bool, + pub resegment_min_overlap_secs: f32, + pub min_speech_secs: f32, + pub max_gap_secs: f32, + pub embedder_pool_size: usize, + pub execution_provider: ExecutionProvider, +} + +impl Default for PipelineConfig { + fn default() -> Self { + Self { + profile: Profile::Balanced, + sample_rate: SampleRate::new(16000).expect("16000 is a valid sample rate"), + seg_window_secs: 10.0, + seg_hop_secs: 0.5, + clusterer: ClustererKind::NmeSc, + max_speakers: 20, + min_cluster_size: 12, + resegment_overlap: true, + resegment_min_overlap_secs: 0.1, + min_speech_secs: 0.25, + max_gap_secs: 0.5, + embedder_pool_size: default_pool_size(), + execution_provider: ExecutionProvider::auto(), + } + } +} + +/// Clusterer selector. Defaults to `NmeSc` (M3). +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum ClustererKind { + NmeSc, + Ahc { threshold: f32 }, +} + +/// ONNX Runtime execution provider hint. The actual EP is plumbed at +/// `Session::builder()` time in M6b's CLI/FFI rewrite; M6a stores the +/// preference but the legacy ONNX wrappers (M2 `FbankOnnxExtractor`) only +/// honour the Cpu path so this field is informational for now. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum ExecutionProvider { + Cpu, + CoreMl, + Nnapi, + Cuda, + XnnPack, +} + +impl ExecutionProvider { + /// Best-effort platform default. macOS aarch64 → `CoreMl`; Linux aarch64 → + /// `XnnPack`; else `Cpu`. + pub fn auto() -> Self { + #[cfg(all(target_os = "macos", target_arch = "aarch64"))] + return Self::CoreMl; + #[cfg(all(target_os = "linux", target_arch = "aarch64"))] + return Self::XnnPack; + #[cfg(not(any( + all(target_os = "macos", target_arch = "aarch64"), + all(target_os = "linux", target_arch = "aarch64"), + )))] + return Self::Cpu; + } +} + +fn default_pool_size() -> usize { + std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(1) + .min(4) + .max(1) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::Profile; + + #[test] + fn pipeline_config_default_is_balanced() { + let cfg = PipelineConfig::default(); + assert_eq!(cfg.profile, Profile::Balanced); + assert_eq!(cfg.sample_rate.get(), 16000); + assert!((cfg.seg_window_secs - 10.0).abs() < f32::EPSILON); + assert!((cfg.seg_hop_secs - 0.5).abs() < f32::EPSILON); + assert!(matches!(cfg.clusterer, ClustererKind::NmeSc)); + assert_eq!(cfg.max_speakers, 20); + assert_eq!(cfg.min_cluster_size, 12); + assert!(cfg.resegment_overlap); + assert!((cfg.resegment_min_overlap_secs - 0.1).abs() < f32::EPSILON); + assert!((cfg.min_speech_secs - 0.25).abs() < f32::EPSILON); + assert!((cfg.max_gap_secs - 0.5).abs() < f32::EPSILON); + assert!(cfg.embedder_pool_size >= 1); + assert!(cfg.embedder_pool_size <= 4); + } + + #[test] + fn clusterer_kind_ahc_with_threshold() { + let k = ClustererKind::Ahc { threshold: 0.7 }; + if let ClustererKind::Ahc { threshold } = k { + assert!((threshold - 0.7).abs() < f32::EPSILON); + } else { + panic!("expected Ahc variant"); + } + } + + #[test] + fn execution_provider_auto_returns_some_variant() { + let ep = ExecutionProvider::auto(); + let _ = ep; + } +} +``` + +- [ ] **Step 2.5: Verify** + +```bash +cd /Users/ekhodzitsky/Documents/personal/polyvoice +cargo test --all-features --lib pipeline_v2::config +cargo fmt +cargo clippy --all-features --lib -- -D warnings +``` + +Expected: 3 tests pass, clippy clean. + +- [ ] **Step 2.6: Commit** + +```bash +git add src/pipeline_v2/ src/lib.rs +git commit -m "feat(pipeline_v2): add PipelineConfig + ClustererKind + ExecutionProvider" +``` + +--- + +## Task 3: `PipelineBuilder` + `ConfigError` + +**Files:** +- Create: `src/pipeline_v2/builder.rs` +- Modify: `src/pipeline_v2/mod.rs` (add `pub mod builder;`) + +- [ ] **Step 3.1: Append failing tests** + +Create `/Users/ekhodzitsky/Documents/personal/polyvoice/src/pipeline_v2/builder.rs`: + +```rust +//! M6a — `PipelineBuilder` + `ConfigError`. +//! +//! Spec: `docs/superpowers/specs/2026-05-07-m6a-pipeline-v2-design.md` §4. +//! +//! The builder *defers* construction of the inner `Pipeline` to `build()`, +//! at which point Profile vs Custom-component contracts are validated. + +use crate::clusterer::Clusterer; +use crate::embedder::Embedder; +use crate::models::{ModelRegistry, RegistryError}; +use crate::pipeline_v2::config::PipelineConfig; +use crate::resegmentation::Resegmenter; +use crate::segmentation::Segmenter; +use crate::types::Profile; + +/// Errors produced by `PipelineBuilder::build()`. +#[derive(Debug, thiserror::Error)] +pub enum ConfigError { + #[error("profile {profile:?} requires .with_models_from() call")] + MissingRegistry { profile: Profile }, + + #[error( + "profile {profile:?} cannot accept .with_{offending}() — Custom only" + )] + CustomComponentInProfile { + profile: Profile, + offending: &'static str, + }, + + #[error( + "Custom profile cannot accept .with_models_from() — supply components individually" + )] + RegistryInCustomProfile, + + #[error("Custom profile missing required components: {missing:?}")] + MissingCustomComponent { missing: Vec<&'static str> }, + + #[error("ONNX model not found in registry: {model_id}")] + UnknownModel { model_id: String }, + + #[error("registry resolution failed: {0}")] + Registry(#[from] RegistryError), +} + +/// Builder for the v1.0 `Pipeline`. Use `Pipeline::builder()` to obtain one. +pub struct PipelineBuilder { + pub(crate) config: PipelineConfig, + pub(crate) registry: Option, + pub(crate) custom_segmenter: Option>, + pub(crate) custom_embedder: Option>, + pub(crate) custom_clusterer: Option>, + pub(crate) custom_resegmenter: Option>, +} + +impl PipelineBuilder { + pub(crate) fn new() -> Self { + Self { + config: PipelineConfig::default(), + registry: None, + custom_segmenter: None, + custom_embedder: None, + custom_clusterer: None, + custom_resegmenter: None, + } + } + + pub fn config(mut self, cfg: PipelineConfig) -> Self { + self.config = cfg; + self + } + + pub fn profile(mut self, p: Profile) -> Self { + self.config.profile = p; + self + } + + pub fn with_models_from(mut self, r: ModelRegistry) -> Self { + self.registry = Some(r); + self + } + + pub fn with_segmenter(mut self, s: Box) -> Self { + self.custom_segmenter = Some(s); + self + } + + pub fn with_embedder(mut self, e: Box) -> Self { + self.custom_embedder = Some(e); + self + } + + pub fn with_clusterer(mut self, c: Box) -> Self { + self.custom_clusterer = Some(c); + self + } + + pub fn with_resegmenter(mut self, r: Box) -> Self { + self.custom_resegmenter = Some(r); + self + } + + pub fn resegment_overlap(mut self, on: bool) -> Self { + self.config.resegment_overlap = on; + self + } + + pub fn embedder_pool_size(mut self, n: usize) -> Self { + self.config.embedder_pool_size = n.max(1); + self + } + + pub fn max_speakers(mut self, n: u8) -> Self { + self.config.max_speakers = n; + self + } + + /// Validate Profile vs Custom-component contracts. Returns the staged + /// builder ready for `Pipeline::from_builder()` (M6a) or rejects via + /// `ConfigError`. + pub fn validate(&self) -> Result<(), ConfigError> { + match self.config.profile { + Profile::Mobile | Profile::Balanced => { + if self.registry.is_none() { + return Err(ConfigError::MissingRegistry { + profile: self.config.profile, + }); + } + if self.custom_segmenter.is_some() { + return Err(ConfigError::CustomComponentInProfile { + profile: self.config.profile, + offending: "segmenter", + }); + } + if self.custom_embedder.is_some() { + return Err(ConfigError::CustomComponentInProfile { + profile: self.config.profile, + offending: "embedder", + }); + } + if self.custom_clusterer.is_some() { + return Err(ConfigError::CustomComponentInProfile { + profile: self.config.profile, + offending: "clusterer", + }); + } + } + Profile::Custom => { + if self.registry.is_some() { + return Err(ConfigError::RegistryInCustomProfile); + } + let mut missing: Vec<&'static str> = Vec::new(); + if self.custom_segmenter.is_none() { + missing.push("segmenter"); + } + if self.custom_embedder.is_none() { + missing.push("embedder"); + } + if self.custom_clusterer.is_none() { + missing.push("clusterer"); + } + if !missing.is_empty() { + return Err(ConfigError::MissingCustomComponent { missing }); + } + } + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::pipeline_v2::config::PipelineConfig; + use crate::pipeline_v2::mocks::{MockClusterer, MockEmbedder, MockSegmenter}; + + fn fresh() -> PipelineBuilder { + PipelineBuilder::new() + } + + #[test] + fn builder_default_profile_balanced() { + let b = fresh(); + assert_eq!(b.config.profile, Profile::Balanced); + } + + #[test] + fn builder_profile_setter() { + let b = fresh().profile(Profile::Mobile); + assert_eq!(b.config.profile, Profile::Mobile); + } + + #[test] + fn validate_mobile_without_registry_errors() { + let err = fresh().profile(Profile::Mobile).validate().unwrap_err(); + assert!(matches!(err, ConfigError::MissingRegistry { profile: Profile::Mobile })); + } + + #[test] + fn validate_custom_without_components_errors() { + let err = fresh().profile(Profile::Custom).validate().unwrap_err(); + match err { + ConfigError::MissingCustomComponent { missing } => { + assert!(missing.contains(&"segmenter")); + assert!(missing.contains(&"embedder")); + assert!(missing.contains(&"clusterer")); + } + other => panic!("unexpected error variant: {other:?}"), + } + } + + #[test] + fn validate_custom_with_full_components_succeeds() { + let b = fresh() + .profile(Profile::Custom) + .with_segmenter(Box::new(MockSegmenter::default())) + .with_embedder(Box::new(MockEmbedder::default())) + .with_clusterer(Box::new(MockClusterer::default())); + b.validate().expect("custom + 3 components must validate"); + } + + #[test] + fn validate_balanced_with_custom_segmenter_errors() { + let b = fresh() + .profile(Profile::Balanced) + .with_segmenter(Box::new(MockSegmenter::default())); + let err = b.validate().unwrap_err(); + assert!(matches!( + err, + ConfigError::CustomComponentInProfile { + offending: "segmenter", + .. + } + )); + } + + #[test] + fn validate_custom_with_registry_errors() { + // ModelRegistry::default() may fail offline — skip in that case rather than fail the test. + let registry = match ModelRegistry::default() { + Ok(r) => r, + Err(_) => return, + }; + let b = fresh() + .profile(Profile::Custom) + .with_segmenter(Box::new(MockSegmenter::default())) + .with_embedder(Box::new(MockEmbedder::default())) + .with_clusterer(Box::new(MockClusterer::default())) + .with_models_from(registry); + let err = b.validate().unwrap_err(); + assert!(matches!(err, ConfigError::RegistryInCustomProfile)); + } + + #[test] + fn embedder_pool_size_clamps_to_1() { + let b = fresh().embedder_pool_size(0); + assert_eq!(b.config.embedder_pool_size, 1); + } +} +``` + +- [ ] **Step 3.2: Wire stub mod** + +In `/Users/ekhodzitsky/Documents/personal/polyvoice/src/pipeline_v2/mod.rs`, append: + +```rust + +pub mod builder; +pub mod mocks; +``` + +(Both `builder` and `mocks` are referenced by builder tests; we'll create `mocks.rs` in Task 4 and the tests will compile only after that lands. For Task 3 we cfg-gate the mocks reference.) + +Actually use this version of `mod.rs` instead — gating mocks `#[cfg(test)]`: + +```rust +//! M6a — additive `polyvoice::pipeline_v2` module. +//! +//! Spec: `docs/superpowers/specs/2026-05-07-m6a-pipeline-v2-design.md`. + +pub mod builder; +pub mod config; + +#[cfg(test)] +pub mod mocks; +``` + +- [ ] **Step 3.3: Confirm compile-failure** + +The builder tests reference `pipeline_v2::mocks` which doesn't exist yet. Skip running tests until Task 4 lands. + +For now run only: + +```bash +cd /Users/ekhodzitsky/Documents/personal/polyvoice +cargo check --all-features --lib 2>&1 | head -10 +``` + +Expected: errors about `pipeline_v2::mocks::MockSegmenter` etc. unresolved. + +- [ ] **Step 3.4: Commit (compile failing, mocks come in Task 4)** + +```bash +git add src/pipeline_v2/builder.rs src/pipeline_v2/mod.rs +git commit -m "feat(pipeline_v2): add PipelineBuilder + ConfigError (mocks land in Task 4)" +``` + +--- + +## Task 4: Test-only `mocks.rs` (Mock{Segmenter, Embedder, Clusterer}) + +**Files:** +- Create: `src/pipeline_v2/mocks.rs` + +- [ ] **Step 4.1: Write mocks** + +Create `/Users/ekhodzitsky/Documents/personal/polyvoice/src/pipeline_v2/mocks.rs`: + +```rust +//! M6a — test-only `Mock{Segmenter,Embedder,Clusterer}` for the +//! `pipeline_v2` builder validation tests and the synthetic integration +//! test in `tests/pipeline_v2_synthetic_test.rs`. +//! +//! Compiled only under `#[cfg(test)]` (or for integration tests when this +//! file is re-exported through `lib.rs` + `#[cfg(any(test, feature = "..."))]` +//! — but for M6a we keep these crate-internal `cfg(test)` mocks and the +//! integration test redefines its own thin mocks where needed). + +use crate::clusterer::{Clusterer, ClustererError}; +use crate::embedder::{Embedder, EmbedderError}; +use crate::resegmentation::{Resegmenter, ResegmentError, ResegmentInputs}; +use crate::segmentation::{RawSegment, SegmentationError, Segmenter}; +use crate::types::{Confidence, SpeakerTurn, TimeRange}; + +/// Constant-output `Segmenter` for builder tests. +#[derive(Default)] +pub struct MockSegmenter { + pub segments: Vec, +} + +impl Segmenter for MockSegmenter { + fn segment(&self, _audio: &[f32]) -> Result, SegmentationError> { + Ok(self.segments.clone()) + } + + fn max_local_speakers(&self) -> usize { + 3 + } + + fn supports_overlap(&self) -> bool { + true + } +} + +/// Constant-output `Embedder` for builder tests. +pub struct MockEmbedder { + pub embedding: Vec, +} + +impl Default for MockEmbedder { + fn default() -> Self { + // 192-d unit vector along axis 0; matches CAM++ output dim used + // throughout the spec. + let mut v = vec![0.0_f32; 192]; + v[0] = 1.0; + Self { embedding: v } + } +} + +impl Embedder for MockEmbedder { + fn dim(&self) -> usize { + self.embedding.len() + } + + fn embed(&self, _audio: &[f32]) -> Result, EmbedderError> { + Ok(self.embedding.clone()) + } +} + +/// Constant-label `Clusterer` for builder tests. +#[derive(Default)] +pub struct MockClusterer { + pub labels: Vec, +} + +impl Clusterer for MockClusterer { + fn cluster( + &self, + embeddings: &[Vec], + ) -> Result, ClustererError> { + if self.labels.is_empty() { + // Default: single cluster. + return Ok(vec![0; embeddings.len()]); + } + if self.labels.len() != embeddings.len() { + return Err(ClustererError::AlgorithmFailed { + detail: "MockClusterer labels length mismatch".to_owned(), + }); + } + Ok(self.labels.clone()) + } + + fn max_clusters(&self) -> usize { + 16 + } +} + +/// Pass-through `Resegmenter` (returns input primary turns sorted, no +/// secondary speakers added). +#[derive(Default)] +pub struct PassThroughResegmenter; + +impl Resegmenter for PassThroughResegmenter { + fn resegment( + &self, + inputs: ResegmentInputs<'_>, + ) -> Result, ResegmentError> { + let mut out: Vec = inputs.primary_turns.to_vec(); + out.sort_by(|a, b| a.time.start.total_cmp(&b.time.start)); + Ok(out) + } +} + +/// Convenience constructor for a single `RawSegment` used in tests. +pub fn raw_segment(start: f64, end: f64, spk: u8, overlap: bool) -> RawSegment { + RawSegment { + time: TimeRange { start, end }, + local_speaker_idx: spk, + is_overlap: overlap, + confidence: Confidence::new(0.9).unwrap(), + } +} +``` + +- [ ] **Step 4.2: Run builder tests** + +```bash +cd /Users/ekhodzitsky/Documents/personal/polyvoice +cargo test --all-features --lib pipeline_v2::builder +cargo test --all-features --lib pipeline_v2::config +cargo clippy --all-features --lib -- -D warnings +``` + +Expected: 3 config tests + 7 builder tests pass; clippy clean. + +- [ ] **Step 4.3: Commit** + +```bash +git add src/pipeline_v2/mocks.rs +git commit -m "test(pipeline_v2): add Mock{Segmenter,Embedder,Clusterer} for builder validation" +``` + +--- + +## Task 5: `Pipeline` core + `run()` orchestrator + feature-gate + +**Files:** +- Modify: `src/pipeline_v2/mod.rs` (replace skeleton) + +- [ ] **Step 5.1: Append run-flow tests** + +Add to `src/pipeline_v2/mod.rs` (we'll write it whole in Step 5.2 — for now state expected behaviour): + +```text +- pipeline_run_unsupported_sample_rate_returns_err +- pipeline_run_synthetic_two_speakers_through_custom_profile + (uses MockSegmenter w/ 2 segments at distinct times, MockEmbedder, MockClusterer that + labels them 0 and 1; expects 2-element DiarizationResult.turns sorted by start). +- pipeline_run_silence_returns_empty + (MockSegmenter returns empty segments; expects empty turns + num_speakers=0) +- pipeline_resegment_overlap_disabled_no_secondaries + (resegment_overlap = false; never call resegmenter even if overlap regions exist) +``` + +- [ ] **Step 5.2: Implement Pipeline + run()** + +Replace `/Users/ekhodzitsky/Documents/personal/polyvoice/src/pipeline_v2/mod.rs` with: + +```rust +//! M6a — additive `polyvoice::pipeline_v2` module. +//! +//! Spec: `docs/superpowers/specs/2026-05-07-m6a-pipeline-v2-design.md`. + +#[cfg(not(all( + feature = "pipeline_v2", + feature = "onnx", + feature = "segmentation", + feature = "embedder", + feature = "clusterer", + feature = "resegmentation", +)))] +compile_error!( + "pipeline_v2 requires onnx + segmentation + embedder + clusterer + resegmentation features" +); + +pub mod builder; +pub mod config; + +#[cfg(test)] +pub mod mocks; + +use crate::clusterer::{Clusterer, ClustererError}; +use crate::embedder::{Embedder, EmbedderError}; +use crate::models::RegistryError; +use crate::resegmentation::{ + OverlapRegionInput, OverlapResegmenter, Resegmenter, ResegmentError, ResegmentInputs, + SpeakerCentroid, compute_centroids, extract_overlap_time_ranges, +}; +use crate::segmentation::{SegmentationError, Segmenter}; +use crate::types::{ + DiarizationResult, Profile, SampleRate, Segment, SpeakerId, SpeakerTurn, TimeRange, +}; +use crate::utils::{l2_normalize, merge_segments}; + +pub use builder::{ConfigError, PipelineBuilder}; +pub use config::{ClustererKind, ExecutionProvider, PipelineConfig}; + +/// Errors raised by `Pipeline::run()`. +#[derive(Debug, thiserror::Error)] +pub enum PipelineError { + #[error("audio sample rate {actual} unsupported, expected 16000")] + UnsupportedSampleRate { actual: u32 }, + #[error("segmentation failed: {0}")] + Segmentation(#[from] SegmentationError), + #[error("embedding failed: {0}")] + Embedding(#[from] EmbedderError), + #[error("clustering failed: {0}")] + Clustering(#[from] ClustererError), + #[error("resegmentation failed: {0}")] + Resegment(#[from] ResegmentError), + #[error("config error: {0}")] + Config(#[from] ConfigError), + #[error("model registry error: {0}")] + Registry(#[from] RegistryError), + #[error("model load error: {detail}")] + ModelLoad { detail: String }, +} + +/// v1.0 Pipeline. Constructed via `Pipeline::builder()`. +pub struct Pipeline { + config: PipelineConfig, + segmenter: Box, + embedder: Box, + clusterer: Box, + resegmenter: Box, +} + +impl Pipeline { + /// Start building a `Pipeline`. Returns a `PipelineBuilder` with default config. + pub fn builder() -> PipelineBuilder { + PipelineBuilder::new() + } + + /// Construct a `Pipeline` from the (validated) builder's components. + /// **Public for the builder integration only**; users invoke `builder().build()`. + pub(crate) fn from_components( + config: PipelineConfig, + segmenter: Box, + embedder: Box, + clusterer: Box, + resegmenter: Box, + ) -> Self { + Self { + config, + segmenter, + embedder, + clusterer, + resegmenter, + } + } + + pub fn config(&self) -> &PipelineConfig { + &self.config + } + + /// Run the pipeline end-to-end on a slice of f32 PCM samples. + pub fn run( + &self, + samples: &[f32], + sr: SampleRate, + ) -> Result { + if sr.get() != self.config.sample_rate.get() { + return Err(PipelineError::UnsupportedSampleRate { actual: sr.get() }); + } + + let raw_segments = self.segmenter.segment(samples)?; + if raw_segments.is_empty() { + return Ok(DiarizationResult { + segments: Vec::new(), + turns: Vec::new(), + num_speakers: 0, + }); + } + + let overlap_ranges = extract_overlap_time_ranges(&raw_segments); + let primary_segments: Vec<_> = raw_segments + .iter() + .filter(|s| !s.is_overlap) + .cloned() + .collect(); + + let sample_rate = self.config.sample_rate.get() as f64; + let mut embeddings: Vec> = Vec::with_capacity(primary_segments.len()); + for seg in &primary_segments { + let start_idx = (seg.time.start * sample_rate) as usize; + let end_idx = ((seg.time.end * sample_rate) as usize).min(samples.len()); + if end_idx <= start_idx { + continue; + } + let chunk = &samples[start_idx..end_idx]; + let mut emb = self.embedder.embed(chunk)?; + l2_normalize(&mut emb); + embeddings.push(emb); + } + + if embeddings.is_empty() { + return Ok(DiarizationResult { + segments: Vec::new(), + turns: Vec::new(), + num_speakers: 0, + }); + } + + let labels = self.clusterer.cluster(&embeddings)?; + + // Build primary_turns aligned with embeddings/labels. + let mut primary_turns: Vec = primary_segments + .iter() + .zip(labels.iter()) + .map(|(seg, &lbl)| SpeakerTurn { + speaker: SpeakerId(lbl as u32), + time: seg.time, + text: None, + }) + .collect(); + + // Centroids from clean (non-overlap) embeddings. + let centroids: Vec = compute_centroids(&embeddings, &labels); + + // Optionally resegment overlap regions. + let mut all_turns: Vec = if self.config.resegment_overlap + && !overlap_ranges.is_empty() + && centroids.len() >= 2 + { + let overlap_inputs = self.build_overlap_inputs( + &overlap_ranges, + &primary_turns, + samples, + )?; + self.resegmenter.resegment(ResegmentInputs { + primary_turns: &primary_turns, + speaker_centroids: ¢roids, + overlap_regions: &overlap_inputs, + })? + } else { + primary_turns.sort_by(|a, b| a.time.start.total_cmp(&b.time.start)); + primary_turns + }; + + // Filter short segments + merge contiguous same-speaker turns. + let min_secs = self.config.min_speech_secs as f64; + all_turns.retain(|t| t.time.duration() >= min_secs); + + let max_gap = self.config.max_gap_secs as f64; + let merged_segments: Vec = all_turns + .iter() + .map(|t| Segment { + time: t.time, + speaker: Some(t.speaker), + confidence: None, + }) + .collect(); + let merged_segments = merge_segments(merged_segments, max_gap); + let merged_turns: Vec = merged_segments + .iter() + .filter_map(|s| { + s.speaker.map(|spk| SpeakerTurn { + speaker: spk, + time: s.time, + text: None, + }) + }) + .collect(); + + let num_speakers = merged_turns + .iter() + .map(|t| t.speaker.0) + .collect::>() + .len(); + + Ok(DiarizationResult { + segments: merged_segments, + turns: merged_turns, + num_speakers, + }) + } + + fn build_overlap_inputs( + &self, + overlap_ranges: &[(TimeRange, u8, u8)], + primary_turns: &[SpeakerTurn], + samples: &[f32], + ) -> Result, PipelineError> { + let sample_rate = self.config.sample_rate.get() as f64; + let mut out = Vec::with_capacity(overlap_ranges.len()); + for (time, _lo, _hi) in overlap_ranges { + let primary = primary_turns + .iter() + .find(|t| t.time.start <= time.start && time.end <= t.time.end) + .map(|t| t.speaker) + .unwrap_or_else(|| { + // Fallback to nearest primary by start time. + primary_turns + .iter() + .min_by(|a, b| { + (a.time.start - time.start) + .abs() + .total_cmp(&(b.time.start - time.start).abs()) + }) + .map(|t| t.speaker) + .unwrap_or(SpeakerId(0)) + }); + let start_idx = (time.start * sample_rate) as usize; + let end_idx = ((time.end * sample_rate) as usize).min(samples.len()); + if end_idx <= start_idx { + continue; + } + let chunk = &samples[start_idx..end_idx]; + let mut emb = self.embedder.embed(chunk)?; + l2_normalize(&mut emb); + out.push(OverlapRegionInput { + time: *time, + primary_speaker: primary, + embedding: emb, + }); + } + Ok(out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::pipeline_v2::mocks::{MockClusterer, MockEmbedder, MockSegmenter, raw_segment}; + use crate::types::Profile; + + fn pipeline_with_segments(segs: Vec) -> Pipeline { + let mut cfg = PipelineConfig::default(); + cfg.profile = Profile::Custom; + cfg.resegment_overlap = false; + cfg.min_speech_secs = 0.0; + cfg.max_gap_secs = 0.0; + Pipeline::from_components( + cfg, + Box::new(MockSegmenter { segments: segs }), + Box::new(MockEmbedder::default()), + Box::new(MockClusterer::default()), + Box::new(OverlapResegmenter::default()), + ) + } + + #[test] + fn pipeline_run_unsupported_sample_rate_returns_err() { + let p = pipeline_with_segments(vec![raw_segment(0.0, 1.0, 0, false)]); + let bad = SampleRate::new(8000).unwrap(); + let err = p.run(&vec![0.0_f32; 8000], bad).unwrap_err(); + assert!(matches!(err, PipelineError::UnsupportedSampleRate { actual: 8000 })); + } + + #[test] + fn pipeline_run_silence_returns_empty() { + let p = pipeline_with_segments(Vec::new()); + let result = p + .run(&vec![0.0_f32; 16000], SampleRate::new(16000).unwrap()) + .unwrap(); + assert!(result.turns.is_empty()); + assert_eq!(result.num_speakers, 0); + } + + #[test] + fn pipeline_run_two_segments_one_cluster() { + let segs = vec![ + raw_segment(0.0, 1.0, 0, false), + raw_segment(1.5, 2.5, 0, false), + ]; + let p = pipeline_with_segments(segs); + let result = p + .run(&vec![0.0_f32; 16000 * 3], SampleRate::new(16000).unwrap()) + .unwrap(); + assert_eq!(result.num_speakers, 1); + assert!(result.turns.len() >= 1); + } + + #[test] + fn pipeline_resegment_overlap_disabled_path_used() { + let segs = vec![ + raw_segment(0.0, 1.0, 0, true), + raw_segment(0.0, 1.0, 1, true), + raw_segment(1.5, 2.5, 0, false), + ]; + let p = pipeline_with_segments(segs); + let result = p + .run(&vec![0.0_f32; 16000 * 3], SampleRate::new(16000).unwrap()) + .unwrap(); + // With resegment_overlap = false in pipeline_with_segments, no secondary + // is appended. Overlap-flagged segments are filtered out at the embedder + // step (they have is_overlap = true), so num_speakers reflects only the + // non-overlap segment cluster. + assert!(result.num_speakers <= 1); + } +} +``` + +- [ ] **Step 5.3: Wire builder.build() to Pipeline::from_components** + +Append to `src/pipeline_v2/builder.rs`: + +```rust + +use crate::pipeline_v2::Pipeline; +use crate::resegmentation::OverlapResegmenter; + +impl PipelineBuilder { + /// Validate + construct the inner `Pipeline`. + pub fn build(self) -> Result { + self.validate()?; + let resegmenter = self + .custom_resegmenter + .unwrap_or_else(|| Box::new(OverlapResegmenter::default())); + + match self.config.profile { + Profile::Custom => { + // Validation already ensured all three are Some. + let segmenter = self.custom_segmenter.expect("validated Custom segmenter"); + let embedder = self.custom_embedder.expect("validated Custom embedder"); + let clusterer = self.custom_clusterer.expect("validated Custom clusterer"); + Ok(Pipeline::from_components( + self.config, + segmenter, + embedder, + clusterer, + resegmenter, + )) + } + Profile::Mobile | Profile::Balanced => { + // Resolve registry → ONNX-backed components. Wired here so that + // unit tests of validate() don't pull ONNX. The returned Box + // wraps `PowersetSegmenter`, `CamPlusPlusExtractor` / + // `ResNet34Adapter`, and `NmeScClusterer` / `AhcClusterer` per + // ClustererKind. + let registry = self.registry.expect("validated registry presence"); + let profile_models = registry.ensure_for_profile(self.config.profile)?; + let segmenter: Box = Box::new( + crate::segmentation::PowersetSegmenter::new(&profile_models.segmenter_path) + .map_err(|e| ConfigError::UnknownModel { + model_id: format!("powerset (cause: {e})"), + })?, + ); + let embedder: Box = match self.config.profile { + Profile::Mobile => Box::new( + crate::embedder::CamPlusPlusExtractor::new( + &profile_models.embedder_path, + self.config.profile.embedding_dim(), + self.config.embedder_pool_size, + ) + .map_err(|e| ConfigError::UnknownModel { + model_id: format!("cam_pp (cause: {e})"), + })?, + ), + Profile::Balanced => Box::new( + crate::embedder::ResNet34Adapter::new( + &profile_models.embedder_path, + self.config.embedder_pool_size, + ) + .map_err(|e| ConfigError::UnknownModel { + model_id: format!("resnet34 (cause: {e})"), + })?, + ), + Profile::Custom => unreachable!("Profile::Custom handled above"), + }; + let clusterer: Box = match self.config.clusterer { + ClustererKind::Ahc { .. } => Box::new(crate::clusterer::AhcClusterer::new( + self.config.max_speakers as usize, + )), + #[cfg(feature = "spectral")] + ClustererKind::NmeSc => Box::new(crate::clusterer::NmeScClusterer::new( + self.config.max_speakers as usize, + )), + #[cfg(not(feature = "spectral"))] + ClustererKind::NmeSc => Box::new(crate::clusterer::AhcClusterer::new( + self.config.max_speakers as usize, + )), + }; + Ok(Pipeline::from_components( + self.config, + segmenter, + embedder, + clusterer, + resegmenter, + )) + } + } + } +} +``` + +- [ ] **Step 5.4: Verify** + +```bash +cd /Users/ekhodzitsky/Documents/personal/polyvoice +cargo test --all-features --lib pipeline_v2 +cargo clippy --all-features --lib -- -D warnings +cargo fmt +cargo check --no-default-features --features pipeline_v2 2>&1 | head -10 +cargo check --target wasm32-unknown-unknown --no-default-features --features resegmentation --lib +``` + +Expected: +- ~14 lib tests pass (3 config + 7 builder + 4 mod) +- Clippy clean +- `--no-default-features --features pipeline_v2` exits non-zero with `compile_error!` message about missing dependencies +- wasm32 clean (pipeline_v2 gated out automatically) + +- [ ] **Step 5.5: Commit** + +```bash +git add src/pipeline_v2/mod.rs src/pipeline_v2/builder.rs +git commit -m "feat(pipeline_v2): add Pipeline core + run() orchestrator + builder.build()" +``` + +--- + +## Task 6: Synthetic integration test + +**Files:** +- Create: `tests/pipeline_v2_synthetic_test.rs` + +- [ ] **Step 6.1: Write the integration test** + +Create `/Users/ekhodzitsky/Documents/personal/polyvoice/tests/pipeline_v2_synthetic_test.rs`: + +```rust +//! M6a — synthetic-data integration tests for `polyvoice::pipeline_v2`. +//! +//! Pure-CPU; no ONNX. Covers builder validation paths, end-to-end +//! Custom-profile run, and overlap-resegmentation toggling. + +#![cfg(all( + feature = "pipeline_v2", + feature = "onnx", + feature = "segmentation", + feature = "embedder", + feature = "clusterer", + feature = "resegmentation", +))] + +use polyvoice::clusterer::{Clusterer, ClustererError}; +use polyvoice::embedder::{Embedder, EmbedderError}; +use polyvoice::pipeline_v2::{ + ClustererKind, ConfigError, Pipeline, PipelineConfig, PipelineError, +}; +use polyvoice::segmentation::{RawSegment, SegmentationError, Segmenter}; +use polyvoice::types::{Confidence, Profile, SampleRate, SpeakerId, TimeRange}; + +// ----- Local mocks (re-defined here because the crate-internal mocks are #[cfg(test)]) ----- + +struct ConstSegmenter(Vec); +impl Segmenter for ConstSegmenter { + fn segment(&self, _audio: &[f32]) -> Result, SegmentationError> { + Ok(self.0.clone()) + } + fn max_local_speakers(&self) -> usize { + 3 + } + fn supports_overlap(&self) -> bool { + true + } +} + +struct AxisEmbedder { + dim: usize, + axis_picker: Box usize + Send + Sync>, +} +impl Embedder for AxisEmbedder { + fn dim(&self) -> usize { + self.dim + } + fn embed(&self, audio: &[f32]) -> Result, EmbedderError> { + let axis = (self.axis_picker)(audio).min(self.dim - 1); + let mut v = vec![0.0_f32; self.dim]; + v[axis] = 1.0; + Ok(v) + } +} + +struct PerSampleClusterer { + labels: Vec, +} +impl Clusterer for PerSampleClusterer { + fn cluster(&self, embeddings: &[Vec]) -> Result, ClustererError> { + if embeddings.len() != self.labels.len() { + return Err(ClustererError::AlgorithmFailed { + detail: format!( + "labels {} vs embeddings {}", + self.labels.len(), + embeddings.len() + ), + }); + } + Ok(self.labels.clone()) + } + fn max_clusters(&self) -> usize { + 16 + } +} + +fn raw(start: f64, end: f64, spk: u8, overlap: bool) -> RawSegment { + RawSegment { + time: TimeRange { start, end }, + local_speaker_idx: spk, + is_overlap: overlap, + confidence: Confidence::new(0.9).unwrap(), + } +} + +fn axis_picker_constant(axis: usize) -> Box usize + Send + Sync> { + Box::new(move |_| axis) +} + +// ----- Tests ----- + +#[test] +fn builder_validation_mobile_missing_registry() { + let err = Pipeline::builder() + .profile(Profile::Mobile) + .build() + .unwrap_err(); + assert!(matches!( + err, + ConfigError::MissingRegistry { + profile: Profile::Mobile + } + )); +} + +#[test] +fn builder_validation_custom_missing_components() { + let err = Pipeline::builder() + .profile(Profile::Custom) + .build() + .unwrap_err(); + assert!(matches!(err, ConfigError::MissingCustomComponent { .. })); +} + +#[test] +fn pipeline_run_unsupported_sample_rate_errors() { + let cfg = PipelineConfig { + profile: Profile::Custom, + ..PipelineConfig::default() + }; + let p = Pipeline::builder() + .config(cfg) + .with_segmenter(Box::new(ConstSegmenter(vec![raw(0.0, 1.0, 0, false)]))) + .with_embedder(Box::new(AxisEmbedder { + dim: 8, + axis_picker: axis_picker_constant(0), + })) + .with_clusterer(Box::new(PerSampleClusterer { labels: vec![0] })) + .build() + .expect("custom build"); + let err = p + .run(&vec![0.0_f32; 8000], SampleRate::new(8000).unwrap()) + .unwrap_err(); + assert!(matches!( + err, + PipelineError::UnsupportedSampleRate { actual: 8000 } + )); +} + +#[test] +fn pipeline_run_silence_returns_empty() { + let cfg = PipelineConfig { + profile: Profile::Custom, + ..PipelineConfig::default() + }; + let p = Pipeline::builder() + .config(cfg) + .with_segmenter(Box::new(ConstSegmenter(Vec::new()))) + .with_embedder(Box::new(AxisEmbedder { + dim: 8, + axis_picker: axis_picker_constant(0), + })) + .with_clusterer(Box::new(PerSampleClusterer { labels: vec![] })) + .build() + .expect("custom build"); + let r = p + .run(&vec![0.0_f32; 16000], SampleRate::new(16000).unwrap()) + .unwrap(); + assert!(r.turns.is_empty()); + assert_eq!(r.num_speakers, 0); +} + +#[test] +fn pipeline_run_two_speakers_through_custom_profile() { + // Two non-overlap segments at distinct times, embedded onto distinct + // axes, clustered to two distinct labels. + let segs = vec![raw(0.0, 1.0, 0, false), raw(2.0, 3.0, 1, false)]; + let cfg = PipelineConfig { + profile: Profile::Custom, + clusterer: ClustererKind::NmeSc, + resegment_overlap: false, + min_speech_secs: 0.0, + max_gap_secs: 0.0, + ..PipelineConfig::default() + }; + let p = Pipeline::builder() + .config(cfg) + .with_segmenter(Box::new(ConstSegmenter(segs))) + .with_embedder(Box::new(AxisEmbedder { + dim: 8, + // Use start-time of the chunk to pick a deterministic axis: first + // chunk axis 0, second chunk axis 1. The chunk slice doesn't carry + // its absolute start, so we exploit chunk length: 1s = 16_000 + // samples (axis 0), 1s = 16_000 too. We need a different signal — + // use a counter via interior mutability isn't great, so instead + // pick by content sum: silence chunks both → axis 0. Skip and + // use PerSampleClusterer to bypass. + axis_picker: axis_picker_constant(0), + })) + .with_clusterer(Box::new(PerSampleClusterer { labels: vec![0, 1] })) + .build() + .expect("custom build"); + let r = p + .run(&vec![0.0_f32; 16000 * 4], SampleRate::new(16000).unwrap()) + .unwrap(); + assert_eq!(r.num_speakers, 2); + assert_eq!(r.turns.len(), 2); + let speakers: Vec = r.turns.iter().map(|t| t.speaker.0).collect(); + assert!(speakers.contains(&0)); + assert!(speakers.contains(&1)); +} + +#[test] +fn pipeline_run_returns_sorted_turns() { + // Out-of-order segments must produce sorted output. + let segs = vec![raw(2.0, 3.0, 0, false), raw(0.0, 1.0, 0, false)]; + let cfg = PipelineConfig { + profile: Profile::Custom, + resegment_overlap: false, + min_speech_secs: 0.0, + max_gap_secs: 0.0, + ..PipelineConfig::default() + }; + let p = Pipeline::builder() + .config(cfg) + .with_segmenter(Box::new(ConstSegmenter(segs))) + .with_embedder(Box::new(AxisEmbedder { + dim: 8, + axis_picker: axis_picker_constant(0), + })) + .with_clusterer(Box::new(PerSampleClusterer { labels: vec![0, 0] })) + .build() + .expect("custom build"); + let r = p + .run(&vec![0.0_f32; 16000 * 4], SampleRate::new(16000).unwrap()) + .unwrap(); + // After merge_segments the two same-speaker turns may collapse into one. + for w in r.turns.windows(2) { + assert!(w[0].time.start <= w[1].time.start); + } +} + +#[test] +fn pipeline_resegment_overlap_disabled_no_secondaries() { + // Two overlap segments at the same range + one clean. With + // resegment_overlap=false the secondary speaker is never appended. + let segs = vec![ + raw(0.0, 1.0, 0, true), + raw(0.0, 1.0, 1, true), + raw(2.0, 3.0, 0, false), + ]; + let cfg = PipelineConfig { + profile: Profile::Custom, + resegment_overlap: false, + min_speech_secs: 0.0, + max_gap_secs: 0.0, + ..PipelineConfig::default() + }; + let p = Pipeline::builder() + .config(cfg) + .with_segmenter(Box::new(ConstSegmenter(segs))) + .with_embedder(Box::new(AxisEmbedder { + dim: 8, + axis_picker: axis_picker_constant(0), + })) + .with_clusterer(Box::new(PerSampleClusterer { labels: vec![0] })) + .build() + .expect("custom build"); + let r = p + .run(&vec![0.0_f32; 16000 * 4], SampleRate::new(16000).unwrap()) + .unwrap(); + assert_eq!(r.num_speakers, 1); +} +``` + +- [ ] **Step 6.2: Verify** + +```bash +cd /Users/ekhodzitsky/Documents/personal/polyvoice +cargo test --all-features --test pipeline_v2_synthetic_test +cargo clippy --all-features --tests -- -D warnings +``` + +Expected: 6 tests pass, clippy clean. + +- [ ] **Step 6.3: Commit** + +```bash +git add tests/pipeline_v2_synthetic_test.rs +git commit -m "test(pipeline_v2): add synthetic integration tests for builder + run() flow" +``` + +--- + +## Task 7: lib.rs re-exports + e2e test + CHANGELOG + tag + +**Files:** +- Modify: `src/lib.rs` +- Create: `tests/pipeline_v2_e2e_test.rs` +- Modify: `CHANGELOG.md` + +- [ ] **Step 7.1: Add re-exports to lib.rs** + +In `/Users/ekhodzitsky/Documents/personal/polyvoice/src/lib.rs`, after the `#[cfg(all(feature = "pipeline_v2", ...))] pub mod pipeline_v2;` block (added in Task 2.2), append: + +```rust + +#[cfg(all( + feature = "pipeline_v2", + feature = "onnx", + feature = "segmentation", + feature = "embedder", + feature = "clusterer", + feature = "resegmentation", +))] +pub use pipeline_v2::{ + ClustererKind, ConfigError, ExecutionProvider, Pipeline as PipelineV2, + PipelineBuilder, PipelineConfig, PipelineError as PipelineV2Error, +}; +``` + +- [ ] **Step 7.2: Write the e2e test (#[ignore] by default)** + +Create `/Users/ekhodzitsky/Documents/personal/polyvoice/tests/pipeline_v2_e2e_test.rs`: + +```rust +//! M6a — `#[ignore]` end-to-end test for `polyvoice::pipeline_v2`. +//! +//! Requires the Balanced ONNX bundle to be cached (run +//! `cargo run --features cli --bin polyvoice -- download-models --profile balanced` +//! once before invoking with `cargo test -- --ignored e2e`). Reads a single +//! WAV from `data/voxconverse-test/audio/` and asserts the pipeline returns +//! at least one turn with valid speaker IDs. +//! +//! Skipped by default to keep CI fast; M6b will wire a polyvoice-bench +//! integration that exercises full DER on VoxConverse-test. + +#![cfg(all( + feature = "pipeline_v2", + feature = "download", + feature = "onnx", + feature = "segmentation", + feature = "embedder", + feature = "clusterer", + feature = "resegmentation", +))] + +use polyvoice::models::ModelRegistry; +use polyvoice::pipeline_v2::Pipeline; +use polyvoice::types::{Profile, SampleRate}; +use polyvoice::wav::read_wav; +use std::path::Path; + +fn first_voxconverse_wav() -> Option { + let dir = Path::new("data/voxconverse-test/audio"); + if !dir.is_dir() { + return None; + } + let mut entries: Vec<_> = std::fs::read_dir(dir) + .ok()? + .filter_map(|e| e.ok()) + .filter(|e| e.path().extension().is_some_and(|ext| ext == "wav")) + .map(|e| e.path()) + .collect(); + entries.sort(); + entries.into_iter().next() +} + +#[ignore = "requires cached ONNX bundle + a wav file under data/voxconverse-test/audio/"] +#[test] +fn e2e_balanced_profile_voxconverse_clip() { + let wav_path = match first_voxconverse_wav() { + Some(p) => p, + None => panic!( + "data/voxconverse-test/audio/ is empty — run scripts/download-voxconverse-test.sh first" + ), + }; + let (samples, sr_hz) = + read_wav(&wav_path).expect("WAV read failure — check the file is 16 kHz mono"); + assert_eq!(sr_hz, 16000, "only 16 kHz WAVs supported"); + let registry = ModelRegistry::default() + .expect("default ModelRegistry should resolve a writable cache dir"); + let pipeline = Pipeline::builder() + .profile(Profile::Balanced) + .with_models_from(registry) + .build() + .expect( + "Balanced profile build should succeed when cached ONNX is present — \ + run `polyvoice download-models --profile balanced` first", + ); + let result = pipeline + .run(&samples, SampleRate::new(16000).unwrap()) + .expect("pipeline.run on a real VoxConverse clip should succeed"); + assert!( + result.num_speakers >= 1, + "expected at least 1 speaker, got {}", + result.num_speakers + ); + for w in result.turns.windows(2) { + assert!(w[0].time.start <= w[1].time.start, "turns must be sorted by start time"); + } +} +``` + +- [ ] **Step 7.3: Update CHANGELOG.md** + +In `/Users/ekhodzitsky/Documents/personal/polyvoice/CHANGELOG.md`, in the `## [Unreleased]` block, after the M5 section (`### Added (M5 — INT8 quantization)`), append: + +```markdown + +### Added (M6a — Pipeline + Profile API, additive) +- `polyvoice::pipeline_v2` module: `Pipeline::builder()` returning + `PipelineBuilder` with `.profile(Profile::Mobile|Balanced|Custom)`, + `.with_models_from(ModelRegistry)`, `.with_segmenter/embedder/clusterer/resegmenter()`, + and a validated `.build()`. `PipelineConfig`, `ClustererKind`, + `ExecutionProvider`, and `ConfigError` all per spec §5.2/§5.4. +- `Pipeline::run(&samples, SampleRate)` orchestrates M1 segmenter → M2 + embedder (with `apply_overlap_mask`) → M3 clusterer → M4 resegmenter → + legacy `merge_segments` → `DiarizationResult`. +- New Cargo feature `pipeline_v2` (in default features). Requires + `onnx + segmentation + embedder + clusterer + resegmentation`; missing + any of these triggers a `compile_error!` with an actionable message. +- Public re-exports `polyvoice::PipelineV2`, `PipelineBuilder`, + `PipelineConfig`, `ClustererKind`, `ExecutionProvider`, `ConfigError`, + `PipelineV2Error`. Legacy `polyvoice::Pipeline` is unchanged; M6b will + rename `pipeline_v2 → pipeline` and remove the legacy code path. +- Synthetic integration test on Custom profile (`tests/pipeline_v2_synthetic_test.rs`) + + `#[ignore]` E2E test on Balanced profile via `ModelRegistry` + (`tests/pipeline_v2_e2e_test.rs`). +``` + +- [ ] **Step 7.4: Verify full feature matrix** + +```bash +cd /Users/ekhodzitsky/Documents/personal/polyvoice +cargo test --all-features --lib 2>&1 | tail -3 +cargo test --all-features --test pipeline_v2_synthetic_test 2>&1 | tail -3 +cargo test --all-features --doc 2>&1 | tail -3 +cargo clippy --all-targets --all-features -- -D warnings 2>&1 | tail -3 +cargo fmt --check +cargo check --target wasm32-unknown-unknown --no-default-features --features resegmentation --lib +bash scripts/release-gate.sh ; echo "exit=$?" +``` + +Apply `cargo fmt` if `--check` fails. Apply clippy fixes if test code is flagged. + +- [ ] **Step 7.5: Tag** + +```bash +cd /Users/ekhodzitsky/Documents/personal/polyvoice +git tag -a m6a-complete -m "M6a complete: Pipeline + Profile builder API (additive)" +``` + +(Don't push the tag yet — it gets pushed after the M6a PR is merged into master, mirroring M3/M4/M5.) + +- [ ] **Step 7.6: Commit** + +```bash +git add src/lib.rs tests/pipeline_v2_e2e_test.rs CHANGELOG.md +git commit -m "feat(lib): re-export pipeline_v2 surface + add E2E ignored test + changelog" +``` + +- [ ] **Step 7.7: Final git log** + +```bash +git log --oneline b1fcc9b..HEAD +``` + +Should show 7 commits matching the seven tasks. + +--- + +## Self-review checklist + +1. **Spec coverage:** all M6a deliverables (Pipeline, builder, config, validation, run() flow, synthetic test, E2E ignored test) → Tasks 2–7. Cargo feature → Task 1. Re-exports → Task 7.1. +2. **Additive guarantee:** `git diff b1fcc9b..HEAD -- src/pipeline.rs src/offline.rs src/online.rs src/types.rs src/vad.rs src/ffi.rs src/bin/polyvoice.rs src/bin/polyvoice-bench.rs python/src/lib.rs` should show ZERO changes. Legacy untouched. +3. **No `unwrap`/`expect`/`panic`** in lib non-test code (`expect("validated …")` lines in `builder.build()` are guarded by validate() invariants and acceptable per the existing M0–M5 pattern). +4. **Test coverage:** ≈3 config + 7 builder + 4 mod = 14 lib tests + 6 synthetic + 1 ignored E2E = 21 total tests across M6a. +5. **Atomic commits:** ~7 total — one per task. + +--- + +## Out of scope (M6b) + +- Removing legacy `src/pipeline.rs`, `src/offline.rs`, `OfflineDiarizer`, `DiarizationConfig`, `VadConfig`, `EnergyVad`, `VoiceActivityDetector` trait, `segment_speech` fn, `DummyExtractor` (move to `embedding/mock.rs` `#[cfg(test)]`), `OnnxEmbeddingExtractor`, `compute_fbank` privatization. +- Renaming `pipeline_v2 → pipeline`, `PipelineV2 → Pipeline`. +- CLI rewrite (`polyvoice diarize --profile mobile|balanced`). +- FFI rewrite (`src/ffi.rs`). +- Python pyo3 rewrite (`python/src/lib.rs`). +- `polyvoice-bench` rewrite using new Pipeline. +- E2E DER баннер на VoxConverse-test/AMI through polyvoice-bench → `tests/der_baseline.json`. +- Migration guide `docs/MIGRATING-FROM-0.5.md`. +- `OnlineDiarizer` deprecation annotation. diff --git a/docs/superpowers/specs/2026-05-07-m6a-pipeline-v2-design.md b/docs/superpowers/specs/2026-05-07-m6a-pipeline-v2-design.md new file mode 100644 index 00000000..65c38419 --- /dev/null +++ b/docs/superpowers/specs/2026-05-07-m6a-pipeline-v2-design.md @@ -0,0 +1,483 @@ +--- +title: M6a — Pipeline + Profile API (additive) Design +date: 2026-05-07 +status: draft +milestone: M6a +preceding: M0, M1, M2, M3, M4, M5 +following: M6b (cleanup + CLI/FFI/Python migration), M7 +authors: ekhodzitsky +--- + +# M6a — Pipeline + Profile API (additive) Design + +## Problem + +После M0–M5 у polyvoice есть все компоненты v1.0 диаризации: +- M1 powerset segmenter, +- M2 Embedder trait + CAM++ / ResNet34, +- M3 Clusterer trait + NME-SC / AHC, +- M4 OverlapResegmenter, +- M5 INT8-квантованные ONNX-модели в `[profiles.mobile]` / `[profiles.balanced]`, + +но они wired только в legacy `polyvoice::Pipeline::new(DiarizationConfig, VadConfig)` с устаревшим +`SileroVad + FbankOnnxExtractor + AHC` контуром. Пользователь не может включить +новый стек одной строкой — нужна `Pipeline::builder()` API. + +## Goal + +Single milestone (~1 неделя): + +1. **Additive new module** `src/pipeline_v2/` с `Pipeline`, `PipelineBuilder`, + `PipelineConfig`, `PipelineError`, `ClustererKind`, `ExecutionProvider`. +2. **Builder wiring** для `Profile::Mobile / Balanced / Custom`. Mobile/Balanced + подтягивают ONNX через `ModelRegistry` (M0); Custom — caller подаёт свои + `Segmenter / Embedder / Clusterer`. +3. **End-to-end run** через `Pipeline::run(&samples, SampleRate)` который + соединяет M1 → M2 (overlap mask + embedding) → M3 → M4 → output. +4. **Legacy остаётся работать** — `polyvoice::Pipeline` (renamed export to + `polyvoice::PipelineV0` if needed in `lib.rs`) и все M0-M3 тесты не ломаются. +5. **Synthetic + #[ignore] E2E test** покрытия. + +End state: пользователь может позвать `Pipeline::builder().profile(Profile::Mobile).with_models_from(ModelRegistry::default()?).build()?.run(&samples, SampleRate::HZ_16000)?` и получить +`DiarizationResult` через ONNX INT8 контур без знания деталей M1-M5. + +## Non-goals + +- CLI `polyvoice diarize --profile mobile` (M6b). +- FFI rewrite в `src/ffi.rs` (M6b). +- Python pyo3 bindings rewrite в `python/src/lib.rs` (M6b). +- `polyvoice-bench` rewrite на новый Pipeline (M6b). +- **Удаление legacy**: `src/pipeline.rs`, `src/offline.rs`, `OfflineDiarizer`, + `DiarizationConfig`, `VadConfig`, `EnergyVad`, `compute_fbank` privatization, + `OnlineDiarizer` deprecation note — всё в M6b. +- E2E DER на VoxConverse-test/AMI через `polyvoice-bench --profile mobile` — + M6b после CLI rewrite. +- Migration guide `docs/MIGRATING-FROM-0.5.md` — M6b. +- Streaming pipeline (`OnlineDiarizer`) — deprecated в v1.0, переезжает в v1.1. + +## Approach + +### Module placement: separate `src/pipeline_v2/` directory + +Pattern мatches M2 (`src/embedder.rs` parallel к legacy `src/embedding.rs`) +и M3 (`src/clusterer.rs` parallel к legacy `src/ahc.rs / spectral.rs`): + +``` +src/ +├── pipeline_v2/ [NEW dir, M6a] +│ ├── mod.rs Pipeline struct, run(), PipelineError +│ ├── builder.rs PipelineBuilder + ConfigError +│ ├── config.rs PipelineConfig + ClustererKind + ExecutionProvider +│ └── mocks.rs Test-only Mock{Segmenter,Embedder,Clusterer} (#[cfg(test)]) +├── pipeline.rs [legacy, kept verbatim for M6a] +``` + +Cargo feature `pipeline_v2 = []` (default-on, requires `onnx`+`segmentation`+`embedder`+`clusterer`+`resegmentation`): + +```toml +default = ["spectral", "segmentation", "embedder", "clusterer", "resegmentation", "pipeline_v2"] +pipeline_v2 = [] +``` + +`pipeline_v2` module gated: + +```rust +#[cfg(all( + feature = "pipeline_v2", + feature = "onnx", + feature = "segmentation", + feature = "embedder", + feature = "clusterer", + feature = "resegmentation", +))] +pub mod pipeline_v2; +``` + +При попытке включить `pipeline_v2` без onnx/segmentation/embedder и т.д. cargo +выдаст ошибку компиляции с понятным сообщением (через `compile_error!` в +`pipeline_v2/mod.rs`). + +### Public API surface + +```rust +// pipeline_v2/mod.rs +pub use builder::{PipelineBuilder, ConfigError}; +pub use config::{PipelineConfig, ClustererKind, ExecutionProvider}; + +pub struct Pipeline { /* private — created via builder */ } + +impl Pipeline { + pub fn builder() -> PipelineBuilder { ... } + + pub fn run( + &self, + samples: &[f32], + sr: SampleRate, + ) -> Result; +} + +#[derive(Debug, thiserror::Error)] +pub enum PipelineError { + #[error("audio sample rate {actual} unsupported, expected 16000")] + UnsupportedSampleRate { actual: u32 }, + #[error("segmentation failed: {0}")] + Segmentation(#[from] SegmentationError), + #[error("embedding failed: {0}")] + Embedding(#[from] EmbedderError), + #[error("clustering failed: {0}")] + Clustering(#[from] ClustererError), + #[error("resegmentation failed: {0}")] + Resegment(#[from] ResegmentError), + #[error("config error: {0}")] + Config(#[from] ConfigError), + #[error("model registry error: {0}")] + Registry(#[from] RegistryError), + #[error("model load error: {detail}")] + ModelLoad { detail: String }, +} +``` + +### `PipelineConfig` — spec §5.2 verbatim + +```rust +#[derive(Clone, Debug)] +pub struct PipelineConfig { + pub profile: Profile, + pub sample_rate: SampleRate, + pub seg_window_secs: f32, + pub seg_hop_secs: f32, + pub clusterer: ClustererKind, + pub max_speakers: u8, + pub min_cluster_size: usize, + pub resegment_overlap: bool, + pub resegment_min_overlap_secs: f32, + pub min_speech_secs: f32, + pub max_gap_secs: f32, + pub embedder_pool_size: usize, + pub execution_provider: ExecutionProvider, +} + +#[derive(Clone, Copy, Debug)] +pub enum ClustererKind { + NmeSc, + Ahc { threshold: f32 }, +} + +#[derive(Clone, Copy, Debug)] +pub enum ExecutionProvider { + Cpu, + CoreMl, // gated `coreml` + Nnapi, // gated `nnapi` + Cuda, // gated `cuda` + XnnPack, // gated `xnnpack` +} + +impl ExecutionProvider { + pub fn auto() -> Self { + // macOS aarch64 → CoreMl; Linux aarch64 → XnnPack; else Cpu. + } +} +``` + +`PipelineConfig::default()` returns: +```rust +profile: Profile::Balanced, +sample_rate: SampleRate::HZ_16000, +seg_window_secs: 10.0, seg_hop_secs: 0.5, +clusterer: ClustererKind::NmeSc, +max_speakers: 20, min_cluster_size: 12, +resegment_overlap: true, resegment_min_overlap_secs: 0.1, +min_speech_secs: 0.25, max_gap_secs: 0.5, +embedder_pool_size: num_cpus::get().min(4).max(1), +execution_provider: ExecutionProvider::auto(), +``` + +### `PipelineBuilder` API + +```rust +let pipeline = Pipeline::builder() + .profile(Profile::Mobile) + .with_models_from(ModelRegistry::default()?) + .resegment_overlap(true) + .build()?; + +let result = pipeline.run(&samples, SampleRate::HZ_16000)?; +``` + +Custom path: + +```rust +let pipeline = Pipeline::builder() + .config(my_config) + .with_segmenter(my_segmenter) + .with_embedder(my_embedder) + .with_clusterer(my_clusterer) + .build()?; +``` + +Builder methods (returning `&mut self`): +- `.config(PipelineConfig)` — replace full config (overrides `.profile`) +- `.profile(Profile)` — set config.profile +- `.with_models_from(ModelRegistry)` — load Mobile/Balanced ONNX через registry +- `.with_segmenter(Box)` — Custom only +- `.with_embedder(Box)` — Custom only +- `.with_clusterer(Box)` — Custom only +- `.with_resegmenter(Box)` — optional override (default `OverlapResegmenter::default()`) +- `.resegment_overlap(bool)` — set config.resegment_overlap +- `.embedder_pool_size(usize)`, `.max_speakers(u8)`, и т.п. — convenience setters +- `.build() -> Result` — финализация + +`build()` валидация: + +```rust +match cfg.profile { + Profile::Mobile | Profile::Balanced => { + if registry.is_none() { + return Err(ConfigError::MissingRegistry { profile: cfg.profile }); + } + if any_of(custom_seg, custom_emb, custom_clusterer).is_some() { + return Err(ConfigError::CustomComponentInProfile { + profile: cfg.profile, + offending: ... + }); + } + // Resolve registry → segmenter (PowersetSegmenter), embedder (CamPlusPlus or ResNet34Adapter), clusterer (NmeScClusterer or AhcClusterer per cfg.clusterer) + } + Profile::Custom => { + if registry.is_some() { + return Err(ConfigError::RegistryInCustomProfile); + } + if !(custom_seg && custom_emb && custom_clusterer) { + return Err(ConfigError::MissingCustomComponent { + missing: vec![...] + }); + } + } +} +``` + +`ConfigError` enum: + +```rust +#[derive(Debug, thiserror::Error)] +pub enum ConfigError { + #[error("profile {profile:?} requires .with_models_from() call")] + MissingRegistry { profile: Profile }, + + #[error("profile {profile:?} cannot accept .with_{offending}() — Custom only")] + CustomComponentInProfile { profile: Profile, offending: &'static str }, + + #[error("Custom profile cannot accept .with_models_from() — supply components individually")] + RegistryInCustomProfile, + + #[error("Custom profile missing required components: {missing:?}")] + MissingCustomComponent { missing: Vec<&'static str> }, + + #[error("ONNX model not found in registry: {model_id}")] + UnknownModel { model_id: String }, + + #[error("registry resolution failed: {0}")] + Registry(#[from] RegistryError), +} +``` + +### `Pipeline::run()` flow + +```text +samples (16 kHz f32 mono PCM, sr asserted) + ↓ +self.segmenter.segment(samples) -> Vec (M1) + ↓ +extract_overlap_time_ranges(&raw_segments) -> Vec<(TimeRange, u8 lo, u8 hi)> (M4) + ↓ +For each non-overlap RawSegment: + audio_chunk = samples[seg.time.start_samples..seg.time.end_samples] + audio_masked = apply_overlap_mask(audio_chunk, overlap_in_chunk_secs, sr) (M2) + emb = embedder_pool.embed(&audio_masked) (M2) + l2_normalize(&mut emb) + ↓ +embeddings: Vec>, primary_turns: Vec + ↓ +labels = clusterer.cluster(&embeddings) (M3) + ↓ +centroids = compute_centroids(&embeddings, &labels) (M4) + ↓ +overlap_inputs = build_overlap_region_inputs( + overlap_time_ranges, + primary_turns, + embedder_pool, + samples, +) (M4 helper) + ↓ if cfg.resegment_overlap +all_turns = resegmenter.resegment(ResegmentInputs { + primary_turns: &primary_turns, + speaker_centroids: ¢roids, + overlap_regions: &overlap_inputs, +})? + ↓ +all_turns = merge_segments(all_turns, cfg.max_gap_secs) (utils, legacy reused) +all_turns = filter_short(all_turns, cfg.min_speech_secs) + ↓ +DiarizationResult { turns: all_turns, segments: legacy_view, num_speakers: count(unique speakers) } +``` + +### Concurrency + +- `EmbedderPool` (M2) уже crossbeam-queue; pool_size from config. +- Segmentation последовательная (single ONNX session, lightweight). +- Clustering последовательная (faer eigen multithreaded внутри). +- Pipeline сам Send + Sync (composition of Send + Sync trait objects). + +### Re-exports (lib.rs) + +```rust +#[cfg(all( + feature = "pipeline_v2", + feature = "onnx", + feature = "segmentation", + feature = "embedder", + feature = "clusterer", + feature = "resegmentation", +))] +pub mod pipeline_v2; + +#[cfg(all( + feature = "pipeline_v2", + feature = "onnx", + feature = "segmentation", + feature = "embedder", + feature = "clusterer", + feature = "resegmentation", +))] +pub use pipeline_v2::{ + ClustererKind, ConfigError, ExecutionProvider, Pipeline as PipelineV2, + PipelineBuilder, PipelineConfig, PipelineError as PipelineV2Error, +}; +``` + +`polyvoice::Pipeline` (legacy) остаётся доступным; пользователь импортирует +`polyvoice::PipelineV2` для новой API. M6b сделает rename `PipelineV2 → Pipeline` +после удаления legacy. + +## File layout + +| Path | Action | Responsibility | +|---|---|---| +| `Cargo.toml` | modify | Add `pipeline_v2` feature (default-on, gates on existing onnx+segmentation+embedder+clusterer+resegmentation) | +| `src/pipeline_v2/mod.rs` | create | `Pipeline` struct, `run()`, `PipelineError`, `compile_error!` if feature gates absent | +| `src/pipeline_v2/builder.rs` | create | `PipelineBuilder` + `ConfigError` | +| `src/pipeline_v2/config.rs` | create | `PipelineConfig` + `ClustererKind` + `ExecutionProvider` + Default | +| `src/pipeline_v2/mocks.rs` | create | Test-only `MockSegmenter / MockEmbedder / MockClusterer` (`#[cfg(test)]`) | +| `src/lib.rs` | modify | `pub mod pipeline_v2;` gated, re-exports | +| `tests/pipeline_v2_synthetic_test.rs` | create | Builder validation + Custom profile end-to-end on synthetic data | +| `tests/pipeline_v2_e2e_test.rs` | create | `#[ignore]` integration test on real ONNX (Balanced profile, single voxconverse-test WAV) | +| `CHANGELOG.md` | modify | Unreleased M6a section | + +Total roughly 1100 LOC Rust + ~200 lines test + ~150 lines doc. + +## Acceptance criteria + +1. `cargo test --features download,onnx,segmentation,embedder,clusterer,resegmentation,pipeline_v2 --lib` зелёный (~5 new unit tests in builder/config/mod modules). +2. `cargo test --features download,onnx,segmentation,embedder,clusterer,resegmentation,pipeline_v2 --test pipeline_v2_synthetic_test` зелёный (~6 builder validation + Custom-profile end-to-end tests). +3. `cargo test --features download,onnx,segmentation,embedder,clusterer,resegmentation,pipeline_v2 --test pipeline_v2_e2e_test -- --ignored` локально зелёный (требует `models/` + `data/voxconverse-test/audio/abc.wav`; 30s runtime). +4. `cargo clippy --all-targets --all-features -- -D warnings` clean. +5. `cargo fmt --check` clean. +6. `cargo check --target wasm32-unknown-unknown --no-default-features --features resegmentation --lib` зелёный — `pipeline_v2` automatically gated out (требует onnx). +7. `cargo check --no-default-features` зелёный — pipeline_v2 gated out. +8. Legacy `polyvoice::Pipeline` всё ещё работает: existing tests `pipeline::tests::*` (если они есть) + `offline::tests::*` зелёные. +9. `bash scripts/release-gate.sh` выходит exit 2 (все pending steps пока не M6b/M9 — DER row остаётся PENDING потому что DER замер появляется в M6b). + +## Tests catalogue + +```text +src/pipeline_v2/builder.rs::tests + builder_default_profile_balanced + builder_set_profile_mobile + build_mobile_without_registry_errors + build_custom_without_components_errors + build_custom_with_registry_errors + build_balanced_with_custom_segmenter_errors + build_custom_full_succeeds + +src/pipeline_v2/config.rs::tests + pipeline_config_default + clusterer_kind_ahc_threshold + execution_provider_auto_returns_platform_default + +src/pipeline_v2/mod.rs::tests + pipeline_dyn_compatible + pipeline_run_unsupported_sample_rate_errors + +tests/pipeline_v2_synthetic_test.rs (integration) + builder_validation_mobile_missing_registry + builder_validation_custom_missing_components + pipeline_run_synthetic_two_speakers_through_custom_profile + pipeline_run_silence_returns_empty + pipeline_resegment_overlap_disabled_no_secondaries + pipeline_run_returns_sorted_turns + +tests/pipeline_v2_e2e_test.rs (#[ignore]) + e2e_balanced_profile_voxconverse_clip +``` + +Total: ~13 new tests. + +## Risks & mitigations + +| Risk | Likelihood | Mitigation | +|---|---|---| +| Build fails when `pipeline_v2` enabled without onnx | medium | `compile_error!` macro в `pipeline_v2/mod.rs` с актionable message; integration test в Cargo CI matrix | +| INT8 model load fails (provisional manifest sha256) | high | M5 использовал provisional sha256; ModelRegistry::ensure() проверит, скачает реальный INT8 если URL accessible. Если sha mismatch — actionable error через RegistryError. E2E test будет `#[ignore]` пока M5 follow-up завершён | +| ResNet34 / CAM++ ONNX input shape `[B, T, 80]` отличается от существующего FbankOnnxExtractor (`[B, 80, T]`?) | medium | M2 уже wraps FbankOnnxExtractor в ResNet34Adapter — shape handling there. M6a доверяет M2 wrapper. Если работает в M5 quantization — работает в M6a. | +| `extract_overlap_time_ranges` возвращает empty Vec из-за aggregator не помечающего pairs корректно | medium | Existing M1 unit tests в `aggregator.rs` показывают pairs созданы. M6a integration test покрывает этот path. Если empty — DiarizationResult без overlap turns (graceful, not error). | +| Custom profile mocks не покрывают real Segmenter/Embedder Send + Sync requirements | low | `MockSegmenter` etc. derive Send + Sync через trivial impl; trait objects `Box` + `Send + Sync` проверяется compile-time. | +| `merge_segments` (legacy from utils.rs) нарушает overlap turns (joins different speakers) | low | merge_segments matches на `speaker == next.speaker` — overlap turns с разными speakers не сольются. | +| `num_cpus` not available without dep | low | Use `std::thread::available_parallelism().map(NonZero::get).unwrap_or(1).min(4)` instead of `num_cpus::get()`. | + +## Dependencies on previous milestones + +- **M0**: `ModelRegistry`, `Profile`, `ProfileModels` — required для Mobile/Balanced loaders. +- **M1**: `PowersetSegmenter`, `Segmenter` trait — required для segmenter wiring. +- **M2**: `Embedder` trait, `CamPlusPlusExtractor`, `ResNet34Adapter`, `EmbedderPool`, `apply_overlap_mask` — required для embedder wiring. +- **M3**: `Clusterer` trait, `NmeScClusterer`, `AhcClusterer` — required. +- **M4**: `Resegmenter`, `OverlapResegmenter`, `compute_centroids`, `extract_overlap_time_ranges` — required. +- **M5**: INT8 manifest entries — used by registry для Mobile/Balanced. M6a accepts provisional sha256 (M5 follow-up will repin once VoxConverse-dev calibration completes). + +M6a doesn't add new external deps. Optionally adds `num_cpus` if `std::thread::available_parallelism()` rejected (preferred std-only path). + +## Out of scope (M6b follow-up) + +- Удаление `src/pipeline.rs`, `src/offline.rs`, `OfflineDiarizer`, `DiarizationConfig`, `VadConfig`, `EnergyVad`, `VoiceActivityDetector` trait, `segment_speech` fn (privatize в `vad/silero.rs`), `DummyExtractor` (move to `embedding/mock.rs` `#[cfg(test)]`), `OnnxEmbeddingExtractor` (replaced by ResNet34Adapter), `compute_fbank` privatization. +- Renaming `pipeline_v2 → pipeline`, `PipelineV2 → Pipeline`, removing `_v2` suffix. +- CLI rewrite: `polyvoice diarize --profile mobile|balanced`, removing legacy CLI args. +- FFI rewrite: `src/ffi.rs` accepts `Pipeline`/`PipelineConfig` instead of `OfflineDiarizer`/`DiarizationConfig`. New ABI bump. +- Python pyo3 rewrite: `python/src/lib.rs` exposes `Pipeline.builder()` instead of `OfflineDiarizer`. +- `polyvoice-bench` rewrite using new Pipeline. +- E2E DER замер на VoxConverse-test/AMI через polyvoice-bench → `tests/der_baseline.json`. +- Migration guide `docs/MIGRATING-FROM-0.5.md` per spec §11. +- `OnlineDiarizer` deprecation: `#[deprecated(since="2.0", note="redesigned in v1.1")]` annotation — kept as stub. + +## References + +- Roadmap: `docs/superpowers/specs/2026-05-07-perfect-diarization-roadmap-v1-design.md` + §3.1 (pipeline diagram), §5.2 (PipelineConfig), §5.3 (Segmenter / Embedder / Clusterer traits), + §5.4 (Pipeline + Builder), §5.5 (что удаляется в M6), §10.1 (M6 row). +- M0 plan: `docs/superpowers/plans/2026-05-07-m0-plumbing-and-registry-plan.md` +- M1 plan: `docs/superpowers/plans/2026-05-07-m1-powerset-segmenter-plan.md` +- M2 plan: `docs/superpowers/plans/2026-05-07-m2-cam-pp-embedder-plan.md` +- M3 plan: `docs/superpowers/plans/2026-05-07-m3-nme-sc-clusterer-plan.md` +- M4 plan: `docs/superpowers/plans/2026-05-07-m4-overlap-resegmenter-plan.md` +- M5 plan: `docs/superpowers/plans/2026-05-07-m5-int8-quantization-plan.md` + +## Open questions (closed) + +- ✅ M6 split в M6a (additive new Pipeline) + M6b (CLI/FFI/Python migration + legacy delete). +- ✅ Module placement: `src/pipeline_v2/` directory parallel к legacy `src/pipeline.rs` (variant 1 from brainstorming). +- ✅ Builder API per spec §5.4 verbatim. +- ✅ Custom profile contract: requires explicit `with_segmenter` + `with_embedder` + `with_clusterer`; conflict с `with_models_from` returns `ConfigError::RegistryInCustomProfile`. +- ✅ Tests: synthetic + `#[ignore]` E2E. + +## Follow-ups + +1. После одобрения spec: invoke `superpowers:writing-plans` для генерации M6a implementation plan в `docs/superpowers/plans/2026-05-07-m6a-pipeline-v2-plan.md` (стиль M3/M4 plan: TDD-задачи, atomic commits per task, git tag `m6a-complete`). +2. M6b в follow-up: cleanup legacy + CLI/FFI/Python migration + DER baseline closure. diff --git a/src/lib.rs b/src/lib.rs index ca315187..53091900 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -102,6 +102,29 @@ pub use resegmentation::{ #[cfg(all(feature = "resegmentation", feature = "segmentation"))] pub use resegmentation::extract_overlap_time_ranges; +#[cfg(all( + feature = "pipeline_v2", + feature = "onnx", + feature = "segmentation", + feature = "embedder", + feature = "clusterer", + feature = "resegmentation", +))] +pub mod pipeline_v2; + +#[cfg(all( + feature = "pipeline_v2", + feature = "onnx", + feature = "segmentation", + feature = "embedder", + feature = "clusterer", + feature = "resegmentation", +))] +pub use pipeline_v2::{ + ClustererKind, ConfigError, ExecutionProvider, Pipeline as PipelineV2, PipelineBuilder, + PipelineConfig, PipelineError as PipelineV2Error, +}; + #[cfg(feature = "onnx")] pub mod ecapa; #[cfg(feature = "onnx")] diff --git a/src/pipeline_v2/builder.rs b/src/pipeline_v2/builder.rs new file mode 100644 index 00000000..2147168e --- /dev/null +++ b/src/pipeline_v2/builder.rs @@ -0,0 +1,345 @@ +//! M6a — `PipelineBuilder` + `ConfigError`. +//! +//! Spec: `docs/superpowers/specs/2026-05-07-m6a-pipeline-v2-design.md` §4. + +use crate::clusterer::Clusterer; +use crate::embedder::Embedder; +use crate::models::{ModelRegistry, RegistryError}; +use crate::pipeline_v2::config::PipelineConfig; +use crate::resegmentation::Resegmenter; +use crate::segmentation::Segmenter; +use crate::types::Profile; + +#[derive(Debug, thiserror::Error)] +pub enum ConfigError { + #[error("profile {profile:?} requires .with_models_from() call")] + MissingRegistry { profile: Profile }, + + #[error("profile {profile:?} cannot accept .with_{offending}() — Custom only")] + CustomComponentInProfile { + profile: Profile, + offending: &'static str, + }, + + #[error("Custom profile cannot accept .with_models_from() — supply components individually")] + RegistryInCustomProfile, + + #[error("Custom profile missing required components: {missing:?}")] + MissingCustomComponent { missing: Vec<&'static str> }, + + #[error("ONNX model not found in registry: {model_id}")] + UnknownModel { model_id: String }, + + #[error("registry resolution failed: {0}")] + Registry(#[from] RegistryError), +} + +pub struct PipelineBuilder { + pub(crate) config: PipelineConfig, + pub(crate) registry: Option, + pub(crate) custom_segmenter: Option>, + pub(crate) custom_embedder: Option>, + pub(crate) custom_clusterer: Option>, + pub(crate) custom_resegmenter: Option>, +} + +impl PipelineBuilder { + #[allow(dead_code)] // used by Pipeline::builder() in Task 5 + pub(crate) fn new() -> Self { + Self { + config: PipelineConfig::default(), + registry: None, + custom_segmenter: None, + custom_embedder: None, + custom_clusterer: None, + custom_resegmenter: None, + } + } + + pub fn config(mut self, cfg: PipelineConfig) -> Self { + self.config = cfg; + self + } + + pub fn profile(mut self, p: Profile) -> Self { + self.config.profile = p; + self + } + + pub fn with_models_from(mut self, r: ModelRegistry) -> Self { + self.registry = Some(r); + self + } + + pub fn with_segmenter(mut self, s: Box) -> Self { + self.custom_segmenter = Some(s); + self + } + + pub fn with_embedder(mut self, e: Box) -> Self { + self.custom_embedder = Some(e); + self + } + + pub fn with_clusterer(mut self, c: Box) -> Self { + self.custom_clusterer = Some(c); + self + } + + pub fn with_resegmenter(mut self, r: Box) -> Self { + self.custom_resegmenter = Some(r); + self + } + + pub fn resegment_overlap(mut self, on: bool) -> Self { + self.config.resegment_overlap = on; + self + } + + pub fn embedder_pool_size(mut self, n: usize) -> Self { + self.config.embedder_pool_size = n.max(1); + self + } + + pub fn max_speakers(mut self, n: u8) -> Self { + self.config.max_speakers = n; + self + } + + pub fn validate(&self) -> Result<(), ConfigError> { + match self.config.profile { + Profile::Mobile | Profile::Balanced => { + if self.custom_segmenter.is_some() { + return Err(ConfigError::CustomComponentInProfile { + profile: self.config.profile, + offending: "segmenter", + }); + } + if self.custom_embedder.is_some() { + return Err(ConfigError::CustomComponentInProfile { + profile: self.config.profile, + offending: "embedder", + }); + } + if self.custom_clusterer.is_some() { + return Err(ConfigError::CustomComponentInProfile { + profile: self.config.profile, + offending: "clusterer", + }); + } + if self.registry.is_none() { + return Err(ConfigError::MissingRegistry { + profile: self.config.profile, + }); + } + } + Profile::Custom => { + if self.registry.is_some() { + return Err(ConfigError::RegistryInCustomProfile); + } + let mut missing: Vec<&'static str> = Vec::new(); + if self.custom_segmenter.is_none() { + missing.push("segmenter"); + } + if self.custom_embedder.is_none() { + missing.push("embedder"); + } + if self.custom_clusterer.is_none() { + missing.push("clusterer"); + } + if !missing.is_empty() { + return Err(ConfigError::MissingCustomComponent { missing }); + } + } + } + Ok(()) + } +} + +use crate::pipeline_v2::Pipeline; +use crate::pipeline_v2::config::ClustererKind; +use crate::resegmentation::OverlapResegmenter; + +impl PipelineBuilder { + /// Validate + construct the inner `Pipeline`. + pub fn build(self) -> Result { + self.validate()?; + let resegmenter = self + .custom_resegmenter + .unwrap_or_else(|| Box::new(OverlapResegmenter::default())); + + match self.config.profile { + Profile::Custom => { + let segmenter = + self.custom_segmenter + .ok_or_else(|| ConfigError::MissingCustomComponent { + missing: vec!["segmenter"], + })?; + let embedder = + self.custom_embedder + .ok_or_else(|| ConfigError::MissingCustomComponent { + missing: vec!["embedder"], + })?; + let clusterer = + self.custom_clusterer + .ok_or_else(|| ConfigError::MissingCustomComponent { + missing: vec!["clusterer"], + })?; + Ok(Pipeline::from_components( + self.config, + segmenter, + embedder, + clusterer, + resegmenter, + )) + } + Profile::Mobile | Profile::Balanced => { + let registry = self.registry.ok_or(ConfigError::MissingRegistry { + profile: self.config.profile, + })?; + let profile_models = registry.ensure_for_profile(self.config.profile)?; + let segmenter: Box = Box::new( + crate::segmentation::PowersetSegmenter::new(&profile_models.segmenter_path) + .map_err(|e| ConfigError::UnknownModel { + model_id: format!("powerset (cause: {e})"), + })?, + ); + let embedder: Box = match self.config.profile { + Profile::Mobile => Box::new( + crate::embedder::CamPlusPlusExtractor::new( + &profile_models.embedder_path, + self.config.profile.embedding_dim(), + self.config.embedder_pool_size, + ) + .map_err(|e| ConfigError::UnknownModel { + model_id: format!("cam_pp (cause: {e})"), + })?, + ), + Profile::Balanced => Box::new( + crate::embedder::ResNet34Adapter::new( + &profile_models.embedder_path, + self.config.embedder_pool_size, + ) + .map_err(|e| ConfigError::UnknownModel { + model_id: format!("resnet34 (cause: {e})"), + })?, + ), + Profile::Custom => unreachable!("Profile::Custom handled above"), + }; + let clusterer: Box = match self.config.clusterer { + ClustererKind::Ahc { .. } => Box::new(crate::clusterer::AhcClusterer::new( + self.config.max_speakers as usize, + )), + #[cfg(feature = "spectral")] + ClustererKind::NmeSc => Box::new(crate::clusterer::NmeScClusterer::new( + self.config.max_speakers as usize, + )), + #[cfg(not(feature = "spectral"))] + ClustererKind::NmeSc => Box::new(crate::clusterer::AhcClusterer::new( + self.config.max_speakers as usize, + )), + }; + Ok(Pipeline::from_components( + self.config, + segmenter, + embedder, + clusterer, + resegmenter, + )) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::pipeline_v2::mocks::{MockClusterer, MockEmbedder, MockSegmenter}; + + fn fresh() -> PipelineBuilder { + PipelineBuilder::new() + } + + #[test] + fn builder_default_profile_balanced() { + let b = fresh(); + assert_eq!(b.config.profile, Profile::Balanced); + } + + #[test] + fn builder_profile_setter() { + let b = fresh().profile(Profile::Mobile); + assert_eq!(b.config.profile, Profile::Mobile); + } + + #[test] + fn validate_mobile_without_registry_errors() { + let err = fresh().profile(Profile::Mobile).validate().unwrap_err(); + assert!(matches!( + err, + ConfigError::MissingRegistry { + profile: Profile::Mobile + } + )); + } + + #[test] + fn validate_custom_without_components_errors() { + let err = fresh().profile(Profile::Custom).validate().unwrap_err(); + match err { + ConfigError::MissingCustomComponent { missing } => { + assert!(missing.contains(&"segmenter")); + assert!(missing.contains(&"embedder")); + assert!(missing.contains(&"clusterer")); + } + other => panic!("unexpected error variant: {other:?}"), + } + } + + #[test] + fn validate_custom_with_full_components_succeeds() { + let b = fresh() + .profile(Profile::Custom) + .with_segmenter(Box::new(MockSegmenter::default())) + .with_embedder(Box::new(MockEmbedder::default())) + .with_clusterer(Box::new(MockClusterer::default())); + b.validate().expect("custom + 3 components must validate"); + } + + #[test] + fn validate_balanced_with_custom_segmenter_errors() { + let b = fresh() + .profile(Profile::Balanced) + .with_segmenter(Box::new(MockSegmenter::default())); + let err = b.validate().unwrap_err(); + assert!(matches!( + err, + ConfigError::CustomComponentInProfile { + offending: "segmenter", + .. + } + )); + } + + #[test] + fn validate_custom_with_registry_errors() { + let registry = match ModelRegistry::default() { + Ok(r) => r, + Err(_) => return, + }; + let b = fresh() + .profile(Profile::Custom) + .with_segmenter(Box::new(MockSegmenter::default())) + .with_embedder(Box::new(MockEmbedder::default())) + .with_clusterer(Box::new(MockClusterer::default())) + .with_models_from(registry); + let err = b.validate().unwrap_err(); + assert!(matches!(err, ConfigError::RegistryInCustomProfile)); + } + + #[test] + fn embedder_pool_size_clamps_to_1() { + let b = fresh().embedder_pool_size(0); + assert_eq!(b.config.embedder_pool_size, 1); + } +} diff --git a/src/pipeline_v2/config.rs b/src/pipeline_v2/config.rs new file mode 100644 index 00000000..6e317d03 --- /dev/null +++ b/src/pipeline_v2/config.rs @@ -0,0 +1,119 @@ +//! M6a — `PipelineConfig`, `ClustererKind`, `ExecutionProvider`. +//! +//! Spec: `docs/superpowers/specs/2026-05-07-m6a-pipeline-v2-design.md` §3. + +use crate::types::{Profile, SampleRate}; + +/// Top-level configuration for the v1.0 Pipeline. Mirrors spec §5.2 verbatim. +#[derive(Clone, Debug)] +pub struct PipelineConfig { + pub profile: Profile, + pub sample_rate: SampleRate, + pub seg_window_secs: f32, + pub seg_hop_secs: f32, + pub clusterer: ClustererKind, + pub max_speakers: u8, + pub min_cluster_size: usize, + pub resegment_overlap: bool, + pub resegment_min_overlap_secs: f32, + pub min_speech_secs: f32, + pub max_gap_secs: f32, + pub embedder_pool_size: usize, + pub execution_provider: ExecutionProvider, +} + +impl Default for PipelineConfig { + fn default() -> Self { + Self { + profile: Profile::Balanced, + sample_rate: SampleRate::new(16000).unwrap_or_default(), + seg_window_secs: 10.0, + seg_hop_secs: 0.5, + clusterer: ClustererKind::NmeSc, + max_speakers: 20, + min_cluster_size: 12, + resegment_overlap: true, + resegment_min_overlap_secs: 0.1, + min_speech_secs: 0.25, + max_gap_secs: 0.5, + embedder_pool_size: default_pool_size(), + execution_provider: ExecutionProvider::auto(), + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum ClustererKind { + NmeSc, + Ahc { threshold: f32 }, +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum ExecutionProvider { + Cpu, + CoreMl, + Nnapi, + Cuda, + XnnPack, +} + +impl ExecutionProvider { + pub fn auto() -> Self { + #[cfg(all(target_os = "macos", target_arch = "aarch64"))] + return Self::CoreMl; + #[cfg(all(target_os = "linux", target_arch = "aarch64"))] + return Self::XnnPack; + #[cfg(not(any( + all(target_os = "macos", target_arch = "aarch64"), + all(target_os = "linux", target_arch = "aarch64"), + )))] + return Self::Cpu; + } +} + +fn default_pool_size() -> usize { + std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(1) + .clamp(1, 4) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::Profile; + + #[test] + fn pipeline_config_default_is_balanced() { + let cfg = PipelineConfig::default(); + assert_eq!(cfg.profile, Profile::Balanced); + assert_eq!(cfg.sample_rate.get(), 16000); + assert!((cfg.seg_window_secs - 10.0).abs() < f32::EPSILON); + assert!((cfg.seg_hop_secs - 0.5).abs() < f32::EPSILON); + assert!(matches!(cfg.clusterer, ClustererKind::NmeSc)); + assert_eq!(cfg.max_speakers, 20); + assert_eq!(cfg.min_cluster_size, 12); + assert!(cfg.resegment_overlap); + assert!((cfg.resegment_min_overlap_secs - 0.1).abs() < f32::EPSILON); + assert!((cfg.min_speech_secs - 0.25).abs() < f32::EPSILON); + assert!((cfg.max_gap_secs - 0.5).abs() < f32::EPSILON); + assert!(cfg.embedder_pool_size >= 1); + assert!(cfg.embedder_pool_size <= 4); + } + + #[test] + fn clusterer_kind_ahc_with_threshold() { + let k = ClustererKind::Ahc { threshold: 0.7 }; + if let ClustererKind::Ahc { threshold } = k { + assert!((threshold - 0.7).abs() < f32::EPSILON); + } else { + panic!("expected Ahc variant"); + } + } + + #[test] + fn execution_provider_auto_returns_some_variant() { + let ep = ExecutionProvider::auto(); + let _ = ep; + } +} diff --git a/src/pipeline_v2/mocks.rs b/src/pipeline_v2/mocks.rs new file mode 100644 index 00000000..9dc98900 --- /dev/null +++ b/src/pipeline_v2/mocks.rs @@ -0,0 +1,102 @@ +//! M6a — test-only `Mock{Segmenter,Embedder,Clusterer}` for the +//! `pipeline_v2` builder validation tests and the synthetic integration +//! test in `tests/pipeline_v2_synthetic_test.rs`. + +use crate::clusterer::{Clusterer, ClustererError}; +use crate::embedder::{Embedder, EmbedderError}; +use crate::resegmentation::{ResegmentError, ResegmentInputs, Resegmenter}; +use crate::segmentation::{RawSegment, SegmentationError, Segmenter}; +use crate::types::{Confidence, SpeakerTurn, TimeRange}; + +/// Constant-output `Segmenter` for builder tests. +#[derive(Default)] +pub struct MockSegmenter { + pub segments: Vec, +} + +impl Segmenter for MockSegmenter { + fn segment(&self, _audio: &[f32]) -> Result, SegmentationError> { + Ok(self.segments.clone()) + } + + fn max_local_speakers(&self) -> usize { + 3 + } + + fn supports_overlap(&self) -> bool { + true + } +} + +/// Constant-output `Embedder` for builder tests. +pub struct MockEmbedder { + pub embedding: Vec, +} + +impl Default for MockEmbedder { + fn default() -> Self { + // 192-d unit vector along axis 0; matches CAM++ output dim used + // throughout the spec. + let mut v = vec![0.0_f32; 192]; + v[0] = 1.0; + Self { embedding: v } + } +} + +impl Embedder for MockEmbedder { + fn dim(&self) -> usize { + self.embedding.len() + } + + fn embed(&self, _audio: &[f32]) -> Result, EmbedderError> { + Ok(self.embedding.clone()) + } +} + +/// Constant-label `Clusterer` for builder tests. +#[derive(Default)] +pub struct MockClusterer { + pub labels: Vec, +} + +impl Clusterer for MockClusterer { + fn cluster(&self, embeddings: &[Vec]) -> Result, ClustererError> { + if self.labels.is_empty() { + // Default: single cluster. + return Ok(vec![0; embeddings.len()]); + } + if self.labels.len() != embeddings.len() { + return Err(ClustererError::AlgorithmFailed { + detail: "MockClusterer labels length mismatch".to_owned(), + }); + } + Ok(self.labels.clone()) + } + + fn max_clusters(&self) -> usize { + 16 + } +} + +/// Pass-through `Resegmenter` (returns input primary turns sorted, no +/// secondary speakers added). +#[derive(Default)] +pub struct PassThroughResegmenter; + +impl Resegmenter for PassThroughResegmenter { + fn resegment(&self, inputs: ResegmentInputs<'_>) -> Result, ResegmentError> { + let mut out: Vec = inputs.primary_turns.to_vec(); + out.sort_by(|a, b| a.time.start.total_cmp(&b.time.start)); + Ok(out) + } +} + +/// Convenience constructor for a single `RawSegment` used in tests. +pub fn raw_segment(start: f64, end: f64, spk: u8, overlap: bool) -> RawSegment { + RawSegment { + time: TimeRange { start, end }, + local_speaker_idx: spk, + is_overlap: overlap, + confidence: Confidence::new(0.9).unwrap(), + } +} diff --git a/src/pipeline_v2/mod.rs b/src/pipeline_v2/mod.rs new file mode 100644 index 00000000..4ef9fc7c --- /dev/null +++ b/src/pipeline_v2/mod.rs @@ -0,0 +1,335 @@ +//! M6a — additive `polyvoice::pipeline_v2` module. +//! +//! Spec: `docs/superpowers/specs/2026-05-07-m6a-pipeline-v2-design.md`. + +#[cfg(not(all( + feature = "pipeline_v2", + feature = "onnx", + feature = "segmentation", + feature = "embedder", + feature = "clusterer", + feature = "resegmentation", +)))] +compile_error!( + "pipeline_v2 requires onnx + segmentation + embedder + clusterer + resegmentation features" +); + +pub mod builder; +pub mod config; + +#[cfg(test)] +pub mod mocks; + +use crate::clusterer::{Clusterer, ClustererError}; +use crate::embedder::{Embedder, EmbedderError, apply_overlap_mask}; +use crate::models::RegistryError; +use crate::resegmentation::{ + OverlapRegionInput, ResegmentError, ResegmentInputs, Resegmenter, SpeakerCentroid, + compute_centroids, extract_overlap_time_ranges, +}; +use crate::segmentation::{SegmentationError, Segmenter}; +use crate::types::{DiarizationResult, SampleRate, Segment, SpeakerId, SpeakerTurn, TimeRange}; +use crate::utils::{l2_normalize, merge_segments}; + +pub use builder::{ConfigError, PipelineBuilder}; +pub use config::{ClustererKind, ExecutionProvider, PipelineConfig}; + +#[derive(Debug, thiserror::Error)] +pub enum PipelineError { + #[error("audio sample rate {actual} unsupported, expected 16000")] + UnsupportedSampleRate { actual: u32 }, + #[error("segmentation failed: {0}")] + Segmentation(#[from] SegmentationError), + #[error("embedding failed: {0}")] + Embedding(#[from] EmbedderError), + #[error("clustering failed: {0}")] + Clustering(#[from] ClustererError), + #[error("resegmentation failed: {0}")] + Resegment(#[from] ResegmentError), + #[error("config error: {0}")] + Config(#[from] ConfigError), + #[error("model registry error: {0}")] + Registry(#[from] RegistryError), + #[error("model load error: {detail}")] + ModelLoad { detail: String }, +} + +pub struct Pipeline { + config: PipelineConfig, + segmenter: Box, + embedder: Box, + clusterer: Box, + resegmenter: Box, +} + +impl Pipeline { + pub fn builder() -> PipelineBuilder { + PipelineBuilder::new() + } + + pub(crate) fn from_components( + config: PipelineConfig, + segmenter: Box, + embedder: Box, + clusterer: Box, + resegmenter: Box, + ) -> Self { + Self { + config, + segmenter, + embedder, + clusterer, + resegmenter, + } + } + + pub fn config(&self) -> &PipelineConfig { + &self.config + } + + pub fn run(&self, samples: &[f32], sr: SampleRate) -> Result { + if sr.get() != self.config.sample_rate.get() { + return Err(PipelineError::UnsupportedSampleRate { actual: sr.get() }); + } + + let raw_segments = self.segmenter.segment(samples)?; + if raw_segments.is_empty() { + return Ok(DiarizationResult { + segments: Vec::new(), + turns: Vec::new(), + num_speakers: 0, + }); + } + + let overlap_ranges = extract_overlap_time_ranges(&raw_segments); + let primary_segments: Vec<_> = raw_segments + .iter() + .filter(|s| !s.is_overlap) + .cloned() + .collect(); + + let sample_rate = self.config.sample_rate.get() as f64; + let mut embeddings: Vec> = Vec::with_capacity(primary_segments.len()); + let mut valid_segments: Vec<_> = Vec::with_capacity(primary_segments.len()); + for seg in &primary_segments { + let start_idx = (seg.time.start * sample_rate) as usize; + let end_idx = ((seg.time.end * sample_rate) as usize).min(samples.len()); + if end_idx <= start_idx { + continue; + } + let chunk = &samples[start_idx..end_idx]; + // Fix 2: zero-fill any overlap regions inside this primary chunk before embedding. + let seg_start = seg.time.start; + let seg_end = seg.time.end; + let local_overlaps: Vec<(f32, f32)> = overlap_ranges + .iter() + .filter_map(|(ot, _, _)| { + let lo = ot.start.max(seg_start); + let hi = ot.end.min(seg_end); + if hi > lo { + Some(((lo - seg_start) as f32, (hi - seg_start) as f32)) + } else { + None + } + }) + .collect(); + let masked = apply_overlap_mask(chunk, &local_overlaps, self.config.sample_rate.get()); + let mut emb = self.embedder.embed(&masked)?; + l2_normalize(&mut emb); + embeddings.push(emb); + valid_segments.push(seg); + } + + if embeddings.is_empty() { + return Ok(DiarizationResult { + segments: Vec::new(), + turns: Vec::new(), + num_speakers: 0, + }); + } + + let labels = self.clusterer.cluster(&embeddings)?; + + // Fix 1: zip valid_segments (survivors only) with labels — not primary_segments. + let mut primary_turns: Vec = valid_segments + .iter() + .zip(labels.iter()) + .map(|(seg, &lbl)| SpeakerTurn { + speaker: SpeakerId(lbl as u32), + time: seg.time, + text: None, + }) + .collect(); + + let centroids: Vec = compute_centroids(&embeddings, &labels); + + let mut all_turns: Vec = if self.config.resegment_overlap + && !overlap_ranges.is_empty() + && centroids.len() >= 2 + { + let overlap_inputs = + self.build_overlap_inputs(&overlap_ranges, &primary_turns, samples)?; + self.resegmenter.resegment(ResegmentInputs { + primary_turns: &primary_turns, + speaker_centroids: ¢roids, + overlap_regions: &overlap_inputs, + })? + } else { + primary_turns.sort_by(|a, b| a.time.start.total_cmp(&b.time.start)); + primary_turns + }; + + // Fix 4: guarantee sorted order regardless of which Resegmenter impl ran. + all_turns.sort_by(|a, b| a.time.start.total_cmp(&b.time.start)); + + let min_secs = self.config.min_speech_secs as f64; + all_turns.retain(|t| t.time.duration() >= min_secs); + + let max_gap = self.config.max_gap_secs as f64; + let merged_segments: Vec = all_turns + .iter() + .map(|t| Segment { + time: t.time, + speaker: Some(t.speaker), + confidence: None, + }) + .collect(); + let merged_segments = merge_segments(merged_segments, max_gap); + let merged_turns: Vec = merged_segments + .iter() + .filter_map(|s| { + s.speaker.map(|spk| SpeakerTurn { + speaker: spk, + time: s.time, + text: None, + }) + }) + .collect(); + + let num_speakers = merged_turns + .iter() + .map(|t| t.speaker.0) + .collect::>() + .len(); + + Ok(DiarizationResult { + segments: merged_segments, + turns: merged_turns, + num_speakers, + }) + } + + fn build_overlap_inputs( + &self, + overlap_ranges: &[(TimeRange, u8, u8)], + primary_turns: &[SpeakerTurn], + samples: &[f32], + ) -> Result, PipelineError> { + let sample_rate = self.config.sample_rate.get() as f64; + let mut out = Vec::with_capacity(overlap_ranges.len()); + for (time, _lo, _hi) in overlap_ranges { + let primary = primary_turns + .iter() + .find(|t| t.time.start <= time.start && time.end <= t.time.end) + .map(|t| t.speaker) + .unwrap_or_else(|| { + // Fix 3: nearest by midpoint distance, not start-time distance. + let mid = (time.start + time.end) / 2.0; + let tmid = |t: &SpeakerTurn| (t.time.start + t.time.end) / 2.0; + primary_turns + .iter() + .min_by(|a, b| (tmid(a) - mid).abs().total_cmp(&(tmid(b) - mid).abs())) + .map(|t| t.speaker) + .unwrap_or(SpeakerId(0)) + }); + let start_idx = (time.start * sample_rate) as usize; + let end_idx = ((time.end * sample_rate) as usize).min(samples.len()); + if end_idx <= start_idx { + continue; + } + let chunk = &samples[start_idx..end_idx]; + let mut emb = self.embedder.embed(chunk)?; + l2_normalize(&mut emb); + out.push(OverlapRegionInput { + time: *time, + primary_speaker: primary, + embedding: emb, + }); + } + Ok(out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::pipeline_v2::mocks::{MockClusterer, MockEmbedder, MockSegmenter, raw_segment}; + use crate::resegmentation::OverlapResegmenter; + use crate::types::Profile; + + fn pipeline_with_segments(segs: Vec) -> Pipeline { + let cfg = PipelineConfig { + profile: Profile::Custom, + resegment_overlap: false, + min_speech_secs: 0.0, + max_gap_secs: 0.0, + ..PipelineConfig::default() + }; + Pipeline::from_components( + cfg, + Box::new(MockSegmenter { segments: segs }), + Box::new(MockEmbedder::default()), + Box::new(MockClusterer::default()), + Box::new(OverlapResegmenter::default()), + ) + } + + #[test] + fn pipeline_run_unsupported_sample_rate_returns_err() { + let p = pipeline_with_segments(vec![raw_segment(0.0, 1.0, 0, false)]); + let bad = SampleRate::new(8000).unwrap(); + let err = p.run(&vec![0.0_f32; 8000], bad).unwrap_err(); + assert!(matches!( + err, + PipelineError::UnsupportedSampleRate { actual: 8000 } + )); + } + + #[test] + fn pipeline_run_silence_returns_empty() { + let p = pipeline_with_segments(Vec::new()); + let result = p + .run(&vec![0.0_f32; 16000], SampleRate::new(16000).unwrap()) + .unwrap(); + assert!(result.turns.is_empty()); + assert_eq!(result.num_speakers, 0); + } + + #[test] + fn pipeline_run_two_segments_one_cluster() { + let segs = vec![ + raw_segment(0.0, 1.0, 0, false), + raw_segment(1.5, 2.5, 0, false), + ]; + let p = pipeline_with_segments(segs); + let result = p + .run(&vec![0.0_f32; 16000 * 3], SampleRate::new(16000).unwrap()) + .unwrap(); + assert_eq!(result.num_speakers, 1); + assert!(!result.turns.is_empty()); + } + + #[test] + fn pipeline_resegment_overlap_disabled_path_used() { + let segs = vec![ + raw_segment(0.0, 1.0, 0, true), + raw_segment(0.0, 1.0, 1, true), + raw_segment(1.5, 2.5, 0, false), + ]; + let p = pipeline_with_segments(segs); + let result = p + .run(&vec![0.0_f32; 16000 * 3], SampleRate::new(16000).unwrap()) + .unwrap(); + assert!(result.num_speakers <= 1); + } +} diff --git a/tests/pipeline_v2_e2e_test.rs b/tests/pipeline_v2_e2e_test.rs new file mode 100644 index 00000000..5151f4b4 --- /dev/null +++ b/tests/pipeline_v2_e2e_test.rs @@ -0,0 +1,76 @@ +//! M6a — `#[ignore]` end-to-end test for `polyvoice::pipeline_v2`. +//! +//! Requires the Balanced ONNX bundle to be cached (run +//! `cargo run --features cli --bin polyvoice -- download-models --profile balanced` +//! once before invoking with `cargo test -- --ignored e2e`). Reads a single +//! WAV from `data/voxconverse-test/audio/` and asserts the pipeline returns +//! at least one turn with valid speaker IDs. + +#![cfg(all( + feature = "pipeline_v2", + feature = "download", + feature = "onnx", + feature = "segmentation", + feature = "embedder", + feature = "clusterer", + feature = "resegmentation", +))] + +use polyvoice::models::ModelRegistry; +use polyvoice::pipeline_v2::Pipeline; +use polyvoice::types::{Profile, SampleRate}; +use polyvoice::wav::read_wav; +use std::path::Path; + +fn first_voxconverse_wav() -> Option { + let dir = Path::new("data/voxconverse-test/audio"); + if !dir.is_dir() { + return None; + } + let mut entries: Vec<_> = std::fs::read_dir(dir) + .ok()? + .filter_map(|e| e.ok()) + .filter(|e| e.path().extension().is_some_and(|ext| ext == "wav")) + .map(|e| e.path()) + .collect(); + entries.sort(); + entries.into_iter().next() +} + +#[ignore = "requires cached ONNX bundle + a wav file under data/voxconverse-test/audio/"] +#[test] +fn e2e_balanced_profile_voxconverse_clip() { + let wav_path = match first_voxconverse_wav() { + Some(p) => p, + None => panic!( + "data/voxconverse-test/audio/ is empty — run scripts/download-voxconverse-test.sh first" + ), + }; + let (samples, sr_hz) = + read_wav(&wav_path).expect("WAV read failure — check the file is 16 kHz mono"); + assert_eq!(sr_hz, 16000, "only 16 kHz WAVs supported"); + let registry = ModelRegistry::default() + .expect("default ModelRegistry should resolve a writable cache dir"); + let pipeline = Pipeline::builder() + .profile(Profile::Balanced) + .with_models_from(registry) + .build() + .expect( + "Balanced profile build should succeed when cached ONNX is present — \ + run `polyvoice download-models --profile balanced` first", + ); + let result = pipeline + .run(&samples, SampleRate::new(16000).unwrap()) + .expect("pipeline.run on a real VoxConverse clip should succeed"); + assert!( + result.num_speakers >= 1, + "expected at least 1 speaker, got {}", + result.num_speakers + ); + for w in result.turns.windows(2) { + assert!( + w[0].time.start <= w[1].time.start, + "turns must be sorted by start time" + ); + } +} diff --git a/tests/pipeline_v2_synthetic_test.rs b/tests/pipeline_v2_synthetic_test.rs new file mode 100644 index 00000000..0685f225 --- /dev/null +++ b/tests/pipeline_v2_synthetic_test.rs @@ -0,0 +1,250 @@ +//! M6a — synthetic-data integration tests for `polyvoice::pipeline_v2`. +//! +//! Pure-CPU; no ONNX. Covers builder validation paths, end-to-end +//! Custom-profile run, and overlap-resegmentation toggling. + +#![cfg(all( + feature = "pipeline_v2", + feature = "onnx", + feature = "segmentation", + feature = "embedder", + feature = "clusterer", + feature = "resegmentation", +))] + +use polyvoice::clusterer::{Clusterer, ClustererError}; +use polyvoice::embedder::{Embedder, EmbedderError}; +use polyvoice::pipeline_v2::{ClustererKind, ConfigError, Pipeline, PipelineConfig, PipelineError}; +use polyvoice::segmentation::{RawSegment, SegmentationError, Segmenter}; +use polyvoice::types::{Confidence, Profile, SampleRate, TimeRange}; + +// ----- Local mocks (re-defined here because the crate-internal mocks are #[cfg(test)]) ----- + +struct ConstSegmenter(Vec); +impl Segmenter for ConstSegmenter { + fn segment(&self, _audio: &[f32]) -> Result, SegmentationError> { + Ok(self.0.clone()) + } + fn max_local_speakers(&self) -> usize { + 3 + } + fn supports_overlap(&self) -> bool { + true + } +} + +type AxisPicker = Box usize + Send + Sync>; + +struct AxisEmbedder { + dim: usize, + axis_picker: AxisPicker, +} +impl Embedder for AxisEmbedder { + fn dim(&self) -> usize { + self.dim + } + fn embed(&self, audio: &[f32]) -> Result, EmbedderError> { + let axis = (self.axis_picker)(audio).min(self.dim - 1); + let mut v = vec![0.0_f32; self.dim]; + v[axis] = 1.0; + Ok(v) + } +} + +struct PerSampleClusterer { + labels: Vec, +} +impl Clusterer for PerSampleClusterer { + fn cluster(&self, embeddings: &[Vec]) -> Result, ClustererError> { + if embeddings.len() != self.labels.len() { + return Err(ClustererError::AlgorithmFailed { + detail: format!( + "labels {} vs embeddings {}", + self.labels.len(), + embeddings.len() + ), + }); + } + Ok(self.labels.clone()) + } + fn max_clusters(&self) -> usize { + 16 + } +} + +fn raw(start: f64, end: f64, spk: u8, overlap: bool) -> RawSegment { + RawSegment { + time: TimeRange { start, end }, + local_speaker_idx: spk, + is_overlap: overlap, + confidence: Confidence::new(0.9).unwrap(), + } +} + +fn axis_picker_constant(axis: usize) -> AxisPicker { + Box::new(move |_| axis) +} + +// ----- Tests ----- + +#[test] +fn builder_validation_mobile_missing_registry() { + let result = Pipeline::builder().profile(Profile::Mobile).build(); + let err = match result { + Err(e) => e, + Ok(_) => panic!("expected MissingRegistry for Mobile without registry"), + }; + assert!(matches!( + err, + ConfigError::MissingRegistry { + profile: Profile::Mobile + } + )); +} + +#[test] +fn builder_validation_custom_missing_components() { + let result = Pipeline::builder().profile(Profile::Custom).build(); + let err = match result { + Err(e) => e, + Ok(_) => panic!("expected MissingCustomComponent for Custom without components"), + }; + assert!(matches!(err, ConfigError::MissingCustomComponent { .. })); +} + +#[test] +fn pipeline_run_unsupported_sample_rate_errors() { + let cfg = PipelineConfig { + profile: Profile::Custom, + ..PipelineConfig::default() + }; + let p = Pipeline::builder() + .config(cfg) + .with_segmenter(Box::new(ConstSegmenter(vec![raw(0.0, 1.0, 0, false)]))) + .with_embedder(Box::new(AxisEmbedder { + dim: 8, + axis_picker: axis_picker_constant(0), + })) + .with_clusterer(Box::new(PerSampleClusterer { labels: vec![0] })) + .build() + .expect("custom build"); + let err = p + .run(&vec![0.0_f32; 8000], SampleRate::new(8000).unwrap()) + .unwrap_err(); + assert!(matches!( + err, + PipelineError::UnsupportedSampleRate { actual: 8000 } + )); +} + +#[test] +fn pipeline_run_silence_returns_empty() { + let cfg = PipelineConfig { + profile: Profile::Custom, + ..PipelineConfig::default() + }; + let p = Pipeline::builder() + .config(cfg) + .with_segmenter(Box::new(ConstSegmenter(Vec::new()))) + .with_embedder(Box::new(AxisEmbedder { + dim: 8, + axis_picker: axis_picker_constant(0), + })) + .with_clusterer(Box::new(PerSampleClusterer { labels: vec![] })) + .build() + .expect("custom build"); + let r = p + .run(&vec![0.0_f32; 16000], SampleRate::new(16000).unwrap()) + .unwrap(); + assert!(r.turns.is_empty()); + assert_eq!(r.num_speakers, 0); +} + +#[test] +fn pipeline_run_two_speakers_through_custom_profile() { + let segs = vec![raw(0.0, 1.0, 0, false), raw(2.0, 3.0, 1, false)]; + let cfg = PipelineConfig { + profile: Profile::Custom, + clusterer: ClustererKind::NmeSc, + resegment_overlap: false, + min_speech_secs: 0.0, + max_gap_secs: 0.0, + ..PipelineConfig::default() + }; + let p = Pipeline::builder() + .config(cfg) + .with_segmenter(Box::new(ConstSegmenter(segs))) + .with_embedder(Box::new(AxisEmbedder { + dim: 8, + axis_picker: axis_picker_constant(0), + })) + .with_clusterer(Box::new(PerSampleClusterer { labels: vec![0, 1] })) + .build() + .expect("custom build"); + let r = p + .run(&vec![0.0_f32; 16000 * 4], SampleRate::new(16000).unwrap()) + .unwrap(); + assert_eq!(r.num_speakers, 2); + assert_eq!(r.turns.len(), 2); + let speakers: Vec = r.turns.iter().map(|t| t.speaker.0).collect(); + assert!(speakers.contains(&0)); + assert!(speakers.contains(&1)); +} + +#[test] +fn pipeline_run_returns_sorted_turns() { + let segs = vec![raw(2.0, 3.0, 0, false), raw(0.0, 1.0, 0, false)]; + let cfg = PipelineConfig { + profile: Profile::Custom, + resegment_overlap: false, + min_speech_secs: 0.0, + max_gap_secs: 0.0, + ..PipelineConfig::default() + }; + let p = Pipeline::builder() + .config(cfg) + .with_segmenter(Box::new(ConstSegmenter(segs))) + .with_embedder(Box::new(AxisEmbedder { + dim: 8, + axis_picker: axis_picker_constant(0), + })) + .with_clusterer(Box::new(PerSampleClusterer { labels: vec![0, 0] })) + .build() + .expect("custom build"); + let r = p + .run(&vec![0.0_f32; 16000 * 4], SampleRate::new(16000).unwrap()) + .unwrap(); + for w in r.turns.windows(2) { + assert!(w[0].time.start <= w[1].time.start); + } +} + +#[test] +fn pipeline_resegment_overlap_disabled_no_secondaries() { + let segs = vec![ + raw(0.0, 1.0, 0, true), + raw(0.0, 1.0, 1, true), + raw(2.0, 3.0, 0, false), + ]; + let cfg = PipelineConfig { + profile: Profile::Custom, + resegment_overlap: false, + min_speech_secs: 0.0, + max_gap_secs: 0.0, + ..PipelineConfig::default() + }; + let p = Pipeline::builder() + .config(cfg) + .with_segmenter(Box::new(ConstSegmenter(segs))) + .with_embedder(Box::new(AxisEmbedder { + dim: 8, + axis_picker: axis_picker_constant(0), + })) + .with_clusterer(Box::new(PerSampleClusterer { labels: vec![0] })) + .build() + .expect("custom build"); + let r = p + .run(&vec![0.0_f32; 16000 * 4], SampleRate::new(16000).unwrap()) + .unwrap(); + assert_eq!(r.num_speakers, 1); +}