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

aria_kit/speaker/turns/
gate.rs
