iamdoubz
a481f5dcbf
Phase 0 of the speaker-accuracy plan. sherpa fast-clustering threshold 0.5 over summed mic+loopback mono exploded a 2-speaker call into 83 clusters; raise to 0.7 (larger threshold -> fewer clusters) and drop sub-700ms spans whose embeddings are unstable and only churn labels in assign_by_overlap.
234 lines
8.4 KiB
Rust
234 lines
8.4 KiB
Rust
//! Speaker diarization (Phase 4, FR-SPK-*). sherpa-onnx: segmentation +
|
||
//! embedding + clustering, fully offline (ADR-0005).
|
||
//!
|
||
//! Runs as a post-stop pass over the recorded WAV (audio = source of truth),
|
||
//! then aligns speaker spans to transcript segments by timestamp overlap.
|
||
//! Speaker IDs are internal ("S1"…); names are applied at render time and never
|
||
//! destructively rewritten onto segments (FR-SPK-5).
|
||
|
||
use crate::models::{SpeakerSpan, TranscriptSegment};
|
||
use std::path::Path;
|
||
|
||
pub mod models;
|
||
pub mod voiceprint;
|
||
|
||
#[derive(Debug, thiserror::Error)]
|
||
pub enum DiarError {
|
||
#[error("model load failed: {0}")]
|
||
Load(String),
|
||
#[error("diarization failed: {0}")]
|
||
Run(String),
|
||
}
|
||
|
||
pub trait Diarizer: Send + Sync {
|
||
/// Partition audio into speaker spans.
|
||
fn diarize(&self, wav: &Path) -> Result<Vec<SpeakerSpan>, DiarError>;
|
||
/// Assign speaker labels to transcript segments by overlap with spans.
|
||
fn assign(&self, segments: &mut [TranscriptSegment], spans: &[SpeakerSpan]);
|
||
}
|
||
|
||
/// Labels each segment with whichever span overlaps it most, in milliseconds
|
||
/// (T4.2, FR-SPK-1). A segment with no overlapping span (e.g. it falls in a
|
||
/// gap between spans) keeps its prior speaker label — the "S1" placeholder
|
||
/// every segment starts with pre-diarization — rather than guessing.
|
||
///
|
||
/// Pure timestamp arithmetic, so it's engine-agnostic: every `Diarizer` impl
|
||
/// can share it instead of reimplementing overlap math.
|
||
pub fn assign_by_overlap(segments: &mut [TranscriptSegment], spans: &[SpeakerSpan]) {
|
||
for segment in segments.iter_mut() {
|
||
let best_span = spans
|
||
.iter()
|
||
.map(|span| {
|
||
let overlap_start = segment.start_ms.max(span.start_ms);
|
||
let overlap_end = segment.end_ms.min(span.end_ms);
|
||
(overlap_end.saturating_sub(overlap_start), span)
|
||
})
|
||
.filter(|(overlap, _)| *overlap > 0)
|
||
.max_by_key(|(overlap, _)| *overlap)
|
||
.map(|(_, span)| span);
|
||
|
||
if let Some(span) = best_span {
|
||
segment.speaker = span.speaker.clone();
|
||
}
|
||
}
|
||
}
|
||
|
||
/// sherpa-onnx-backed diarizer: pyannote segmentation + speaker-embedding +
|
||
/// fast clustering (ADR-0005, T4.1). `Diarize::compute` needs `&mut self`; it's
|
||
/// wrapped in a `Mutex` to satisfy `Diarizer: Sync` — diarization is a
|
||
/// once-per-meeting post-pass (never a hot path), so lock contention is moot.
|
||
#[cfg(feature = "diarization")]
|
||
pub struct SherpaDiarizer {
|
||
engine: std::sync::Mutex<sherpa_rs::diarize::Diarize>,
|
||
}
|
||
|
||
#[cfg(feature = "diarization")]
|
||
impl SherpaDiarizer {
|
||
pub fn new(segmentation_model: &Path, embedding_model: &Path) -> Result<Self, DiarError> {
|
||
let config = sherpa_rs::diarize::DiarizeConfig {
|
||
// A meeting's speaker count isn't known ahead of time: <= 0 tells
|
||
// sherpa-onnx to pick the cluster count itself from `threshold`
|
||
// instead of forcing a fixed number of speakers.
|
||
num_clusters: Some(-1),
|
||
// Tuning evidence: 0.5 over the summed mic+loopback mono WAV
|
||
// exploded a 2-speaker conversation into 83 clusters (overlapped
|
||
// speech + short pyannote chunks yield unstable embeddings). sherpa
|
||
// fast-clustering: larger threshold → fewer clusters. 0.7 lands a
|
||
// 2-person meeting near 2–4. ponytail: code constant, no settings
|
||
// knob until 0.7 proves insufficient in the field.
|
||
threshold: Some(0.7),
|
||
..Default::default()
|
||
};
|
||
let engine = sherpa_rs::diarize::Diarize::new(segmentation_model, embedding_model, config)
|
||
.map_err(|e| DiarError::Load(e.to_string()))?;
|
||
Ok(Self {
|
||
engine: std::sync::Mutex::new(engine),
|
||
})
|
||
}
|
||
}
|
||
|
||
#[cfg(feature = "diarization")]
|
||
impl Diarizer for SherpaDiarizer {
|
||
fn diarize(&self, wav: &Path) -> Result<Vec<SpeakerSpan>, DiarError> {
|
||
let samples =
|
||
crate::audio::read_wav_mono_16k(wav).map_err(|e| DiarError::Run(e.to_string()))?;
|
||
let mut engine = self.engine.lock().unwrap_or_else(|e| e.into_inner());
|
||
let segments = engine
|
||
.compute(samples, None)
|
||
.map_err(|e| DiarError::Run(e.to_string()))?;
|
||
// Drop sub-700ms spans: their embeddings are unstable and only cause
|
||
// label churn in assign_by_overlap (which keeps the prior label when no
|
||
// span overlaps). See MIN_SPAN_MS.
|
||
Ok(segments
|
||
.into_iter()
|
||
.map(segment_to_span)
|
||
.filter(span_long_enough)
|
||
.collect())
|
||
}
|
||
|
||
fn assign(&self, segments: &mut [TranscriptSegment], spans: &[SpeakerSpan]) {
|
||
assign_by_overlap(segments, spans);
|
||
}
|
||
}
|
||
|
||
/// Minimum span duration kept from a diarization pass. Sub-second pyannote
|
||
/// chunks carry unstable ERes2Net embeddings and only churn labels.
|
||
#[cfg(feature = "diarization")]
|
||
const MIN_SPAN_MS: u64 = 700;
|
||
|
||
/// A span is kept only if it lasts at least [`MIN_SPAN_MS`].
|
||
#[cfg(feature = "diarization")]
|
||
fn span_long_enough(span: &SpeakerSpan) -> bool {
|
||
span.end_ms.saturating_sub(span.start_ms) >= MIN_SPAN_MS
|
||
}
|
||
|
||
/// sherpa-onnx speaker indices are 0-based; WA's internal labels are 1-based ("S1"…).
|
||
#[cfg(feature = "diarization")]
|
||
fn segment_to_span(seg: sherpa_rs::diarize::Segment) -> SpeakerSpan {
|
||
SpeakerSpan {
|
||
start_ms: (seg.start.max(0.0) * 1000.0) as u64,
|
||
end_ms: (seg.end.max(0.0) * 1000.0) as u64,
|
||
speaker: format!("S{}", seg.speaker + 1),
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod overlap_tests {
|
||
use super::*;
|
||
|
||
fn segment(start_ms: u64, end_ms: u64) -> TranscriptSegment {
|
||
TranscriptSegment {
|
||
id: 0,
|
||
start_ms,
|
||
end_ms,
|
||
speaker: "S1".to_string(),
|
||
text: String::new(),
|
||
confidence: None,
|
||
interim: false,
|
||
}
|
||
}
|
||
|
||
fn span(start_ms: u64, end_ms: u64, speaker: &str) -> SpeakerSpan {
|
||
SpeakerSpan {
|
||
start_ms,
|
||
end_ms,
|
||
speaker: speaker.to_string(),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn assigns_the_span_that_overlaps_most() {
|
||
let mut segments = vec![segment(0, 1000), segment(1000, 2000)];
|
||
let spans = vec![span(0, 1000, "S1"), span(1000, 2000, "S2")];
|
||
assign_by_overlap(&mut segments, &spans);
|
||
assert_eq!(segments[0].speaker, "S1");
|
||
assert_eq!(segments[1].speaker, "S2");
|
||
}
|
||
|
||
#[test]
|
||
fn a_segment_straddling_two_spans_picks_the_larger_overlap() {
|
||
// 700ms in S1's span (300-1000), 300ms in S2's span (1000-1300).
|
||
let mut segments = vec![segment(300, 1300)];
|
||
let spans = vec![span(0, 1000, "S1"), span(1000, 2000, "S2")];
|
||
assign_by_overlap(&mut segments, &spans);
|
||
assert_eq!(segments[0].speaker, "S1");
|
||
}
|
||
|
||
#[test]
|
||
fn a_segment_with_no_overlapping_span_keeps_its_prior_label() {
|
||
let mut segments = vec![segment(5000, 6000)];
|
||
segments[0].speaker = "S9".to_string(); // distinct from any span below
|
||
let spans = vec![span(0, 1000, "S1")];
|
||
assign_by_overlap(&mut segments, &spans);
|
||
assert_eq!(segments[0].speaker, "S9"); // untouched, not overwritten with a guess
|
||
}
|
||
|
||
#[test]
|
||
fn no_spans_at_all_leaves_segments_untouched() {
|
||
let mut segments = vec![segment(0, 1000)];
|
||
assign_by_overlap(&mut segments, &[]);
|
||
assert_eq!(segments[0].speaker, "S1");
|
||
}
|
||
}
|
||
|
||
#[cfg(all(test, feature = "diarization"))]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn segment_to_span_maps_0_based_speaker_index_to_1_based_label() {
|
||
let seg = sherpa_rs::diarize::Segment {
|
||
start: 1.5,
|
||
end: 3.25,
|
||
speaker: 0,
|
||
};
|
||
let span = segment_to_span(seg);
|
||
assert_eq!(span.start_ms, 1500);
|
||
assert_eq!(span.end_ms, 3250);
|
||
assert_eq!(span.speaker, "S1");
|
||
}
|
||
|
||
#[test]
|
||
fn span_long_enough_filters_sub_minimum_spans() {
|
||
let short = SpeakerSpan {
|
||
start_ms: 0,
|
||
end_ms: MIN_SPAN_MS - 1,
|
||
speaker: "S1".to_string(),
|
||
};
|
||
let exact = SpeakerSpan {
|
||
start_ms: 0,
|
||
end_ms: MIN_SPAN_MS,
|
||
speaker: "S1".to_string(),
|
||
};
|
||
assert!(!span_long_enough(&short));
|
||
assert!(span_long_enough(&exact));
|
||
}
|
||
|
||
#[test]
|
||
fn new_surfaces_a_load_error_for_missing_models_instead_of_panicking() {
|
||
let result =
|
||
SherpaDiarizer::new(Path::new("no-such-seg.onnx"), Path::new("no-such-emb.onnx"));
|
||
assert!(matches!(result, Err(DiarError::Load(_))));
|
||
}
|
||
}
|