sicada_decode/
nbest.rs

1//! Reading answers out of a compact lattice.
2//!
3//! *n*-best is only meaningful once the alignments are collapsed. Over a raw
4//! lattice the *n* best *paths* are usually *n* ways of cutting up the same
5//! sentence; over a compact one, which is deterministic on words, every path is
6//! a different word sequence, so the *n* best paths are *n* different answers.
7//! Nothing here does anything clever for that reason: it is
8//! [`shortest_path`] over an
9//! FST that has already been put in the right shape.
10//!
11//! What *is* here is the rescoring knob the whole compact-lattice semiring
12//! exists for: [`scale`] rescales the acoustic half against the graph half
13//! without touching the alignments, and the answers move accordingly.
14
15use sicada::algorithms::shortest_path::{ShortestPathOptions, shortest_path};
16use sicada::arc::{Arc, ArcLabel, ArcStateId, ArcTpl};
17use sicada::error::OpenFstError;
18use sicada::fst::{ExpandedFst, Fst, MutableFst};
19use sicada::fsts::vector_fst::VectorFst;
20use sicada::weight::Weight;
21
22use crate::compact_lattice_weight::CompactLatticeWeight;
23use crate::lattice_weight::LatticeWeight;
24
25/// One answer read off a compact lattice.
26#[derive(Debug, Clone, PartialEq)]
27pub struct Hypothesis<L: ArcLabel> {
28    /// The words, epsilons removed.
29    pub words: Vec<L>,
30    /// The cost, still in its two halves, with the frames each word spanned.
31    pub weight: CompactLatticeWeight<L>,
32}
33
34impl<L: ArcLabel> Hypothesis<L> {
35    /// The frames this hypothesis used, in order.
36    #[inline]
37    pub fn alignment(&self) -> &[L] {
38        self.weight.alignment()
39    }
40
41    /// The combined graph and acoustic cost.
42    #[inline]
43    pub fn cost(&self) -> f32 {
44        self.weight.weight().total()
45    }
46}
47
48/// The `n` best word sequences, cheapest first.
49///
50/// Fewer than `n` come back when the lattice has fewer paths. Two answers are
51/// distinct word sequences as long as `lattice` is deterministic on words,
52/// as [`determinize_lattice`](crate::compact::determinize_lattice) makes it;
53/// over a lattice that is not, this returns the `n` best *paths* and
54/// they may repeat themselves.
55pub fn n_best<L, S>(
56    lattice: &VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>>,
57    n: usize,
58) -> Result<Vec<Hypothesis<L>>, OpenFstError>
59where
60    L: ArcLabel,
61    S: ArcStateId,
62{
63    if n == 0 || lattice.start().is_none() {
64        return Ok(Vec::new());
65    }
66
67    let mut best: VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>> = VectorFst::new();
68    shortest_path(
69        lattice,
70        &mut best,
71        &ShortestPathOptions {
72            nshortest: n,
73            ..ShortestPathOptions::default()
74        },
75    )?;
76
77    let mut found = enumerate(&best);
78    // `shortest_path` returns the paths as one FST, in no particular order.
79    found.sort_by(|a, b| a.cost().total_cmp(&b.cost()));
80    found.truncate(n);
81    Ok(found)
82}
83
84/// Every path of an acyclic FST, as a hypothesis.
85///
86/// The result of `shortest_path` is a tree of at most `n` paths, so walking it
87/// exhaustively is bounded by what was asked for.
88fn enumerate<L, S>(fst: &VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>>) -> Vec<Hypothesis<L>>
89where
90    L: ArcLabel,
91    S: ArcStateId,
92{
93    let mut found = Vec::new();
94    let Some(start) = fst.start() else {
95        return found;
96    };
97    let zero = CompactLatticeWeight::<L>::zero();
98    let mut stack = vec![(start, Vec::new(), CompactLatticeWeight::<L>::one())];
99    while let Some((state, words, weight)) = stack.pop() {
100        let final_weight = fst.final_weight(state);
101        if final_weight.is_member() && final_weight != zero {
102            found.push(Hypothesis {
103                words: words.clone(),
104                weight: weight.times(&final_weight),
105            });
106        }
107        for arc in fst.arcs(state) {
108            let mut next = words.clone();
109            if arc.olabel() != L::epsilon() {
110                next.push(arc.olabel());
111            }
112            stack.push((arc.nextstate(), next, weight.times(arc.weight())));
113        }
114    }
115    found
116}
117
118/// Rescales the two halves of every weight.
119///
120/// `acoustic` multiplies the acoustic cost and `graph` the graph cost. This is
121/// what a compact lattice is *for*: the first pass decoded under one balance
122/// between the acoustic model and the language model, and a second pass can ask
123/// what the answer would have been under another, without decoding again.
124///
125/// Note that the halves are scaled where they sit, so the result is still a
126/// compact lattice and still says which frames each word spanned. Only what
127/// counts as *best* moves.
128pub fn scale<L, S>(
129    lattice: &mut VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>>,
130    acoustic: f32,
131    graph: f32,
132) where
133    L: ArcLabel,
134    S: ArcStateId,
135{
136    let rescale = |weight: &CompactLatticeWeight<L>| {
137        CompactLatticeWeight::new(
138            LatticeWeight::new(
139                graph * weight.weight().graph,
140                acoustic * weight.weight().acoustic,
141            ),
142            weight.alignment().iter().copied().collect(),
143        )
144    };
145
146    for state in 0..lattice.num_states() {
147        let state = S::from_usize(state);
148        let final_weight = lattice.final_weight(state);
149        if final_weight.is_member() && final_weight != CompactLatticeWeight::zero() {
150            lattice.set_final(state, rescale(&final_weight));
151        }
152        for arc in lattice.arcs_mut(state) {
153            arc.weight = rescale(&arc.weight);
154        }
155    }
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161    use sicada::arc::StdArc;
162    use sicada::fsts::vector_fst::StdVectorFst;
163    use sicada::properties::K_FST_PROPERTIES;
164    use sicada::weights::float_weight::TropicalWeight;
165
166    use crate::compact::{DeterminizeLatticeOptions, determinize_lattice};
167    use crate::dense::DenseFst;
168    use crate::lattice::{LatticeDecodeOptions, lattice_decode};
169
170    type Compact = VectorFst<ArcTpl<CompactLatticeWeight<i32>, i32, i32>>;
171
172    // Three symbols, each its own word, any sequence allowed.
173    fn graph() -> StdVectorFst {
174        let mut fst: StdVectorFst = VectorFst::new();
175        fst.add_state();
176        fst.set_start(0);
177        fst.set_final(0, TropicalWeight::one());
178        for label in 1..=3 {
179            fst.add_arc(0, StdArc::new(label, label * 10, TropicalWeight::one(), 0));
180        }
181        fst.properties(K_FST_PROPERTIES, true);
182        fst
183    }
184
185    fn compact_of(scores: &[f32], frames: usize) -> Compact {
186        let dense = DenseFst::<StdArc>::new(scores, frames, 3).unwrap();
187        let lattice = lattice_decode(&graph(), &dense, &LatticeDecodeOptions::exhaustive())
188            .unwrap()
189            .expect("a lattice");
190        determinize_lattice(&lattice, &DeterminizeLatticeOptions::default()).unwrap()
191    }
192
193    // Two frames, so nine word sequences, whose costs are the two frames'
194    // scores added.
195    const SCORES: [f32; 6] = [
196        0.0, 1.0, 2.0, //
197        0.0, 0.5, 3.0,
198    ];
199
200    #[test]
201    fn it_returns_distinct_word_sequences_cheapest_first() {
202        let compact = compact_of(&SCORES, 2);
203        let best = n_best(&compact, 4).expect("four answers");
204
205        assert_eq!(best.len(), 4);
206        let words: Vec<&[i32]> = best.iter().map(|h| h.words.as_slice()).collect();
207        assert_eq!(
208            words,
209            vec![
210                &[10, 10][..], // 0.0 + 0.0
211                &[10, 20][..], // 0.0 + 0.5
212                &[20, 10][..], // 1.0 + 0.0
213                &[20, 20][..], // 1.0 + 0.5
214            ]
215        );
216        for pair in best.windows(2) {
217            assert!(pair[0].cost() <= pair[1].cost(), "not sorted");
218        }
219        assert!((best[0].cost() - 0.0).abs() < 1e-6);
220        assert!((best[3].cost() - 1.5).abs() < 1e-6);
221    }
222
223    // The distinctness is the whole reason to determinize first, so it is
224    // asserted rather than assumed.
225    #[test]
226    fn no_two_answers_say_the_same_thing() {
227        let compact = compact_of(&SCORES, 2);
228        let best = n_best(&compact, 9).expect("nine answers");
229        assert_eq!(best.len(), 9, "three symbols over two frames");
230
231        let mut seen: Vec<&[i32]> = best.iter().map(|h| h.words.as_slice()).collect();
232        seen.sort_unstable();
233        let before = seen.len();
234        seen.dedup();
235        assert_eq!(seen.len(), before, "an answer was repeated");
236    }
237
238    #[test]
239    fn asking_for_more_than_there_are_returns_what_there_is() {
240        let compact = compact_of(&SCORES, 2);
241        assert_eq!(n_best(&compact, 100).unwrap().len(), 9);
242        assert!(n_best(&compact, 0).unwrap().is_empty());
243    }
244
245    #[test]
246    fn every_answer_carries_the_frames_it_used() {
247        let compact = compact_of(&SCORES, 2);
248        for hypothesis in n_best(&compact, 9).unwrap() {
249            assert_eq!(
250                hypothesis.alignment().len(),
251                2,
252                "two frames were decoded: {hypothesis:?}"
253            );
254            // Word 10n came from label n, so the alignment says the words back.
255            let from_alignment: Vec<i32> = hypothesis
256                .alignment()
257                .iter()
258                .map(|label| label * 10)
259                .collect();
260            assert_eq!(from_alignment, hypothesis.words);
261        }
262    }
263
264    // What keeping the halves apart is for: the same lattice, a different
265    // balance between the models, and a different answer, with no decoding.
266    #[test]
267    fn rescaling_the_acoustic_half_changes_which_answer_wins() {
268        // Word 10 is cheap acoustically and dear in the graph; word 20 is the
269        // other way round. Under equal weight, 10 wins by a hair.
270        let mut fst: StdVectorFst = VectorFst::new();
271        fst.add_state();
272        fst.set_start(0);
273        fst.set_final(0, TropicalWeight::one());
274        fst.add_arc(0, StdArc::new(1, 10, TropicalWeight(1.0), 0));
275        fst.add_arc(0, StdArc::new(2, 20, TropicalWeight(0.0), 0));
276        fst.properties(K_FST_PROPERTIES, true);
277
278        // One frame: symbol 1 costs 0.0, symbol 2 costs 1.2.
279        let scores = [0.0, 1.2, 9.0];
280        let dense = DenseFst::<StdArc>::new(&scores, 1, 3).unwrap();
281        let lattice = lattice_decode(&fst, &dense, &LatticeDecodeOptions::exhaustive())
282            .unwrap()
283            .unwrap();
284        let compact = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default()).unwrap();
285
286        // graph 1.0 + acoustic 0.0 = 1.0 beats graph 0.0 + acoustic 1.2.
287        assert_eq!(n_best(&compact, 1).unwrap()[0].words, vec![10]);
288
289        // Halve what the acoustic model has to say and the graph decides.
290        let mut quieter = compact.clone();
291        scale(&mut quieter, 0.5, 1.0);
292        assert_eq!(n_best(&quieter, 1).unwrap()[0].words, vec![20]);
293
294        // The alignments are untouched by the rescaling.
295        assert_eq!(n_best(&quieter, 1).unwrap()[0].alignment(), &[2]);
296    }
297
298    #[test]
299    fn scaling_leaves_the_alignments_alone() {
300        let mut compact = compact_of(&SCORES, 2);
301        let before: Vec<Vec<i32>> = n_best(&compact, 9)
302            .unwrap()
303            .iter()
304            .map(|h| h.alignment().to_vec())
305            .collect();
306        scale(&mut compact, 3.0, 2.0);
307        let after: Vec<Vec<i32>> = n_best(&compact, 9)
308            .unwrap()
309            .iter()
310            .map(|h| h.alignment().to_vec())
311            .collect();
312        assert_eq!(before.len(), after.len());
313        // Scaling is monotone here, so the order is unchanged as well.
314        assert_eq!(before, after);
315    }
316}
sicada_decode/nbest.rs

sicada_decode/
nbest.rs
