aria_kit/speaker/
turns.rs

1//! Who is speaking, cut from the audio itself: turns of speech between
2//! pauses, each judged as the owner's voice or another's. The audio sets the
3//! clock, not the transcriber, so a decoder that finalizes late or in big
4//! pieces cannot merge two voices into one turn.
5
6use std::ops::Range;
7
8use super::{Label, SpeakerName};
9use crate::voice_activity::{Progress, SpeechMap};
10
11mod feed;
12mod gate;
13mod tracker;
14
15pub use feed::{Ahead, TurnFeed};
16pub use gate::{Admission, Decision, Gate, Released};
17pub use tracker::{NoSegmenter, OwnerTrackingError, VoiceprintTracker};
18
19/// Public items here stay unnameable outside this module.
20mod sealed {
21    pub trait Sealed {}
22}
23
24/// Says who spoke when in one recording, reading where anyone spoke from a
25/// [`SpeechMap`]. A recording's audio arrives in order, in pieces of any
26/// size; the map must have scored a frame before the tracker hears it.
27///
28/// Only this crate implements it, since a [`Timeline`] is built here.
29pub trait SpeakerTracker: sealed::Sealed {
30    /// Where the next sample heard sits.
31    fn heard_to(&self) -> u64;
32    /// Takes the audio after what it has heard.
33    fn hear(&mut self, speech: &SpeechMap, audio: &[f32]);
34    /// Gives up on the audio before `at`, which it will never hear.
35    fn skip_to(&mut self, speech: &SpeechMap, at: u64);
36    /// Judges what is still open. Call once the recording ends.
37    fn finish(&mut self);
38    fn timeline(&self) -> &Timeline;
39    /// Forgets the timeline before `at`, or before the open turn's start
40    /// when that comes first. Returns how far the speech map may be
41    /// forgotten: never past what the tracker still reads from it.
42    fn forget_before(&mut self, at: u64) -> u64;
43}
44
45/// One run of speech and whose voice it is.
46#[derive(Debug, Clone, PartialEq)]
47struct Turn {
48    /// Stream samples, through the pause that ended it. Grows while open.
49    span: Range<u64>,
50    judgment: Judgment,
51}
52
53#[derive(Debug, Clone, PartialEq)]
54struct LabeledTurn {
55    span: Range<u64>,
56    label: Option<Label>,
57}
58
59/// The embedding one check computed over a span.
60#[derive(Debug, Clone, PartialEq)]
61struct Embedding {
62    span: Range<u64>,
63    vector: Vec<f32>,
64}
65
66#[derive(Debug, Clone, Copy, PartialEq, Eq)]
67enum Judgment {
68    /// Too little heard to tell.
69    Pending,
70    Unverified,
71    Owner,
72    Other,
73}
74
75/// Whose voice the latest turn is, as judged so far.
76#[derive(Debug, Clone, PartialEq, Eq)]
77pub enum Speaker {
78    /// Owner admission is undecided, though the voice may already have a name.
79    Pending(Option<SpeakerName>),
80    Unverified,
81    /// The owner's, and the name they are enrolled under.
82    Owner(Option<SpeakerName>),
83    /// Another voice, named when it is an enrolled guest's.
84    Other(Option<SpeakerName>),
85}
86
87/// The turns of one recording so far, and how far into it they reach.
88#[derive(Debug, Clone, Default)]
89pub struct Timeline {
90    turns: Vec<Turn>,
91    // Kept apart from `turns`: a label change must not judge a turn or reset
92    // its evidence.
93    labels: Vec<LabeledTurn>,
94    /// Kept only by a tracker that numbers voices, so a session can
95    /// relabel them once it ends.
96    embeddings: Vec<Embedding>,
97    reached: u64,
98    /// Whether the last turn is still open.
99    open: bool,
100    stable_to: u64,
101    progress: Progress,
102}
103
104/// How much of a span lies in admitted turns and how much in rejected ones.
105#[derive(Debug, Clone, Copy, PartialEq, Eq)]
106struct Tally {
107    admitted: u64,
108    rejected: u64,
109}
110
111impl Tally {
112    /// No judged turn lies under the span.
113    fn is_empty(self) -> bool {
114        self.admitted == 0 && self.rejected == 0
115    }
116
117    fn mostly_admitted(self) -> bool {
118        self.admitted > self.rejected
119    }
120
121    fn wholly_admitted(self) -> bool {
122        self.admitted > 0 && self.rejected == 0
123    }
124}
125
126impl Timeline {
127    fn admission(&self, span: &Range<u64>) -> Option<bool> {
128        self.tally(span).map(Tally::mostly_admitted)
129    }
130
131    pub fn wholly_admitted(&self, span: &Range<u64>) -> Option<bool> {
132        if self.is_broken() || span.start > span.end {
133            return None;
134        }
135        self.tally(span).map(Tally::wholly_admitted)
136    }
137
138    fn admitted_so_far(&self, span: &Range<u64>) -> bool {
139        self.tally_stable_to(span, u64::MAX)
140            .is_some_and(Tally::wholly_admitted)
141    }
142
143    fn tally(&self, span: &Range<u64>) -> Option<Tally> {
144        self.tally_stable_to(span, self.stable_to)
145    }
146
147    /// [`Self::tally`], with the open turn stable up to `stable_to`.
148    fn tally_stable_to(&self, span: &Range<u64>, stable_to: u64) -> Option<Tally> {
149        // A word timed at a single instant still sits in a turn.
150        let span = span.start..span.end.max(span.start.saturating_add(1));
151        if span.end > self.reached {
152            return (self.progress == Progress::Finished).then_some(Tally {
153                admitted: 0,
154                rejected: 0,
155            });
156        }
157        let open = self.open.then(|| self.turns.len() - 1);
158        let mut tally = Tally {
159            admitted: 0,
160            rejected: 0,
161        };
162        for (index, turn) in self.turns.iter().enumerate() {
163            let covered = overlap(&span, &turn.span);
164            if covered == 0 {
165                continue;
166            }
167            if open == Some(index) && span.end > stable_to {
168                return None;
169            }
170            match turn.judgment {
171                Judgment::Pending => return None,
172                Judgment::Owner | Judgment::Unverified => tally.admitted += covered,
173                Judgment::Other => tally.rejected += covered,
174            }
175        }
176        Some(tally)
177    }
178
179    /// The end of the last whole frame the tracker has heard.
180    pub fn reached(&self) -> u64 {
181        self.reached
182    }
183
184    /// Forgets turns, labels and embeddings that ended before `at`, keeping
185    /// the latest turn and the latest label. Only for sessions that never
186    /// ask about the whole recording.
187    pub(crate) fn forget_before(&mut self, at: u64) {
188        let ended = |span: &Range<u64>| span.end <= at;
189        let turns = self
190            .turns
191            .iter()
192            .take_while(|turn| ended(&turn.span))
193            .count();
194        self.turns
195            .drain(..turns.min(self.turns.len().saturating_sub(1)));
196        let labels = self
197            .labels
198            .iter()
199            .take_while(|turn| ended(&turn.span))
200            .count();
201        self.labels
202            .drain(..labels.min(self.labels.len().saturating_sub(1)));
203        self.embeddings.retain(|embedding| !ended(&embedding.span));
204    }
205
206    pub fn is_broken(&self) -> bool {
207        self.progress == Progress::Broken
208    }
209
210    fn break_off(&mut self) {
211        self.progress = Progress::Broken;
212    }
213
214    /// A timeline nothing can be judged by: every segment held against it
215    /// is dropped.
216    pub fn broken() -> Self {
217        Self {
218            progress: Progress::Broken,
219            ..Self::default()
220        }
221    }
222
223    /// The enrolled voice of the turn that covers most of `span`.
224    pub fn speaker(&self, span: &Range<u64>) -> Option<&SpeakerName> {
225        self.label(span)?.enrolled()
226    }
227
228    /// The unit-norm embedding of `span`: every check's embedding over it,
229    /// each weighted by how much of `span` it covers.
230    fn embedding(&self, span: &Range<u64>) -> Option<Vec<f32>> {
231        if self.is_broken() {
232            return None;
233        }
234        let mut sum = Vec::new();
235        for embedding in &self.embeddings {
236            let weight = overlap(span, &embedding.span) as f32;
237            if weight == 0.0 {
238                continue;
239            }
240            sum.resize(embedding.vector.len(), 0.0);
241            for (sum, value) in sum.iter_mut().zip(&embedding.vector) {
242                *sum += value * weight;
243            }
244        }
245        super::normalized(sum)
246    }
247
248    /// The label of the turn that covers most of `span`.
249    fn label(&self, span: &Range<u64>) -> Option<&Label> {
250        if self.is_broken() {
251            return None;
252        }
253        self.labels
254            .iter()
255            .filter(|turn| overlap(span, &turn.span) > 0)
256            .max_by_key(|turn| overlap(span, &turn.span))?
257            .label
258            .as_ref()
259    }
260
261    /// Whose voice the latest turn is. `None` before anyone speaks.
262    pub fn speaking(&self) -> Option<Speaker> {
263        if self.is_broken() {
264            return None;
265        }
266        let turn = self.turns.last()?;
267        let name = self
268            .labels
269            .last()
270            .and_then(|turn| turn.label.as_ref()?.enrolled().cloned());
271        Some(match turn.judgment {
272            Judgment::Pending => Speaker::Pending(name),
273            Judgment::Unverified => Speaker::Unverified,
274            Judgment::Owner => Speaker::Owner(name),
275            Judgment::Other => Speaker::Other(name),
276        })
277    }
278
279    /// The name covering the longest span of admitted owner audio.
280    pub fn owner_name(&self) -> Option<&SpeakerName> {
281        if self.is_broken() {
282            return None;
283        }
284        self.labels
285            .iter()
286            .filter_map(|labeled| {
287                let name = labeled.label.as_ref()?.enrolled()?;
288                let owner_samples: u64 = self
289                    .turns
290                    .iter()
291                    .filter(|turn| turn.judgment == Judgment::Owner)
292                    .map(|turn| overlap(&labeled.span, &turn.span))
293                    .sum();
294                (owner_samples > 0).then_some((owner_samples, name))
295            })
296            .max_by_key(|(samples, _)| *samples)
297            .map(|(_, name)| name)
298    }
299
300    pub fn admitted_is_owner(&self) -> bool {
301        self.progress == Progress::Finished
302            && self
303                .turns
304                .iter()
305                .any(|turn| turn.judgment == Judgment::Owner)
306            && self
307                .turns
308                .iter()
309                .all(|turn| matches!(turn.judgment, Judgment::Owner | Judgment::Other))
310    }
311
312    /// The audio of turns that may be typed: the owner's, and those too
313    /// short to tell.
314    pub fn admitted_audio(&self, samples: &[f32], from: u64) -> Vec<f32> {
315        self.audio_judged(samples, from, |judgment| {
316            matches!(judgment, Judgment::Owner | Judgment::Unverified)
317        })
318    }
319
320    /// The audio of turns judged the owner's, and nothing else: what may
321    /// leave this computer.
322    pub fn owner_audio(&self, samples: &[f32], from: u64) -> Vec<f32> {
323        self.audio_judged(samples, from, |judgment| judgment == Judgment::Owner)
324    }
325
326    /// Whether every turn over `span` is the owner's, once decided.
327    pub fn wholly_owner(&self, span: &Range<u64>) -> Option<bool> {
328        let decided = self.wholly_admitted(span)?;
329        Some(
330            decided
331                && self
332                    .turns
333                    .iter()
334                    .filter(|turn| overlap(span, &turn.span) > 0)
335                    .all(|turn| turn.judgment == Judgment::Owner),
336        )
337    }
338
339    fn audio_judged(
340        &self,
341        samples: &[f32],
342        from: u64,
343        keep: impl Fn(Judgment) -> bool,
344    ) -> Vec<f32> {
345        if self.is_broken() {
346            return Vec::new();
347        }
348        self.turns
349            .iter()
350            .filter(|turn| keep(turn.judgment))
351            .flat_map(|turn| {
352                let index =
353                    |position: u64| (position.saturating_sub(from) as usize).min(samples.len());
354                &samples[index(turn.span.start)..index(turn.span.end)]
355            })
356            .copied()
357            .collect()
358    }
359}
360
361fn overlap(a: &Range<u64>, b: &Range<u64>) -> u64 {
362    a.end.min(b.end).saturating_sub(a.start.max(b.start))
363}
aria_kit/speaker/turns.rs

aria_kit/speaker/
turns.rs
