Skip to main content

sicada_decode/
compact.rs

1//! Collapsing a lattice's alignments, so each word sequence appears once.
2//!
3//! The lattice a decoder produces has one arc per frame. A word that took eight
4//! frames to say is eight arcs, and every way of placing its boundaries is a
5//! separate path, so the same sentence comes back dozens of times, differing
6//! only in where the words were cut. *n*-best over that is not *n* sentences,
7//! and rescoring it does the same work over and over.
8//!
9//! A *compact* lattice has one arc per word, with the frames it spanned moved
10//! into the weight ([`CompactLatticeWeight`]). Once they are in the weight,
11//! ordinary determinization over words does the collapsing: two arcs for the
12//! same word merge, and ⊕ keeps the better alignment instead of both.
13//!
14//! Upstream says the same thing, and says why gallic will not do
15//! (`fstext/determinize-lattice.h`):
16//!
17//! > We determinize this using acceptor determinization with epsilon removal.
18//! > […] `CompactLatticeWeightTpl` has a special kind of semiring where we
19//! > always take the string corresponding to the best cost […] and discard the
20//! > other. […] We couldn't use the Gallic weight for this, or it would die as
21//! > soon as it detected that the input FST was non-functional.
22//!
23//! A lattice is exactly that non-functional transducer: one word sequence, many
24//! alignments.
25//!
26//! What the algorithm needs beyond the semiring is the right *common divisor*.
27//! Determinization normalises each subset by dividing out what its members
28//! share, and for this weight that is the better cost together with the longest
29//! common **prefix** of the alignments, rather than ⊕, whose alignment belongs
30//! to one member and divides none of the others. That is the one piece
31//! [`CompactLatticeCommonDivisor`] supplies; the rest is sicada's own
32//! [`determinize_fsa`], which already takes a divisor because OpenFst's does.
33//!
34//! Kaldi writes its own determinizer instead, to prune as it goes and to keep
35//! the alignments in a shared trie rather than copying them. Neither is here.
36//! What is here is the outer half of the same idea: [`determinize_lattice_pruned`]
37//! narrows the lattice and tries again when determinization runs away, which is
38//! what Kaldi's wrapper does around its own. That much matters, because a bare
39//! CTC topology constrains nothing, so a thousand-frame lattice has
40//! astronomically many symbol sequences inside any generous beam, and
41//! determinizing it whole does not finish.
42
43use sicada::algorithms::connect::connect;
44use sicada::algorithms::determinize::{CommonDivisor, determinize_fsa};
45use sicada::algorithms::prune::{PruneOptions, prune as prune_fst};
46use sicada::algorithms::rmepsilon::rm_epsilon;
47use sicada::arc::{Arc, ArcLabel, ArcStateId, ArcTpl};
48use sicada::error::OpenFstError;
49use sicada::fst::{ExpandedFst, Fst, MutableFst};
50use sicada::fsts::vector_fst::VectorFst;
51use sicada::weight::Weight;
52
53use crate::compact_lattice_weight::{Alignment, CompactLatticeArc, CompactLatticeWeight};
54use crate::lattice_weight::LatticeWeight;
55
56/// A lattice with one arc per word, for a decoding graph over `A`.
57pub type CompactLattice<A> = VectorFst<CompactLatticeArc<A>>;
58
59/// What [`determinize_lattice`] may be asked to do differently.
60#[derive(Debug, Clone, Copy, PartialEq)]
61pub struct DeterminizeLatticeOptions {
62    /// How closely two subsets' weights must agree to count as the same subset.
63    pub delta: f32,
64    /// A cap on the states built, or `None` for no cap.
65    ///
66    /// Determinizing a lattice can produce far more states than it consumed,
67    /// and a decoder that is otherwise bounded should not become unbounded
68    /// here. Reaching the cap is reported rather than truncated: half a lattice
69    /// looks exactly like a whole one to everything downstream.
70    pub max_states: Option<usize>,
71}
72
73impl Default for DeterminizeLatticeOptions {
74    fn default() -> Self {
75        Self {
76            delta: 1.0 / 32.0,
77            max_states: Some(1 << 20),
78        }
79    }
80}
81
82/// What the members of a determinized subset share.
83///
84/// The better cost, and the longest common prefix of the alignments. ⊕ will not
85/// do: its alignment is one member's whole sequence, which does not divide any
86/// of the others, and determinization would find every subset unnormalisable.
87#[derive(Debug, Clone, Copy, Default)]
88pub struct CompactLatticeCommonDivisor;
89
90impl<L: ArcLabel> CommonDivisor<CompactLatticeWeight<L>> for CompactLatticeCommonDivisor {
91    fn divisor(
92        &self,
93        w1: &CompactLatticeWeight<L>,
94        w2: &CompactLatticeWeight<L>,
95    ) -> CompactLatticeWeight<L> {
96        // Zero contributes nothing, so the other one stands whole. This is the
97        // same convention `LabelCommonDivisor` follows, and it lets the divisor
98        // be folded over a subset starting from zero.
99        let zero = CompactLatticeWeight::zero();
100        match (w1 == &zero, w2 == &zero) {
101            (true, true) => return zero,
102            (true, false) => return w2.clone(),
103            (false, true) => return w1.clone(),
104            (false, false) => {}
105        }
106        let shared = w1
107            .alignment()
108            .iter()
109            .zip(w2.alignment())
110            .take_while(|(a, b)| a == b)
111            .map(|(a, _)| *a)
112            .collect();
113        CompactLatticeWeight::new(w1.weight().plus(w2.weight()), shared)
114    }
115}
116
117/// Moves each arc's input label into its weight, leaving the words on the arcs.
118///
119/// The result has the same shape as `lattice`; it is [`determinize_lattice`]
120/// that collapses it. An input epsilon contributes nothing to the alignment,
121/// having consumed no frame.
122///
123/// Written over the label and state-id types rather than over the graph's arc:
124/// a lattice mentions that arc only behind its associated types, which do not
125/// determine it, so a caller would have had to name it.
126pub fn to_compact<L, S>(
127    lattice: &VectorFst<ArcTpl<LatticeWeight, L, S>>,
128) -> VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>>
129where
130    L: ArcLabel,
131    S: ArcStateId,
132{
133    let mut compact: VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>> = VectorFst::new();
134    compact.reserve_states(lattice.num_states());
135    for _ in 0..lattice.num_states() {
136        compact.add_state();
137    }
138    if let Some(start) = lattice.start() {
139        compact.set_start(start);
140    }
141    compact.set_input_symbols(lattice.output_symbols());
142    compact.set_output_symbols(lattice.output_symbols());
143
144    for state in lattice.states() {
145        let final_weight = lattice.final_weight(state);
146        if final_weight.is_member() && final_weight != LatticeWeight::zero() {
147            compact.set_final(state, CompactLatticeWeight::from_weight(final_weight));
148        }
149        for arc in lattice.arcs(state) {
150            let mut alignment = Alignment::new();
151            if arc.ilabel() != L::epsilon() {
152                alignment.push(arc.ilabel());
153            }
154            compact.add_arc(
155                state,
156                ArcTpl::new(
157                    // An acceptor over words: what determinization merges on.
158                    arc.olabel(),
159                    arc.olabel(),
160                    CompactLatticeWeight::new(*arc.weight(), alignment),
161                    arc.nextstate(),
162                ),
163            );
164        }
165    }
166    compact
167}
168
169/// Rewrites `lattice` so that each word sequence appears once, with its best
170/// alignment.
171///
172/// # Errors
173///
174/// Reaching `opts.max_states` is an error, not a truncation.
175pub fn determinize_lattice<L, S>(
176    lattice: &VectorFst<ArcTpl<LatticeWeight, L, S>>,
177    opts: &DeterminizeLatticeOptions,
178) -> Result<VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>>, OpenFstError>
179where
180    L: ArcLabel,
181    S: ArcStateId,
182{
183    let mut compact = to_compact(lattice);
184
185    // A word-epsilon arc is one the determinization cannot merge on, so it has
186    // to go first. Removing it is also what carries its frames onto whatever
187    // word comes next, which is where they belong.
188    rm_epsilon(&mut compact, true)?;
189
190    let mut determinized: VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>> = VectorFst::new();
191    determinize_fsa(
192        &compact,
193        &mut determinized,
194        &CompactLatticeCommonDivisor,
195        opts.delta,
196        opts.max_states,
197    )?;
198    connect(&mut determinized);
199    Ok(determinized)
200}
201
202/// What [`determinize_lattice_pruned`] may be asked to do differently.
203#[derive(Debug, Clone, Copy, PartialEq)]
204pub struct PrunedDeterminizeOptions {
205    /// The beam to try first: paths worse than the best by more than this are
206    /// dropped before determinizing.
207    pub beam: f32,
208    /// What to multiply the beam by when an attempt runs away. Kaldi's wrapper
209    /// halves it.
210    pub beam_ratio: f32,
211    /// How many times to narrow and try again.
212    pub max_retries: usize,
213    /// Passed through to [`determinize_lattice`].
214    pub determinize: DeterminizeLatticeOptions,
215}
216
217impl Default for PrunedDeterminizeOptions {
218    fn default() -> Self {
219        Self {
220            beam: 8.0,
221            beam_ratio: 0.5,
222            max_retries: 6,
223            determinize: DeterminizeLatticeOptions::default(),
224        }
225    }
226}
227
228/// A collapsed lattice, and the beam it took to get one.
229#[derive(Debug, Clone)]
230pub struct PrunedLattice<L: ArcLabel, S: ArcStateId> {
231    /// The determinized compact lattice.
232    pub lattice: VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>>,
233    /// The beam actually used, which is `opts.beam` unless it had to narrow.
234    ///
235    /// A narrower beam preserves the best path but may remove alternatives that
236    /// a caller intended to rescore.
237    pub beam: f32,
238    /// How many attempts ran away before one finished.
239    pub narrowed: usize,
240}
241
242/// As [`determinize_lattice`], narrowing the lattice and trying again when an
243/// attempt runs past `max_states`.
244///
245/// Determinization can produce far more states than it consumes, and how many is
246/// not knowable in advance, since it depends on how many distinct symbol
247/// sequences the beam admits and the lattice does not say. The answer is
248/// therefore to try, and to ask for less if it runs away. Upstream's wrapper
249/// does the same around its own determinizer.
250///
251/// # Errors
252///
253/// The last attempt's error, once the retries are used up.
254pub fn determinize_lattice_pruned<L, S>(
255    lattice: &VectorFst<ArcTpl<LatticeWeight, L, S>>,
256    opts: &PrunedDeterminizeOptions,
257) -> Result<PrunedLattice<L, S>, OpenFstError>
258where
259    L: ArcLabel,
260    S: ArcStateId,
261{
262    let mut beam = opts.beam;
263    let mut last: Option<OpenFstError> = None;
264
265    for narrowed in 0..=opts.max_retries {
266        let mut narrowed_lattice = lattice.clone();
267        if beam.is_finite() {
268            prune_fst(
269                &mut narrowed_lattice,
270                &PruneOptions::threshold(LatticeWeight::new(beam, 0.0)),
271            )?;
272        }
273        // Every path is gone, which no narrower beam will undo.
274        if narrowed_lattice.start().is_none() {
275            return Err(OpenFstError::InvalidOperation(format!(
276                "determinize_lattice_pruned: a beam of {beam} left no path at all"
277            )));
278        }
279
280        match determinize_lattice(&narrowed_lattice, &opts.determinize) {
281            Ok(lattice) => {
282                return Ok(PrunedLattice {
283                    lattice,
284                    beam,
285                    narrowed,
286                });
287            }
288            // Any failure is retried: a narrower beam can only make the
289            // determinization smaller, so there is nothing to gain by telling
290            // one kind of failure from another here.
291            Err(error) => {
292                last = Some(error);
293                beam *= opts.beam_ratio;
294            }
295        }
296    }
297
298    Err(last.unwrap_or_else(|| {
299        OpenFstError::InvalidOperation("determinize_lattice_pruned: no attempts were made".into())
300    }))
301}
302
303#[cfg(test)]
304mod tests {
305    use super::*;
306    use rustc_hash::FxHashMap;
307    use sicada::arc::StdArc;
308    use sicada::fsts::vector_fst::StdVectorFst;
309    use sicada::properties::{K_ACYCLIC, K_FST_PROPERTIES, K_I_DETERMINISTIC};
310    use sicada::weights::float_weight::TropicalWeight;
311
312    use crate::dense::DenseFst;
313    use crate::lattice::{Lattice, LatticeDecodeOptions, lattice_decode};
314
315    struct Rng(u64);
316
317    impl Rng {
318        fn next(&mut self) -> u64 {
319            self.0 ^= self.0 << 13;
320            self.0 ^= self.0 >> 7;
321            self.0 ^= self.0 << 17;
322            self.0
323        }
324        fn below(&mut self, n: usize) -> usize {
325            (self.next() % n as u64) as usize
326        }
327        fn cost(&mut self) -> f32 {
328            self.below(256) as f32 / 16.0
329        }
330    }
331
332    // Every word sequence the FST accepts, with the best cost for each.
333    //
334    // Only usable on an acyclic FST, as the lattices below are: their states
335    // are `(frame, graph state)` and no arc goes back a frame.
336    fn word_sequences<W, F>(fst: &F) -> FxHashMap<Vec<i32>, f32>
337    where
338        W: Weight,
339        F: Fst<ArcTpl<W, i32, i32>>,
340        W: TotalCost,
341    {
342        let mut found: FxHashMap<Vec<i32>, f32> = FxHashMap::default();
343        let Some(start) = fst.start() else {
344            return found;
345        };
346        let mut stack = vec![(start, Vec::<i32>::new(), 0.0f32)];
347        while let Some((state, words, cost)) = stack.pop() {
348            let final_weight = fst.final_weight(state);
349            if final_weight.is_member() && final_weight != W::zero() {
350                let total = cost + final_weight.total_cost();
351                found
352                    .entry(words.clone())
353                    .and_modify(|best| *best = best.min(total))
354                    .or_insert(total);
355            }
356            for arc in fst.arcs(state) {
357                let mut next = words.clone();
358                if arc.olabel() != 0 {
359                    next.push(arc.olabel());
360                }
361                stack.push((arc.nextstate(), next, cost + arc.weight().total_cost()));
362            }
363        }
364        found
365    }
366
367    // The one number a path's weight comes down to, whichever of the two
368    // lattice semirings it is in.
369    trait TotalCost {
370        fn total_cost(&self) -> f32;
371    }
372
373    impl TotalCost for LatticeWeight {
374        fn total_cost(&self) -> f32 {
375            self.total()
376        }
377    }
378
379    impl TotalCost for CompactLatticeWeight<i32> {
380        fn total_cost(&self) -> f32 {
381            self.weight().total()
382        }
383    }
384
385    // A graph with no input epsilons, so the lattice comes out acyclic and can
386    // be enumerated. Output epsilons are plentiful, so the alignments
387    // collapse.
388    fn random_graph(rng: &mut Rng, symbols: usize) -> StdVectorFst {
389        let states = 1 + rng.below(4);
390        let mut graph: StdVectorFst = VectorFst::new();
391        for _ in 0..states {
392            graph.add_state();
393        }
394        graph.set_start(0);
395        for from in 0..states as i32 {
396            for _ in 0..1 + rng.below(3) {
397                let ilabel = 1 + rng.below(symbols) as i32;
398                // Half the arcs say no word at all, which is how one word comes
399                // to span several frames.
400                let olabel = if rng.below(2) == 0 {
401                    0
402                } else {
403                    10 * (1 + rng.below(2) as i32)
404                };
405                let to = rng.below(states) as i32;
406                graph.add_arc(
407                    from,
408                    StdArc::new(ilabel, olabel, TropicalWeight(rng.cost()), to),
409                );
410            }
411            if rng.below(2) == 0 {
412                graph.set_final(from, TropicalWeight(rng.cost()));
413            }
414        }
415        graph.properties(K_FST_PROPERTIES, true);
416        graph
417    }
418
419    fn decode(
420        graph: &StdVectorFst,
421        scores: &[f32],
422        frames: usize,
423        symbols: usize,
424    ) -> Option<Lattice<StdArc>> {
425        let dense = DenseFst::<StdArc>::new(scores, frames, symbols).unwrap();
426        lattice_decode(graph, &dense, &LatticeDecodeOptions::exhaustive()).unwrap()
427    }
428
429    // Compare complete word-to-cost maps, not only the best path.
430    #[test]
431    fn it_keeps_every_word_sequence_at_the_same_cost() {
432        let symbols = 3;
433        let mut rng = Rng(0x00D1_5EA5_E1A1_2345);
434        let mut compared = 0;
435
436        for round in 0..150 {
437            let graph = random_graph(&mut rng, symbols);
438            let frames = 1 + rng.below(4);
439            let scores: Vec<f32> = (0..frames * symbols).map(|_| rng.cost()).collect();
440            let Some(lattice) = decode(&graph, &scores, frames, symbols) else {
441                continue;
442            };
443
444            let before = word_sequences(&lattice);
445            let compact = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default())
446                .expect("a determinization");
447            let after = word_sequences(&compact);
448
449            assert_eq!(
450                before.len(),
451                after.len(),
452                "round {round}: {} sequences became {}",
453                before.len(),
454                after.len()
455            );
456            for (words, cost) in &before {
457                let found = after
458                    .get(words)
459                    .unwrap_or_else(|| panic!("round {round}: {words:?} went missing"));
460                assert!(
461                    (found - cost).abs() < 1e-3,
462                    "round {round}: {words:?} cost {found}, was {cost}"
463                );
464            }
465            compared += 1;
466        }
467
468        assert!(compared > 80, "only {compared} rounds produced a lattice");
469    }
470
471    // Every path's `(cost, alignment)`, folded per word sequence by the
472    // semiring's own ⊕.
473    //
474    // That fold *is* the specification: ⊕ keeps the better cost's alignment
475    // whole, and breaks a tie on the shorter one. Writing the oracle as the
476    // fold rather than as "take the minimum cost" makes it check the alignment
477    // half too, ties included.
478    fn best_per_sequence<W, F>(fst: &F) -> FxHashMap<Vec<i32>, CompactLatticeWeight<i32>>
479    where
480        W: Weight + AsCompact,
481        F: Fst<ArcTpl<W, i32, i32>>,
482    {
483        let mut found: FxHashMap<Vec<i32>, CompactLatticeWeight<i32>> = FxHashMap::default();
484        let Some(start) = fst.start() else {
485            return found;
486        };
487        let one = CompactLatticeWeight::<i32>::one();
488        let mut stack = vec![(start, Vec::<i32>::new(), one)];
489        while let Some((state, words, weight)) = stack.pop() {
490            let final_weight = fst.final_weight(state);
491            if final_weight.is_member() && final_weight != W::zero() {
492                let whole = weight.times(&final_weight.as_compact());
493                found
494                    .entry(words.clone())
495                    .and_modify(|best| *best = best.plus(&whole))
496                    .or_insert(whole);
497            }
498            for arc in fst.arcs(state) {
499                let mut next = words.clone();
500                if arc.olabel() != 0 {
501                    next.push(arc.olabel());
502                }
503                stack.push((
504                    arc.nextstate(),
505                    next,
506                    weight.times(&arc.weight().as_compact()),
507                ));
508            }
509        }
510        found
511    }
512
513    // A path's weight as the compact semiring sees it.
514    //
515    // Only the compact weight implements it: the enumeration runs over
516    // [`to_compact`]'s output, which is the raw lattice with each arc's frame
517    // already moved into its weight and nothing else changed.
518    trait AsCompact {
519        fn as_compact(&self) -> CompactLatticeWeight<i32>;
520    }
521
522    impl AsCompact for CompactLatticeWeight<i32> {
523        fn as_compact(&self) -> Self {
524            self.clone()
525        }
526    }
527
528    // The alignment half is the point, so it is checked and not only the cost:
529    // each word sequence must come back with the *best-scoring* alignment the
530    // lattice had for it, chosen by the same ⊕ the semiring defines.
531    #[test]
532    fn it_keeps_the_best_alignment_for_each_word_sequence() {
533        let symbols = 3;
534        let mut rng = Rng(0x00AB_CDEF_0123_4567);
535        let mut compared = 0;
536
537        for round in 0..150 {
538            let graph = random_graph(&mut rng, symbols);
539            let frames = 1 + rng.below(4);
540            let scores: Vec<f32> = (0..frames * symbols).map(|_| rng.cost()).collect();
541            let Some(lattice) = decode(&graph, &scores, frames, symbols) else {
542                continue;
543            };
544
545            // `to_compact` moves each arc's frame into its weight without
546            // changing anything else, so enumerating *that* gives the same
547            // paths as the raw lattice with their alignments attached.
548            let expected = best_per_sequence(&to_compact(&lattice));
549            let compact = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default())
550                .expect("a determinization");
551            let found = best_per_sequence(&compact);
552
553            assert_eq!(expected.len(), found.len(), "round {round}");
554            for (words, want) in &expected {
555                let got = found
556                    .get(words)
557                    .unwrap_or_else(|| panic!("round {round}: {words:?} went missing"));
558                assert_eq!(
559                    got.alignment(),
560                    want.alignment(),
561                    "round {round}: {words:?} kept the wrong alignment"
562                );
563                assert!(
564                    (got.weight().total() - want.weight().total()).abs() < 1e-3,
565                    "round {round}: {words:?} cost {got} vs {want}"
566                );
567            }
568            compared += 1;
569        }
570
571        assert!(compared > 80, "only {compared} rounds produced a lattice");
572    }
573
574    // What "compact" buys: one arc per word, and one path per word sequence.
575    #[test]
576    fn each_word_sequence_is_a_single_path() {
577        let symbols = 3;
578        let mut rng = Rng(0x0FED_CBA9_8765_4321);
579        let mut collapsed = 0;
580
581        for round in 0..150 {
582            let graph = random_graph(&mut rng, symbols);
583            let frames = 1 + rng.below(4);
584            let scores: Vec<f32> = (0..frames * symbols).map(|_| rng.cost()).collect();
585            let Some(lattice) = decode(&graph, &scores, frames, symbols) else {
586                continue;
587            };
588
589            let compact = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default())
590                .expect("a determinization");
591            if compact.start().is_none() {
592                continue;
593            }
594
595            let props = compact.properties(K_I_DETERMINISTIC | K_ACYCLIC, true);
596            assert_ne!(
597                props & K_I_DETERMINISTIC,
598                0,
599                "round {round}: not deterministic, so a word sequence has two paths"
600            );
601
602            // The raw lattice usually had several paths per sequence; the
603            // compact one has exactly as many paths as sequences.
604            let paths_before = count_paths(&lattice);
605            let sequences = word_sequences(&compact).len();
606            let paths_after = count_paths(&compact);
607            assert_eq!(paths_after, sequences, "round {round}");
608            if paths_before > paths_after {
609                collapsed += 1;
610            }
611        }
612
613        assert!(collapsed > 40, "nothing collapsed in {collapsed} rounds");
614    }
615
616    fn count_paths<W, F>(fst: &F) -> usize
617    where
618        W: Weight,
619        F: Fst<ArcTpl<W, i32, i32>>,
620    {
621        let Some(start) = fst.start() else {
622            return 0;
623        };
624        let mut stack = vec![start];
625        let mut paths = 0;
626        while let Some(state) = stack.pop() {
627            let final_weight = fst.final_weight(state);
628            if final_weight.is_member() && final_weight != W::zero() {
629                paths += 1;
630            }
631            for arc in fst.arcs(state) {
632                stack.push(arc.nextstate());
633            }
634        }
635        paths
636    }
637
638    // The frames a word spanned have to survive. A second pass rescores them,
639    // and an alignment is read from them.
640    #[test]
641    fn the_alignment_travels_with_the_word() {
642        // Three frames, one symbol each, all mapping to the same word 10.
643        let mut graph: StdVectorFst = VectorFst::new();
644        graph.add_state();
645        graph.set_start(0);
646        graph.set_final(0, TropicalWeight::one());
647        // Label 1 says the word, labels 2 and 3 continue it silently.
648        graph.add_arc(0, StdArc::new(1, 10, TropicalWeight::one(), 0));
649        graph.add_arc(0, StdArc::new(2, 0, TropicalWeight::one(), 0));
650        graph.add_arc(0, StdArc::new(3, 0, TropicalWeight::one(), 0));
651        graph.properties(K_FST_PROPERTIES, true);
652
653        // Frame 0 wants label 1, frames 1 and 2 want labels 2 and 3.
654        let scores = [
655            0.0, 9.0, 9.0, //
656            9.0, 0.0, 9.0, //
657            9.0, 9.0, 0.0,
658        ];
659        let lattice = decode(&graph, &scores, 3, 3).expect("a lattice");
660        let compact = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default()).unwrap();
661
662        // One word leaves the start: the two silent frames are gone as arcs,
663        // their frames folded into the weights around them.
664        let start = compact.start().expect("a start");
665        let arcs: Vec<_> = compact.arcs(start).collect();
666        assert_eq!(arcs.len(), 1, "one word, one arc");
667        assert_eq!(arcs[0].olabel(), 10);
668
669        // The frames after the last word land on the final weight, which is
670        // Kaldi's arrangement too, so the path's alignment is the ⊗ of the
671        // arc's and the final state's, in that order.
672        let mut best: VectorFst<ArcTpl<CompactLatticeWeight<i32>, i32, i32>> = VectorFst::new();
673        sicada::algorithms::shortest_path::shortest_path(
674            &compact,
675            &mut best,
676            &sicada::algorithms::shortest_path::ShortestPathOptions::default(),
677        )
678        .expect("a best path");
679        let (words, weight) =
680            sicada::string::string_fst_to_output_labels(&best).expect("a single path");
681
682        assert_eq!(words, vec![10], "one word was said");
683        assert_eq!(
684            weight.alignment(),
685            &[1, 2, 3],
686            "the three frames the word spanned"
687        );
688        assert!(weight.weight().total().abs() < 1e-6, "{weight}");
689    }
690
691    // A compact lattice is an FST like any other, so it writes and reads like
692    // one, and the header it writes carries Kaldi's names, which is why the
693    // weight's `type_name` was matched to upstream.
694    #[test]
695    fn it_writes_and_reads_back_as_an_fst() {
696        use sicada::fst::{FstReadOptions, FstWriteOptions};
697        use std::io::Write as _;
698
699        let scores = [0.0, 1.0, 2.0, 0.5, 0.25, 3.0];
700        let mut rng = Rng(0x0011_2233_4455_6677);
701        let graph = random_graph(&mut rng, 3);
702        let Some(lattice) = decode(&graph, &scores, 2, 3) else {
703            return;
704        };
705        let compact = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default()).unwrap();
706        if compact.start().is_none() {
707            return;
708        }
709
710        assert_eq!(
711            <ArcTpl<CompactLatticeWeight<i32>, i32, i32> as Arc>::type_name().as_str(),
712            "compactlattice44",
713            "the name the header records"
714        );
715
716        let mut bytes = Vec::new();
717        compact
718            .write(&mut bytes, &FstWriteOptions::default())
719            .expect("written");
720
721        let directory = tempfile::tempdir().expect("a directory");
722        let path = directory.path().join("lattice.fst");
723        std::fs::File::create(&path)
724            .unwrap()
725            .write_all(&bytes)
726            .unwrap();
727
728        let mut file = std::fs::File::open(&path).unwrap();
729        let read: VectorFst<ArcTpl<CompactLatticeWeight<i32>, i32, i32>> =
730            VectorFst::read(&mut file, &FstReadOptions::default()).expect("read back");
731
732        assert_eq!(read.num_states(), compact.num_states());
733        assert_eq!(read.start(), compact.start());
734        for state in compact.states() {
735            assert_eq!(read.final_weight(state), compact.final_weight(state));
736            assert_eq!(
737                read.arcs(state).collect::<Vec<_>>(),
738                compact.arcs(state).collect::<Vec<_>>(),
739                "state {state}"
740            );
741        }
742    }
743
744    // A failed attempt must retry with a narrower beam and report that beam.
745    #[test]
746    fn it_narrows_the_beam_rather_than_giving_up() {
747        let symbols = 3;
748        let mut rng = Rng(0x0777_8888_9999_AAAA);
749        let graph = random_graph(&mut rng, symbols);
750        let frames = 4;
751        let scores: Vec<f32> = (0..frames * symbols).map(|_| rng.cost()).collect();
752        let Some(lattice) = decode(&graph, &scores, frames, symbols) else {
753            return;
754        };
755
756        // A cap of one state cannot be met at any beam, so every retry runs
757        // away and the last error comes back rather than a truncated lattice.
758        let impossible = determinize_lattice_pruned(
759            &lattice,
760            &PrunedDeterminizeOptions {
761                determinize: DeterminizeLatticeOptions {
762                    max_states: Some(1),
763                    ..DeterminizeLatticeOptions::default()
764                },
765                max_retries: 2,
766                ..PrunedDeterminizeOptions::default()
767            },
768        );
769        assert!(impossible.is_err());
770
771        // With a workable cap the first attempt succeeds and nothing narrows.
772        let fine = determinize_lattice_pruned(&lattice, &PrunedDeterminizeOptions::default())
773            .expect("a lattice");
774        assert_eq!(fine.narrowed, 0);
775        assert_eq!(fine.beam, PrunedDeterminizeOptions::default().beam);
776    }
777
778    // Narrowing keeps the best path, which makes it a safe answer to running
779    // away rather than a wrong one.
780    #[test]
781    fn narrowing_never_loses_the_best_path() {
782        let symbols = 3;
783        let mut rng = Rng(0x00BB_CCDD_EEFF_0011);
784        let mut compared = 0;
785
786        for round in 0..100 {
787            let graph = random_graph(&mut rng, symbols);
788            let frames = 1 + rng.below(4);
789            let scores: Vec<f32> = (0..frames * symbols).map(|_| rng.cost()).collect();
790            let Some(lattice) = decode(&graph, &scores, frames, symbols) else {
791                continue;
792            };
793
794            let whole = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default())
795                .expect("a determinization");
796            let best = best_per_sequence(&whole)
797                .into_values()
798                .map(|weight| weight.weight().total())
799                .fold(f32::INFINITY, f32::min);
800
801            for beam in [8.0f32, 2.0, 0.5] {
802                let narrowed = determinize_lattice_pruned(
803                    &lattice,
804                    &PrunedDeterminizeOptions {
805                        beam,
806                        ..PrunedDeterminizeOptions::default()
807                    },
808                )
809                .expect("a lattice");
810                let after = best_per_sequence(&narrowed.lattice)
811                    .into_values()
812                    .map(|weight| weight.weight().total())
813                    .fold(f32::INFINITY, f32::min);
814                assert!(
815                    (after - best).abs() < 1e-3,
816                    "round {round} at beam {beam}: best became {after}, was {best}"
817                );
818            }
819            compared += 1;
820        }
821
822        assert!(compared > 50, "only {compared} rounds produced a lattice");
823    }
824
825    #[test]
826    fn a_cap_on_the_states_is_reported_rather_than_truncating() {
827        let symbols = 3;
828        let mut rng = Rng(0x0123_4567_89AB_CDEF);
829        let graph = random_graph(&mut rng, symbols);
830        let frames = 4;
831        let scores: Vec<f32> = (0..frames * symbols).map(|_| rng.cost()).collect();
832        let Some(lattice) = decode(&graph, &scores, frames, symbols) else {
833            return;
834        };
835
836        let err = determinize_lattice(
837            &lattice,
838            &DeterminizeLatticeOptions {
839                max_states: Some(1),
840                ..DeterminizeLatticeOptions::default()
841            },
842        );
843        assert!(err.is_err(), "a cap of one state should not be reachable");
844    }
845}