440 lines
16 KiB
Rust
440 lines
16 KiB
Rust
//! Speaker diarization composed from existing in-crate components.
|
||
//!
|
||
//! Pipeline (~the pyannote-community-1 shape, simpler):
|
||
//! 1. Silero V5 VAD → per-32ms speech probability → speech intervals
|
||
//! 2. WavLM-SV → 512-d x-vector per ~2-second window inside each interval
|
||
//! 3. Agglomerative average-linkage clustering on cosine distance →
|
||
//! speaker labels for every window
|
||
//! 4. Merge adjacent same-speaker windows back into time segments
|
||
//!
|
||
//! Pure candle, zero new deps. Uses [`crate::silero_vad::SileroVad`] (Phase 11)
|
||
//! and [`crate::wavlm_sv::WavLmSv`] (Phase 5d). Run as a sidecar binary or
|
||
//! library call during data prep — diarization is offline work that doesn't
|
||
//! belong in the live voice loop.
|
||
//!
|
||
//! Reference: ~80% of the WhisperX / pyannote pipeline without the ort
|
||
//! runtime risk (see `docs/sesame_gap_analysis.md` and the runtime-conflict
|
||
//! notes in `personal_voice_training_guide.md`).
|
||
|
||
use crate::error::{CsmError, Result};
|
||
use crate::silero_vad::SileroVad;
|
||
use crate::wavlm_sv::WavLmSv;
|
||
use candle_core::Device;
|
||
|
||
/// Sample rate the diarizer expects on its input. Both Silero V5 and
|
||
/// WavLM-SV in this crate operate at 16 kHz; the caller is responsible for
|
||
/// resampling 24 kHz CSM-side audio down before calling [`Diarizer::diarize`].
|
||
pub const SAMPLE_RATE: u32 = 16_000;
|
||
|
||
/// Silero V5 chunk = 512 samples = 32 ms at 16 kHz. Mirror the constant
|
||
/// from `silero_vad.rs` so we don't recompute it.
|
||
const VAD_CHUNK_MS: f32 = 32.0;
|
||
|
||
/// One speaker-attributed time interval.
|
||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||
pub struct DiarizedSegment {
|
||
pub start_s: f32,
|
||
pub end_s: f32,
|
||
pub speaker: usize,
|
||
}
|
||
|
||
/// Configuration for the full pipeline.
|
||
#[derive(Debug, Clone)]
|
||
pub struct DiarizationConfig {
|
||
/// VAD probability above which a 32 ms chunk is considered speech.
|
||
/// Silero V5 is well-calibrated; 0.5 is a safe default.
|
||
pub vad_threshold: f32,
|
||
/// Minimum continuous speech length (seconds) for a region to be kept.
|
||
/// Filters out cough / lip-smack / single-word interjections that
|
||
/// would produce noisy embeddings.
|
||
pub min_speech_s: f32,
|
||
/// Maximum silence (seconds) inside a speech interval before it is split.
|
||
/// Tolerates short pauses inside a single utterance.
|
||
pub max_silence_s: f32,
|
||
/// WavLM-SV window length in seconds. 2 s is the sweet spot —
|
||
/// long enough for a discriminative x-vector, short enough that
|
||
/// time resolution survives.
|
||
pub window_s: f32,
|
||
/// Hop between consecutive WavLM-SV windows. 1 s = 50 % overlap.
|
||
pub hop_s: f32,
|
||
/// Cosine-distance threshold above which two clusters are NOT merged.
|
||
/// Lower = more clusters (over-segmentation); higher = fewer clusters
|
||
/// (collapsed speakers). 0.5 is a literature default for WavLM-SV.
|
||
pub cluster_threshold: f32,
|
||
/// When `Some(k)`, force exactly `k` clusters and ignore the threshold.
|
||
/// `None` = auto-detect via the threshold.
|
||
pub n_speakers: Option<usize>,
|
||
/// Drop output segments shorter than this (after merge). Removes
|
||
/// 1-window blips that survived clustering.
|
||
pub min_segment_s: f32,
|
||
}
|
||
|
||
impl Default for DiarizationConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
vad_threshold: 0.5,
|
||
min_speech_s: 0.5,
|
||
max_silence_s: 0.3,
|
||
window_s: 2.0,
|
||
hop_s: 1.0,
|
||
cluster_threshold: 0.5,
|
||
n_speakers: None,
|
||
min_segment_s: 0.5,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Owns the two underlying models. Construct once, reuse across many calls.
|
||
pub struct Diarizer {
|
||
pub vad: SileroVad,
|
||
pub sv: WavLmSv,
|
||
device: Device,
|
||
}
|
||
|
||
impl Diarizer {
|
||
pub fn new(vad: SileroVad, sv: WavLmSv, device: Device) -> Self {
|
||
Self { vad, sv, device }
|
||
}
|
||
|
||
/// Run the full pipeline. `samples_16k` must be 16 kHz mono f32.
|
||
pub fn diarize(
|
||
&mut self,
|
||
samples_16k: &[f32],
|
||
cfg: &DiarizationConfig,
|
||
) -> Result<Vec<DiarizedSegment>> {
|
||
// 1. VAD → speech intervals
|
||
let probs = self
|
||
.vad
|
||
.forward_audio(samples_16k, &self.device)
|
||
.map_err(|e| CsmError::Config(format!("diarize: vad: {e}")))?;
|
||
let intervals = vad_intervals(&probs, cfg);
|
||
if intervals.is_empty() {
|
||
return Ok(Vec::new());
|
||
}
|
||
tracing::debug!("diarize: {} speech intervals", intervals.len());
|
||
|
||
// 2. WavLM-SV embeddings per window inside each interval. VAD
|
||
// interval bounds are chunk-quantized and may exceed
|
||
// `samples_16k.len()`; clamp before slicing.
|
||
let mut windows: Vec<WindowEmb> = Vec::new();
|
||
let win_samples = (cfg.window_s * SAMPLE_RATE as f32) as usize;
|
||
let hop_samples = (cfg.hop_s * SAMPLE_RATE as f32) as usize;
|
||
for &(s_idx, e_idx) in intervals.iter() {
|
||
let s = s_idx.min(samples_16k.len());
|
||
let e = e_idx.min(samples_16k.len());
|
||
if e <= s {
|
||
continue;
|
||
}
|
||
// Interval shorter than one window: take the whole interval as a
|
||
// single sub-window if it's at least half a window long. This is
|
||
// the common case for short utterances.
|
||
if e - s < win_samples {
|
||
if e - s >= win_samples / 2 {
|
||
let chunk = &samples_16k[s..e];
|
||
let emb = self.sv.embed_samples(chunk, &self.device)?;
|
||
let l2 = l2_normalize(&emb);
|
||
windows.push(WindowEmb {
|
||
start_s: s as f32 / SAMPLE_RATE as f32,
|
||
end_s: e as f32 / SAMPLE_RATE as f32,
|
||
emb: l2,
|
||
});
|
||
}
|
||
continue;
|
||
}
|
||
let mut t = s;
|
||
while t + win_samples <= e {
|
||
let chunk = &samples_16k[t..t + win_samples];
|
||
let emb = self.sv.embed_samples(chunk, &self.device)?;
|
||
let l2 = l2_normalize(&emb);
|
||
windows.push(WindowEmb {
|
||
start_s: t as f32 / SAMPLE_RATE as f32,
|
||
end_s: (t + win_samples) as f32 / SAMPLE_RATE as f32,
|
||
emb: l2,
|
||
});
|
||
t += hop_samples;
|
||
}
|
||
// Tail: grab one last full-width window flush against `e` to
|
||
// capture the trailing speech if the hop didn't already cover it.
|
||
if t < e && e - t >= win_samples / 2 {
|
||
let start = e.saturating_sub(win_samples);
|
||
let chunk = &samples_16k[start..e];
|
||
let emb = self.sv.embed_samples(chunk, &self.device)?;
|
||
let l2 = l2_normalize(&emb);
|
||
windows.push(WindowEmb {
|
||
start_s: start as f32 / SAMPLE_RATE as f32,
|
||
end_s: e as f32 / SAMPLE_RATE as f32,
|
||
emb: l2,
|
||
});
|
||
}
|
||
}
|
||
if windows.is_empty() {
|
||
return Ok(Vec::new());
|
||
}
|
||
tracing::debug!("diarize: {} embedding windows", windows.len());
|
||
|
||
// 3. Agglomerative clustering
|
||
let labels = agglomerative_cluster(
|
||
&windows.iter().map(|w| w.emb.as_slice()).collect::<Vec<_>>(),
|
||
cfg.cluster_threshold,
|
||
cfg.n_speakers,
|
||
);
|
||
|
||
// 4. Merge adjacent same-speaker windows
|
||
let mut segments = Vec::new();
|
||
let mut cur_speaker = labels[0];
|
||
let mut cur_start = windows[0].start_s;
|
||
let mut cur_end = windows[0].end_s;
|
||
for (i, w) in windows.iter().enumerate().skip(1) {
|
||
if labels[i] == cur_speaker && w.start_s <= cur_end + 0.01 {
|
||
cur_end = w.end_s.max(cur_end);
|
||
} else {
|
||
if cur_end - cur_start >= cfg.min_segment_s {
|
||
segments.push(DiarizedSegment {
|
||
start_s: cur_start,
|
||
end_s: cur_end,
|
||
speaker: cur_speaker,
|
||
});
|
||
}
|
||
cur_speaker = labels[i];
|
||
cur_start = w.start_s;
|
||
cur_end = w.end_s;
|
||
}
|
||
}
|
||
if cur_end - cur_start >= cfg.min_segment_s {
|
||
segments.push(DiarizedSegment {
|
||
start_s: cur_start,
|
||
end_s: cur_end,
|
||
speaker: cur_speaker,
|
||
});
|
||
}
|
||
Ok(segments)
|
||
}
|
||
}
|
||
|
||
struct WindowEmb {
|
||
start_s: f32,
|
||
end_s: f32,
|
||
emb: Vec<f32>,
|
||
}
|
||
|
||
/// Convert per-chunk Silero V5 probabilities into sample-index intervals
|
||
/// `(start, end)`. Internally chunk-quantized; the returned bounds are
|
||
/// exact-multiple-of-`CHUNK_SAMPLES`.
|
||
fn vad_intervals(probs: &[f32], cfg: &DiarizationConfig) -> Vec<(usize, usize)> {
|
||
let chunk_samples = (VAD_CHUNK_MS * SAMPLE_RATE as f32 / 1000.0) as usize;
|
||
let max_silence_chunks = (cfg.max_silence_s * 1000.0 / VAD_CHUNK_MS).round() as usize;
|
||
let min_speech_chunks = (cfg.min_speech_s * 1000.0 / VAD_CHUNK_MS).ceil() as usize;
|
||
|
||
// Smooth small gaps: turn isolated low-prob chunks inside speech
|
||
// into speech if surrounding `max_silence_chunks` is mostly speech.
|
||
let mut active: Vec<bool> = probs.iter().map(|p| *p >= cfg.vad_threshold).collect();
|
||
let mut last_speech: Option<usize> = None;
|
||
for i in 0..active.len() {
|
||
if active[i] {
|
||
if let Some(prev) = last_speech
|
||
&& i - prev <= max_silence_chunks + 1
|
||
{
|
||
active[(prev + 1)..i].fill(true);
|
||
}
|
||
last_speech = Some(i);
|
||
}
|
||
}
|
||
|
||
// Coalesce contiguous true runs into intervals.
|
||
let mut intervals = Vec::new();
|
||
let mut i = 0;
|
||
while i < active.len() {
|
||
if active[i] {
|
||
let start = i;
|
||
while i < active.len() && active[i] {
|
||
i += 1;
|
||
}
|
||
let len = i - start;
|
||
if len >= min_speech_chunks {
|
||
intervals.push((start * chunk_samples, i * chunk_samples));
|
||
}
|
||
} else {
|
||
i += 1;
|
||
}
|
||
}
|
||
intervals
|
||
}
|
||
|
||
fn l2_normalize(v: &[f32]) -> Vec<f32> {
|
||
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt().max(1e-9);
|
||
v.iter().map(|x| x / norm).collect()
|
||
}
|
||
|
||
/// Cosine *distance* between two L2-normalized vectors: `1 - dot(a, b)`.
|
||
/// Range `[0, 2]`; identical = 0; orthogonal = 1; opposite = 2.
|
||
fn cosine_distance_normed(a: &[f32], b: &[f32]) -> f32 {
|
||
let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
|
||
1.0 - dot
|
||
}
|
||
|
||
/// Agglomerative average-linkage clustering on cosine distance.
|
||
///
|
||
/// Each input vector is a starting cluster. At each step, find the two
|
||
/// clusters with the smallest pairwise average distance and merge them.
|
||
/// Stop when:
|
||
/// - `n_speakers` is `Some(k)` and we have `k` clusters, OR
|
||
/// - the smallest distance exceeds `threshold` (auto-detect mode)
|
||
///
|
||
/// Returns one label per input vector. Labels are remapped to be a dense
|
||
/// `0..n_clusters` range in order of first appearance.
|
||
fn agglomerative_cluster(embs: &[&[f32]], threshold: f32, n_speakers: Option<usize>) -> Vec<usize> {
|
||
let n = embs.len();
|
||
if n == 0 {
|
||
return Vec::new();
|
||
}
|
||
if n == 1 {
|
||
return vec![0];
|
||
}
|
||
|
||
// Each cluster carries: its membership (Vec of input idx) and the
|
||
// current label. We work with cluster ids 0..n that get unioned.
|
||
let mut clusters: Vec<Vec<usize>> = (0..n).map(|i| vec![i]).collect();
|
||
|
||
// Pairwise distance cache, keyed by sorted (smaller_id, larger_id).
|
||
// We rebuild after each merge — n is small (~ 10²) for typical
|
||
// diarization, so O(n²) per step is fine.
|
||
loop {
|
||
if clusters.len() == 1 {
|
||
break;
|
||
}
|
||
if let Some(k) = n_speakers
|
||
&& clusters.len() <= k
|
||
{
|
||
break;
|
||
}
|
||
// Find the closest pair using average linkage.
|
||
let mut best: Option<(usize, usize, f32)> = None;
|
||
for i in 0..clusters.len() {
|
||
for j in (i + 1)..clusters.len() {
|
||
let mut sum = 0.0f32;
|
||
let mut count = 0u32;
|
||
for &a in &clusters[i] {
|
||
for &b in &clusters[j] {
|
||
sum += cosine_distance_normed(embs[a], embs[b]);
|
||
count += 1;
|
||
}
|
||
}
|
||
let d = sum / count as f32;
|
||
if best.map(|(_, _, bd)| d < bd).unwrap_or(true) {
|
||
best = Some((i, j, d));
|
||
}
|
||
}
|
||
}
|
||
let (i, j, d) = best.unwrap();
|
||
if n_speakers.is_none() && d > threshold {
|
||
break;
|
||
}
|
||
// Merge j into i, remove j (j > i so popping order is safe).
|
||
let merged_j = clusters.remove(j);
|
||
clusters[i].extend(merged_j);
|
||
}
|
||
|
||
// Build per-input label mapping. Remap to dense 0..k in order of
|
||
// first appearance in the input.
|
||
let mut raw_label = vec![0usize; n];
|
||
for (cid, members) in clusters.iter().enumerate() {
|
||
for &m in members {
|
||
raw_label[m] = cid;
|
||
}
|
||
}
|
||
let mut remap = std::collections::HashMap::new();
|
||
let mut next_label = 0usize;
|
||
let mut out = Vec::with_capacity(n);
|
||
for &r in &raw_label {
|
||
let lab = *remap.entry(r).or_insert_with(|| {
|
||
let l = next_label;
|
||
next_label += 1;
|
||
l
|
||
});
|
||
out.push(lab);
|
||
}
|
||
out
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn cosine_distance_identical_is_zero() {
|
||
let a = l2_normalize(&[1.0, 2.0, 3.0]);
|
||
let d = cosine_distance_normed(&a, &a);
|
||
assert!(d.abs() < 1e-6, "identical vectors → distance ≈ 0, got {d}");
|
||
}
|
||
|
||
#[test]
|
||
fn cosine_distance_orthogonal_is_one() {
|
||
let a = l2_normalize(&[1.0, 0.0]);
|
||
let b = l2_normalize(&[0.0, 1.0]);
|
||
let d = cosine_distance_normed(&a, &b);
|
||
assert!((d - 1.0).abs() < 1e-6, "orthogonal → distance ≈ 1, got {d}");
|
||
}
|
||
|
||
#[test]
|
||
fn agglomerative_two_clusters_merges_within_threshold() {
|
||
// Two well-separated clusters of 3 vectors each.
|
||
let cluster_a = [
|
||
l2_normalize(&[1.0, 0.0, 0.0]),
|
||
l2_normalize(&[0.95, 0.1, 0.0]),
|
||
l2_normalize(&[0.9, 0.0, 0.1]),
|
||
];
|
||
let cluster_b = [
|
||
l2_normalize(&[0.0, 1.0, 0.0]),
|
||
l2_normalize(&[0.0, 0.95, 0.1]),
|
||
l2_normalize(&[0.05, 0.9, 0.0]),
|
||
];
|
||
let mut owned = Vec::new();
|
||
owned.extend(cluster_a.iter().cloned());
|
||
owned.extend(cluster_b.iter().cloned());
|
||
let refs: Vec<&[f32]> = owned.iter().map(|v| v.as_slice()).collect();
|
||
let labels = agglomerative_cluster(&refs, 0.5, None);
|
||
// First three should share a label; last three should share the other.
|
||
assert_eq!(labels[0], labels[1]);
|
||
assert_eq!(labels[0], labels[2]);
|
||
assert_eq!(labels[3], labels[4]);
|
||
assert_eq!(labels[3], labels[5]);
|
||
assert_ne!(labels[0], labels[3]);
|
||
}
|
||
|
||
#[test]
|
||
fn agglomerative_n_speakers_forces_k_clusters() {
|
||
// 4 well-separated points in different directions; threshold-mode
|
||
// would give 4 clusters, but n_speakers=2 should collapse them
|
||
// to exactly 2.
|
||
let pts = [
|
||
l2_normalize(&[1.0, 0.0, 0.0, 0.0]),
|
||
l2_normalize(&[0.0, 1.0, 0.0, 0.0]),
|
||
l2_normalize(&[0.0, 0.0, 1.0, 0.0]),
|
||
l2_normalize(&[0.0, 0.0, 0.0, 1.0]),
|
||
];
|
||
let refs: Vec<&[f32]> = pts.iter().map(|v| v.as_slice()).collect();
|
||
let labels = agglomerative_cluster(&refs, 0.1, Some(2));
|
||
let unique: std::collections::HashSet<_> = labels.iter().collect();
|
||
assert_eq!(unique.len(), 2, "n_speakers=2 should yield 2 clusters");
|
||
}
|
||
|
||
#[test]
|
||
fn vad_intervals_drops_too_short_speech() {
|
||
// 32 ms chunks, threshold 0.5. Pattern: 0.1 0.9 0.1 0.1 0.9 0.9 0.9 0.9 0.1
|
||
// = one isolated speech chunk (skipped, below min_speech) then 4
|
||
// contiguous (kept). With min_speech_s=0.1 (≈3 chunks at 32ms).
|
||
let probs = [0.1, 0.9, 0.1, 0.1, 0.9, 0.9, 0.9, 0.9, 0.1];
|
||
let cfg = DiarizationConfig {
|
||
vad_threshold: 0.5,
|
||
min_speech_s: 0.1,
|
||
max_silence_s: 0.0,
|
||
..DiarizationConfig::default()
|
||
};
|
||
let intervals = vad_intervals(&probs, &cfg);
|
||
assert_eq!(intervals.len(), 1, "should drop the isolated speech chunk");
|
||
// 4 chunks × 512 samples = 2048 samples
|
||
let (s, e) = intervals[0];
|
||
assert_eq!(e - s, 4 * 512);
|
||
}
|
||
}
|