1use 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
56pub type CompactLattice<A> = VectorFst<CompactLatticeArc<A>>;
58
59#[derive(Debug, Clone, Copy, PartialEq)]
61pub struct DeterminizeLatticeOptions {
62 pub delta: f32,
64 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#[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 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
117pub 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 arc.olabel(),
159 arc.olabel(),
160 CompactLatticeWeight::new(*arc.weight(), alignment),
161 arc.nextstate(),
162 ),
163 );
164 }
165 }
166 compact
167}
168
169pub 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 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#[derive(Debug, Clone, Copy, PartialEq)]
204pub struct PrunedDeterminizeOptions {
205 pub beam: f32,
208 pub beam_ratio: f32,
211 pub max_retries: usize,
213 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#[derive(Debug, Clone)]
230pub struct PrunedLattice<L: ArcLabel, S: ArcStateId> {
231 pub lattice: VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>>,
233 pub beam: f32,
238 pub narrowed: usize,
240}
241
242pub 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 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 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 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 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 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 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 #[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 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 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 #[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 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 #[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 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 #[test]
641 fn the_alignment_travels_with_the_word() {
642 let mut graph: StdVectorFst = VectorFst::new();
644 graph.add_state();
645 graph.set_start(0);
646 graph.set_final(0, TropicalWeight::one());
647 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 let scores = [
655 0.0, 9.0, 9.0, 9.0, 0.0, 9.0, 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 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 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 #[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 #[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 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 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 #[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}