sicada_decode/
frontier.rs

1//! What both decoders keep: the graph states alive at one frame, and the beam.
2//!
3//! A frontier maps a graph state to the best cost of reaching it at this frame,
4//! plus one `u32` the decoder is free to use: a backpointer in
5//! [`viterbi`](crate::viterbi), a lattice state in [`lattice`](crate::lattice).
6//! Both decoders prune identically, so the beam lives here rather than in each.
7
8use rustc_hash::FxHashMap;
9
10/// The cost of reaching a graph state at this frame, and one word of whatever
11/// the decoder needs alongside it.
12#[derive(Debug, Clone, Copy, PartialEq)]
13pub(crate) struct Token {
14    pub cost: f32,
15    pub aux: u32,
16}
17
18/// Not a valid `aux`, for the decoders that need an absent one.
19pub(crate) const NO_AUX: u32 = u32::MAX;
20
21/// How wide to search.
22#[derive(Debug, Clone, Copy, PartialEq)]
23pub struct DecodeOptions {
24    /// Tokens worse than the frame's best by more than this are dropped.
25    ///
26    /// In nats, since the scores are. `f32::INFINITY` searches exhaustively,
27    /// which the tests compare against.
28    pub beam: f32,
29    /// A hard cap on the tokens kept per frame, applied after `beam`.
30    ///
31    /// This is what bounds the work when the beam turns out to be too generous
32    /// for a confusable stretch of audio.
33    pub max_active: usize,
34    /// A floor under the same cap: the beam is never tightened below this many
35    /// tokens, so a confident frame does not narrow the search to nothing.
36    pub min_active: usize,
37}
38
39impl Default for DecodeOptions {
40    fn default() -> Self {
41        // Kaldi's defaults for a first pass, which are a reasonable starting
42        // point for any acoustic model scaled in nats.
43        Self {
44            beam: 16.0,
45            max_active: 7000,
46            min_active: 200,
47        }
48    }
49}
50
51impl DecodeOptions {
52    /// No beam and no cap: every path is kept.
53    ///
54    /// The oracle the decoders are tested against searches everything, so this
55    /// setting is the one that makes the two comparable.
56    pub fn exhaustive() -> Self {
57        Self {
58            beam: f32::INFINITY,
59            max_active: usize::MAX,
60            min_active: 0,
61        }
62    }
63}
64
65#[inline]
66pub(crate) fn relax_cost<S: std::hash::Hash + Eq>(
67    frontier: &mut FxHashMap<S, Token>,
68    state: S,
69    cost: f32,
70    aux: u32,
71) -> bool {
72    match frontier.get_mut(&state) {
73        Some(token) if token.cost <= cost => false,
74        Some(token) => {
75            *token = Token { cost, aux };
76            true
77        }
78        None => {
79            frontier.insert(state, Token { cost, aux });
80            true
81        }
82    }
83}
84
85// `costs` is caller-owned scratch space so pruning allocates nothing per frame.
86pub(crate) fn prune<S: std::hash::Hash + Eq>(
87    frontier: &mut FxHashMap<S, Token>,
88    opts: &DecodeOptions,
89    costs: &mut Vec<f32>,
90) -> f32 {
91    let best = frontier
92        .values()
93        .map(|token| token.cost)
94        .fold(f32::INFINITY, f32::min);
95    let mut cutoff = best + opts.beam;
96
97    // The cap is a *tighter* cutoff, found by asking which cost sits at the
98    // cap's rank. Selecting is O(n); sorting the frontier would not be.
99    let cap = opts.max_active.max(opts.min_active);
100    if frontier.len() > cap {
101        costs.clear();
102        costs.extend(frontier.values().map(|token| token.cost));
103        let rank = cap.min(costs.len() - 1);
104        let (_, &mut nth, _) = costs.select_nth_unstable_by(rank, |a, b| a.total_cmp(b));
105        cutoff = cutoff.min(nth);
106    }
107
108    if cutoff.is_finite() {
109        frontier.retain(|_, token| token.cost <= cutoff);
110    }
111    cutoff
112}
sicada_decode/frontier.rs

sicada_decode/
frontier.rs
