Phase 4

src-tauri/src/diarization/mod.rs

@@ -9,6 +9,8 @@
 use crate::models::{SpeakerSpan, TranscriptSegment};
 use std::path::Path;
 
+pub mod models;
+
 #[derive(Debug, thiserror::Error)]
 pub enum DiarError {
     #[error("model load failed: {0}")]
@@ -24,17 +26,167 @@ pub trait Diarizer: Send + Sync {
     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;
+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),
+            threshold: Some(0.5),
+            ..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> {
-        // T4.1: sherpa-onnx segmentation + embedding + clustering via FFI.
-        todo!("Phase 4 — diarize")
+    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()))?;
+        Ok(segments.into_iter().map(segment_to_span).collect())
     }
-    fn assign(&self, _segments: &mut [TranscriptSegment], _spans: &[SpeakerSpan]) {
-        // T4.2: timestamp-overlap alignment.
-        todo!("Phase 4 — assign speakers to segments")
+
+    fn assign(&self, segments: &mut [TranscriptSegment], spans: &[SpeakerSpan]) {
+        assign_by_overlap(segments, spans);
+    }
+}
+
+/// 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 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(_))));
     }
 }