366 lines
14 KiB
Rust
366 lines
14 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>;
|
||
/// Partition already-loaded 16kHz mono samples into speaker spans. Phase 3
|
||
/// (FR-SPK) diarizes a *masked* far-side signal (the mic's "You" ranges
|
||
/// zeroed), so the caller loads and masks the samples itself rather than
|
||
/// pointing at the raw WAV.
|
||
fn diarize_samples(&self, samples: Vec<f32>) -> 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();
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Split-layout attribution (FR-SPK): decides per segment between "You" (mic
|
||
/// channel voice activity) and the far side's diarized speakers by comparing
|
||
/// the *total* voiced overlap on each channel, not by picking the single
|
||
/// longest span — a long far-side diarizer span could otherwise swallow a
|
||
/// segment the user spoke most of, showing their words under "Speaker N".
|
||
/// The mic channel is physically the user's voice alone, so channel evidence
|
||
/// outranks cluster evidence; ties go to "You" (mislabeling the user's own
|
||
/// words as someone else is the worse failure). A segment with no voiced
|
||
/// overlap on either channel keeps its prior label rather than guessing.
|
||
// ponytail: whole-segment labels — a segment genuinely containing both sides
|
||
// still gets one speaker; the upgrade path is transcribing each channel
|
||
// separately so segments can never mix voices.
|
||
pub fn assign_split(
|
||
segments: &mut [TranscriptSegment],
|
||
you_spans: &[(u64, u64)],
|
||
far_vad: &[(u64, u64)],
|
||
far_spans: &[SpeakerSpan],
|
||
) {
|
||
fn overlap(a0: u64, a1: u64, b0: u64, b1: u64) -> u64 {
|
||
a1.min(b1).saturating_sub(a0.max(b0))
|
||
}
|
||
for seg in segments.iter_mut() {
|
||
let mic_ms: u64 = you_spans
|
||
.iter()
|
||
.map(|&(s, e)| overlap(seg.start_ms, seg.end_ms, s, e))
|
||
.sum();
|
||
let far_ms: u64 = far_vad
|
||
.iter()
|
||
.map(|&(s, e)| overlap(seg.start_ms, seg.end_ms, s, e))
|
||
.sum();
|
||
if mic_ms == 0 && far_ms == 0 {
|
||
continue;
|
||
}
|
||
if mic_ms >= far_ms {
|
||
seg.speaker = "You".to_string();
|
||
} else if let Some(span) = far_spans
|
||
.iter()
|
||
.map(|sp| {
|
||
(
|
||
overlap(seg.start_ms, seg.end_ms, sp.start_ms, sp.end_ms),
|
||
sp,
|
||
)
|
||
})
|
||
.filter(|(o, _)| *o > 0)
|
||
.max_by_key(|(o, _)| *o)
|
||
.map(|(_, sp)| sp)
|
||
{
|
||
seg.speaker = span.speaker.clone();
|
||
}
|
||
// Far side voiced but no diarizer span overlaps (e.g. a sub-700ms span
|
||
// was filtered): keep the prior label rather than guess.
|
||
}
|
||
}
|
||
|
||
/// 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()))?;
|
||
self.diarize_samples(samples)
|
||
}
|
||
|
||
fn diarize_samples(&self, samples: Vec<f32>) -> Result<Vec<SpeakerSpan>, DiarError> {
|
||
// Guard the native FFI against degenerate input (a fully-masked far side
|
||
// is all zeros; a tiny buffer has no usable window). sherpa can crash
|
||
// rather than error on these, so skip cleanly instead.
|
||
if samples.len() < MIN_DIARIZE_SAMPLES {
|
||
return Ok(Vec::new());
|
||
}
|
||
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;
|
||
|
||
/// Below this many 16kHz samples (~1s) the diarizer skips the sherpa FFI: a
|
||
/// fully-masked far side is all zeros and a tiny buffer has no usable
|
||
/// segmentation window, either of which can crash the native code.
|
||
#[cfg(feature = "diarization")]
|
||
const MIN_DIARIZE_SAMPLES: usize = 16_000;
|
||
|
||
/// 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");
|
||
}
|
||
|
||
#[test]
|
||
fn split_labels_a_mic_dominant_segment_you_even_against_a_longer_far_span() {
|
||
// The user spoke 0-4000ms; the far side 4000-6000ms — but the far
|
||
// cluster span covers the whole window, so the old merged max-overlap
|
||
// pick handed the entire segment (the user's words included) to the
|
||
// far speaker. Channel totals must side with the mic instead.
|
||
let mut segments = vec![segment(0, 6000)];
|
||
let you = vec![(0u64, 4000u64)];
|
||
let far_vad = vec![(4000u64, 6000u64)];
|
||
let far_spans = vec![span(0, 6000, "S1")]; // long far cluster span
|
||
assign_split(&mut segments, &you, &far_vad, &far_spans);
|
||
assert_eq!(segments[0].speaker, "You");
|
||
}
|
||
|
||
#[test]
|
||
fn split_ties_go_to_you() {
|
||
let mut segments = vec![segment(0, 2000)];
|
||
let you = vec![(0u64, 1000u64)];
|
||
let far_vad = vec![(1000u64, 2000u64)];
|
||
let far_spans = vec![span(1000, 2000, "S1")];
|
||
assign_split(&mut segments, &you, &far_vad, &far_spans);
|
||
assert_eq!(segments[0].speaker, "You");
|
||
}
|
||
|
||
#[test]
|
||
fn split_assigns_the_best_far_span_when_the_far_side_dominates() {
|
||
let mut segments = vec![segment(0, 3000)];
|
||
let you = vec![(0u64, 500u64)];
|
||
let far_vad = vec![(500u64, 3000u64)];
|
||
let far_spans = vec![span(500, 1000, "S1"), span(1000, 3000, "S2")];
|
||
assign_split(&mut segments, &you, &far_vad, &far_spans);
|
||
assert_eq!(segments[0].speaker, "S2");
|
||
}
|
||
|
||
#[test]
|
||
fn split_keeps_the_prior_label_when_both_channels_are_silent() {
|
||
let mut segments = vec![segment(5000, 6000)];
|
||
segments[0].speaker = "S9".to_string();
|
||
assign_split(
|
||
&mut segments,
|
||
&[(0, 1000)],
|
||
&[(0, 1000)],
|
||
&[span(0, 1000, "S1")],
|
||
);
|
||
assert_eq!(segments[0].speaker, "S9");
|
||
}
|
||
|
||
#[test]
|
||
fn split_keeps_the_prior_label_when_far_is_voiced_but_no_far_span_overlaps() {
|
||
// Far VAD hears speech but every diarizer span was filtered (sub-700ms):
|
||
// don't guess a label.
|
||
let mut segments = vec![segment(0, 1000)];
|
||
segments[0].speaker = "S3".to_string();
|
||
assign_split(&mut segments, &[], &[(0, 1000)], &[span(2000, 3000, "S1")]);
|
||
assert_eq!(segments[0].speaker, "S3");
|
||
}
|
||
}
|
||
|
||
#[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(_))));
|
||
}
|
||
}
|