aria_kit/speaker/turns/
gate.rs

1//! Which decoded words a session lets out, decided against the timeline.
2
3use std::collections::VecDeque;
4use std::ops::Range;
5use std::time::Duration;
6
7use super::{Tally, Timeline};
8use crate::audio::canonical_samples;
9use crate::speaker::Label;
10use crate::transcription::Heard;
11use crate::voice_activity::SpeechMap;
12
13/// Under [`Admission::Spoken`], a segment with less speech than this, whose
14/// segment before it had less too, is music or noise.
15const MIN_SPEECH_SAMPLES: u64 = canonical_samples(Duration::from_millis(200));
16
17/// Which words a session lets out. The tracker reads it to know what to
18/// judge, and the gate reads it to know what to release.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum Admission {
21    /// Only words in turns judged the owner's, or too short to judge.
22    OwnerOnly,
23    /// Every word, named when an enrolled voice was heard in it.
24    Everyone,
25    /// Every word spoken: a segment with no speech under it is dropped.
26    Spoken,
27}
28
29/// Holds decoded segments until the timeline decides them.
30pub struct Gate {
31    admission: Admission,
32    /// Where the recording starts in the stream.
33    from: u64,
34    held: VecDeque<Heard>,
35    provisional: Option<Provisional>,
36    /// Speech under the last segment judged by its speech. A segment short
37    /// of speech still passes when this was enough: decode windows lag the
38    /// words at a boundary.
39    previous_speech: u64,
40}
41
42/// The decoder's latest provisional text.
43struct Provisional {
44    text: String,
45    source_window: Option<Range<u64>>,
46    /// How far the decoder had read when it wrote the text.
47    read_to: u64,
48}
49
50/// A segment let out, and the voice it was heard in.
51pub struct Released {
52    pub heard: Heard,
53    pub speaker: Option<Label>,
54    /// The unit-norm embedding of its audio, when the tracker numbers voices.
55    pub embedding: Option<Vec<f32>>,
56}
57
58impl Released {
59    fn new(heard: Heard, timeline: &Timeline) -> Self {
60        Self {
61            speaker: timeline.label(&heard.window).cloned(),
62            embedding: timeline.embedding(&heard.window),
63            heard,
64        }
65    }
66}
67
68/// What may go out now. Segments still undecided stay held.
69#[derive(Default)]
70pub struct Decision {
71    pub released: Vec<Released>,
72    /// Whole segments kept out: in another voice, or with no speech.
73    pub dropped: Vec<Heard>,
74    /// The provisional text to show, empty to hide it. `None` while the
75    /// decoder has none since its last finalized segment.
76    pub preview: Option<String>,
77}
78
79impl Gate {
80    /// `from` is where the recording starts in the stream.
81    pub fn new(admission: Admission, from: u64) -> Self {
82        Self {
83            admission,
84            from,
85            held: VecDeque::new(),
86            provisional: None,
87            previous_speech: 0,
88        }
89    }
90
91    pub fn admission(&self) -> Admission {
92        self.admission
93    }
94
95    /// Where the oldest held segment's audio starts, while one is held.
96    pub fn undecided_from(&self) -> Option<u64> {
97        self.held.front().map(|heard| heard.window.start)
98    }
99
100    /// A finalized segment joins the queue. It replaces the provisional text it came from.
101    pub fn hold(&mut self, heard: Heard) {
102        self.provisional = None;
103        self.held.push_back(heard);
104    }
105
106    /// The decoder's latest provisional text replaces the last.
107    pub fn propose(&mut self, text: String, source_window: Option<Range<u64>>, read_to: u64) {
108        self.provisional = Some(Provisional {
109            text,
110            source_window,
111            read_to,
112        });
113    }
114
115    /// Decides what the timeline and the speech map allow out now. The
116    /// preview is judged before the release, against the segments held when
117    /// it was asked for. Without a speech map, every recording counts as
118    /// speech.
119    pub fn decide(&mut self, timeline: &Timeline, speech: Option<&SpeechMap>) -> Decision {
120        let mut decision = Decision {
121            preview: self.preview(timeline),
122            ..Decision::default()
123        };
124        match self.admission {
125            Admission::Everyone => decision.released.extend(
126                self.held
127                    .drain(..)
128                    .map(|heard| Released::new(heard, timeline)),
129            ),
130            Admission::OwnerOnly => self.release_owner(timeline, &mut decision),
131            Admission::Spoken => self.release_spoken(timeline, speech, &mut decision),
132        }
133        decision
134    }
135
136    /// Releases segments with speech under them, in order, once the speech
137    /// map has scored their audio.
138    fn release_spoken(
139        &mut self,
140        timeline: &Timeline,
141        speech: Option<&SpeechMap>,
142        decision: &mut Decision,
143    ) {
144        while let Some(heard) = self.held.pop_front() {
145            // A map that could not be judged counts as speech.
146            let spoken = match speech.filter(|speech| !speech.is_broken()) {
147                None => true,
148                Some(speech) => {
149                    if heard.window.end > speech.reached() && !speech.is_finished() {
150                        self.held.push_front(heard);
151                        break;
152                    }
153                    let heard_speech = speech.speech(&heard.window);
154                    let enough = heard_speech.max(self.previous_speech) >= MIN_SPEECH_SAMPLES;
155                    self.previous_speech = heard_speech;
156                    enough
157                }
158            };
159            if spoken {
160                decision.released.push(Released::new(heard, timeline));
161            } else {
162                tracing::debug!("a segment with no speech under it was dropped");
163                decision.dropped.push(heard);
164            }
165        }
166    }
167
168    /// Releases admitted words in order; broken tracking discards the queue.
169    fn release_owner(&mut self, timeline: &Timeline, decision: &mut Decision) {
170        while let Some(mut heard) = self.held.pop_front() {
171            if timeline.is_broken() {
172                tracing::debug!("segment not delivered: nothing can be judged");
173                decision.dropped.push(heard);
174                continue;
175            }
176            let Some(keep) = word_admissions(timeline, &heard, self.from) else {
177                self.held.push_front(heard);
178                break;
179            };
180            let Some(segment) = heard.segment.keeping(&keep) else {
181                tracing::debug!("segment in another voice not delivered");
182                decision.dropped.push(heard);
183                continue;
184            };
185            if segment.word_spans().is_some() {
186                heard.span = segment.span();
187                heard.window = segment.span().unwrap_or(heard.window);
188            }
189            heard.segment = segment;
190            decision.released.push(Released::new(heard, timeline));
191        }
192    }
193
194    fn preview(&self, timeline: &Timeline) -> Option<String> {
195        let provisional = self.provisional.as_ref()?;
196        Some(match self.admission {
197            Admission::Everyone | Admission::Spoken => provisional.text.clone(),
198            Admission::OwnerOnly => self
199                .owner_preview(timeline, provisional)
200                .unwrap_or_default(),
201        })
202    }
203
204    /// The held words and the provisional text after them, while every one of them
205    /// sits in the owner's speech so far.
206    fn owner_preview(&self, timeline: &Timeline, provisional: &Provisional) -> Option<String> {
207        let read_to = provisional.read_to;
208        if timeline.is_broken() || timeline.reached < read_to {
209            return None;
210        }
211        let fallback = self.from..read_to;
212        let source = provisional.source_window.as_ref().unwrap_or(&fallback);
213        if source.start < self.from || source.start > source.end || source.end > read_to {
214            return None;
215        }
216        let provisional_owner = provisional.text.is_empty() || timeline.admitted_so_far(source);
217        let held_owner = self
218            .held
219            .iter()
220            .all(|heard| timeline.admitted_so_far(&heard.window));
221        if !provisional_owner || !held_owner {
222            return None;
223        }
224        let words: Vec<&str> = self
225            .held
226            .iter()
227            .map(|heard| heard.segment.as_str())
228            .chain([provisional.text.as_str()])
229            .filter(|text| !text.is_empty())
230            .collect();
231        Some(words.join(" "))
232    }
233}
234
235fn word_admissions(timeline: &Timeline, heard: &Heard, from: u64) -> Option<Vec<bool>> {
236    let source = heard.segment.source_window();
237    match heard.segment.word_spans() {
238        Some(words) => words
239            .iter()
240            .map(|word| word_admission(timeline, word, source.as_ref(), from))
241            .collect(),
242        None => timeline
243            .wholly_admitted(&heard.window)
244            .map(|owner| vec![owner]),
245    }
246}
247
248fn word_admission(
249    timeline: &Timeline,
250    word: &Range<u64>,
251    source: Option<&Range<u64>>,
252    from: u64,
253) -> Option<bool> {
254    let admitted = timeline.admission(word)?;
255    if admitted
256        || word.start >= timeline.reached
257        || word.end > timeline.reached
258        || !timeline.tally(word).is_some_and(Tally::is_empty)
259    {
260        return Some(admitted);
261    }
262    let Some(source) = source.filter(|source| source.start >= from && covers(source, word)) else {
263        return Some(false);
264    };
265    // A coarse word timestamp can miss speech that its decode window contains.
266    timeline.wholly_admitted(source)
267}
268
269fn covers(source: &Range<u64>, word: &Range<u64>) -> bool {
270    source.start <= word.start && word.start < source.end && word.end <= source.end
271}
aria_kit/speaker/turns/gate.rs

aria_kit/speaker/turns/
gate.rs
