use sicada::algorithms::connect::connect;
use sicada::algorithms::determinize::{CommonDivisor, determinize_fsa};
use sicada::algorithms::prune::{PruneOptions, prune as prune_fst};
use sicada::algorithms::rmepsilon::rm_epsilon;
use sicada::arc::{Arc, ArcLabel, ArcStateId, ArcTpl};
use sicada::error::OpenFstError;
use sicada::fst::{ExpandedFst, Fst, MutableFst};
use sicada::fsts::vector_fst::VectorFst;
use sicada::weight::Weight;
use crate::compact_lattice_weight::{Alignment, CompactLatticeArc, CompactLatticeWeight};
use crate::lattice_weight::LatticeWeight;
pub type CompactLattice<A> = VectorFst<CompactLatticeArc<A>>;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct DeterminizeLatticeOptions {
pub delta: f32,
pub max_states: Option<usize>,
}
impl Default for DeterminizeLatticeOptions {
fn default() -> Self {
Self {
delta: 1.0 / 32.0,
max_states: Some(1 << 20),
}
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct CompactLatticeCommonDivisor;
impl<L: ArcLabel> CommonDivisor<CompactLatticeWeight<L>> for CompactLatticeCommonDivisor {
fn divisor(
&self,
w1: &CompactLatticeWeight<L>,
w2: &CompactLatticeWeight<L>,
) -> CompactLatticeWeight<L> {
let zero = CompactLatticeWeight::zero();
match (w1 == &zero, w2 == &zero) {
(true, true) => return zero,
(true, false) => return w2.clone(),
(false, true) => return w1.clone(),
(false, false) => {}
}
let shared = w1
.alignment()
.iter()
.zip(w2.alignment())
.take_while(|(a, b)| a == b)
.map(|(a, _)| *a)
.collect();
CompactLatticeWeight::new(w1.weight().plus(w2.weight()), shared)
}
}
pub fn to_compact<L, S>(
lattice: &VectorFst<ArcTpl<LatticeWeight, L, S>>,
) -> VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>>
where
L: ArcLabel,
S: ArcStateId,
{
let mut compact: VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>> = VectorFst::new();
compact.reserve_states(lattice.num_states());
for _ in 0..lattice.num_states() {
compact.add_state();
}
if let Some(start) = lattice.start() {
compact.set_start(start);
}
compact.set_input_symbols(lattice.output_symbols());
compact.set_output_symbols(lattice.output_symbols());
for state in lattice.states() {
let final_weight = lattice.final_weight(state);
if final_weight.is_member() && final_weight != LatticeWeight::zero() {
compact.set_final(state, CompactLatticeWeight::from_weight(final_weight));
}
for arc in lattice.arcs(state) {
let mut alignment = Alignment::new();
if arc.ilabel() != L::epsilon() {
alignment.push(arc.ilabel());
}
compact.add_arc(
state,
ArcTpl::new(
arc.olabel(),
arc.olabel(),
CompactLatticeWeight::new(*arc.weight(), alignment),
arc.nextstate(),
),
);
}
}
compact
}
pub fn determinize_lattice<L, S>(
lattice: &VectorFst<ArcTpl<LatticeWeight, L, S>>,
opts: &DeterminizeLatticeOptions,
) -> Result<VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>>, OpenFstError>
where
L: ArcLabel,
S: ArcStateId,
{
let mut compact = to_compact(lattice);
rm_epsilon(&mut compact, true)?;
let mut determinized: VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>> = VectorFst::new();
determinize_fsa(
&compact,
&mut determinized,
&CompactLatticeCommonDivisor,
opts.delta,
opts.max_states,
)?;
connect(&mut determinized);
Ok(determinized)
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct PrunedDeterminizeOptions {
pub beam: f32,
pub beam_ratio: f32,
pub max_retries: usize,
pub determinize: DeterminizeLatticeOptions,
}
impl Default for PrunedDeterminizeOptions {
fn default() -> Self {
Self {
beam: 8.0,
beam_ratio: 0.5,
max_retries: 6,
determinize: DeterminizeLatticeOptions::default(),
}
}
}
#[derive(Debug, Clone)]
pub struct PrunedLattice<L: ArcLabel, S: ArcStateId> {
pub lattice: VectorFst<ArcTpl<CompactLatticeWeight<L>, L, S>>,
pub beam: f32,
pub narrowed: usize,
}
pub fn determinize_lattice_pruned<L, S>(
lattice: &VectorFst<ArcTpl<LatticeWeight, L, S>>,
opts: &PrunedDeterminizeOptions,
) -> Result<PrunedLattice<L, S>, OpenFstError>
where
L: ArcLabel,
S: ArcStateId,
{
let mut beam = opts.beam;
let mut last: Option<OpenFstError> = None;
for narrowed in 0..=opts.max_retries {
let mut narrowed_lattice = lattice.clone();
if beam.is_finite() {
prune_fst(
&mut narrowed_lattice,
&PruneOptions::threshold(LatticeWeight::new(beam, 0.0)),
)?;
}
if narrowed_lattice.start().is_none() {
return Err(OpenFstError::InvalidOperation(format!(
"determinize_lattice_pruned: a beam of {beam} left no path at all"
)));
}
match determinize_lattice(&narrowed_lattice, &opts.determinize) {
Ok(lattice) => {
return Ok(PrunedLattice {
lattice,
beam,
narrowed,
});
}
Err(error) => {
last = Some(error);
beam *= opts.beam_ratio;
}
}
}
Err(last.unwrap_or_else(|| {
OpenFstError::InvalidOperation("determinize_lattice_pruned: no attempts were made".into())
}))
}
#[cfg(test)]
mod tests {
use super::*;
use rustc_hash::FxHashMap;
use sicada::arc::StdArc;
use sicada::fsts::vector_fst::StdVectorFst;
use sicada::properties::{K_ACYCLIC, K_FST_PROPERTIES, K_I_DETERMINISTIC};
use sicada::weights::float_weight::TropicalWeight;
use crate::dense::DenseFst;
use crate::lattice::{Lattice, LatticeDecodeOptions, lattice_decode};
struct Rng(u64);
impl Rng {
fn next(&mut self) -> u64 {
self.0 ^= self.0 << 13;
self.0 ^= self.0 >> 7;
self.0 ^= self.0 << 17;
self.0
}
fn below(&mut self, n: usize) -> usize {
(self.next() % n as u64) as usize
}
fn cost(&mut self) -> f32 {
self.below(256) as f32 / 16.0
}
}
fn word_sequences<W, F>(fst: &F) -> FxHashMap<Vec<i32>, f32>
where
W: Weight,
F: Fst<ArcTpl<W, i32, i32>>,
W: TotalCost,
{
let mut found: FxHashMap<Vec<i32>, f32> = FxHashMap::default();
let Some(start) = fst.start() else {
return found;
};
let mut stack = vec![(start, Vec::<i32>::new(), 0.0f32)];
while let Some((state, words, cost)) = stack.pop() {
let final_weight = fst.final_weight(state);
if final_weight.is_member() && final_weight != W::zero() {
let total = cost + final_weight.total_cost();
found
.entry(words.clone())
.and_modify(|best| *best = best.min(total))
.or_insert(total);
}
for arc in fst.arcs(state) {
let mut next = words.clone();
if arc.olabel() != 0 {
next.push(arc.olabel());
}
stack.push((arc.nextstate(), next, cost + arc.weight().total_cost()));
}
}
found
}
trait TotalCost {
fn total_cost(&self) -> f32;
}
impl TotalCost for LatticeWeight {
fn total_cost(&self) -> f32 {
self.total()
}
}
impl TotalCost for CompactLatticeWeight<i32> {
fn total_cost(&self) -> f32 {
self.weight().total()
}
}
fn random_graph(rng: &mut Rng, symbols: usize) -> StdVectorFst {
let states = 1 + rng.below(4);
let mut graph: StdVectorFst = VectorFst::new();
for _ in 0..states {
graph.add_state();
}
graph.set_start(0);
for from in 0..states as i32 {
for _ in 0..1 + rng.below(3) {
let ilabel = 1 + rng.below(symbols) as i32;
let olabel = if rng.below(2) == 0 {
0
} else {
10 * (1 + rng.below(2) as i32)
};
let to = rng.below(states) as i32;
graph.add_arc(
from,
StdArc::new(ilabel, olabel, TropicalWeight(rng.cost()), to),
);
}
if rng.below(2) == 0 {
graph.set_final(from, TropicalWeight(rng.cost()));
}
}
graph.properties(K_FST_PROPERTIES, true);
graph
}
fn decode(
graph: &StdVectorFst,
scores: &[f32],
frames: usize,
symbols: usize,
) -> Option<Lattice<StdArc>> {
let dense = DenseFst::<StdArc>::new(scores, frames, symbols).unwrap();
lattice_decode(graph, &dense, &LatticeDecodeOptions::exhaustive()).unwrap()
}
#[test]
fn it_keeps_every_word_sequence_at_the_same_cost() {
let symbols = 3;
let mut rng = Rng(0x00D1_5EA5_E1A1_2345);
let mut compared = 0;
for round in 0..150 {
let graph = random_graph(&mut rng, symbols);
let frames = 1 + rng.below(4);
let scores: Vec<f32> = (0..frames * symbols).map(|_| rng.cost()).collect();
let Some(lattice) = decode(&graph, &scores, frames, symbols) else {
continue;
};
let before = word_sequences(&lattice);
let compact = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default())
.expect("a determinization");
let after = word_sequences(&compact);
assert_eq!(
before.len(),
after.len(),
"round {round}: {} sequences became {}",
before.len(),
after.len()
);
for (words, cost) in &before {
let found = after
.get(words)
.unwrap_or_else(|| panic!("round {round}: {words:?} went missing"));
assert!(
(found - cost).abs() < 1e-3,
"round {round}: {words:?} cost {found}, was {cost}"
);
}
compared += 1;
}
assert!(compared > 80, "only {compared} rounds produced a lattice");
}
fn best_per_sequence<W, F>(fst: &F) -> FxHashMap<Vec<i32>, CompactLatticeWeight<i32>>
where
W: Weight + AsCompact,
F: Fst<ArcTpl<W, i32, i32>>,
{
let mut found: FxHashMap<Vec<i32>, CompactLatticeWeight<i32>> = FxHashMap::default();
let Some(start) = fst.start() else {
return found;
};
let one = CompactLatticeWeight::<i32>::one();
let mut stack = vec![(start, Vec::<i32>::new(), one)];
while let Some((state, words, weight)) = stack.pop() {
let final_weight = fst.final_weight(state);
if final_weight.is_member() && final_weight != W::zero() {
let whole = weight.times(&final_weight.as_compact());
found
.entry(words.clone())
.and_modify(|best| *best = best.plus(&whole))
.or_insert(whole);
}
for arc in fst.arcs(state) {
let mut next = words.clone();
if arc.olabel() != 0 {
next.push(arc.olabel());
}
stack.push((
arc.nextstate(),
next,
weight.times(&arc.weight().as_compact()),
));
}
}
found
}
trait AsCompact {
fn as_compact(&self) -> CompactLatticeWeight<i32>;
}
impl AsCompact for CompactLatticeWeight<i32> {
fn as_compact(&self) -> Self {
self.clone()
}
}
#[test]
fn it_keeps_the_best_alignment_for_each_word_sequence() {
let symbols = 3;
let mut rng = Rng(0x00AB_CDEF_0123_4567);
let mut compared = 0;
for round in 0..150 {
let graph = random_graph(&mut rng, symbols);
let frames = 1 + rng.below(4);
let scores: Vec<f32> = (0..frames * symbols).map(|_| rng.cost()).collect();
let Some(lattice) = decode(&graph, &scores, frames, symbols) else {
continue;
};
let expected = best_per_sequence(&to_compact(&lattice));
let compact = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default())
.expect("a determinization");
let found = best_per_sequence(&compact);
assert_eq!(expected.len(), found.len(), "round {round}");
for (words, want) in &expected {
let got = found
.get(words)
.unwrap_or_else(|| panic!("round {round}: {words:?} went missing"));
assert_eq!(
got.alignment(),
want.alignment(),
"round {round}: {words:?} kept the wrong alignment"
);
assert!(
(got.weight().total() - want.weight().total()).abs() < 1e-3,
"round {round}: {words:?} cost {got} vs {want}"
);
}
compared += 1;
}
assert!(compared > 80, "only {compared} rounds produced a lattice");
}
#[test]
fn each_word_sequence_is_a_single_path() {
let symbols = 3;
let mut rng = Rng(0x0FED_CBA9_8765_4321);
let mut collapsed = 0;
for round in 0..150 {
let graph = random_graph(&mut rng, symbols);
let frames = 1 + rng.below(4);
let scores: Vec<f32> = (0..frames * symbols).map(|_| rng.cost()).collect();
let Some(lattice) = decode(&graph, &scores, frames, symbols) else {
continue;
};
let compact = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default())
.expect("a determinization");
if compact.start().is_none() {
continue;
}
let props = compact.properties(K_I_DETERMINISTIC | K_ACYCLIC, true);
assert_ne!(
props & K_I_DETERMINISTIC,
0,
"round {round}: not deterministic, so a word sequence has two paths"
);
let paths_before = count_paths(&lattice);
let sequences = word_sequences(&compact).len();
let paths_after = count_paths(&compact);
assert_eq!(paths_after, sequences, "round {round}");
if paths_before > paths_after {
collapsed += 1;
}
}
assert!(collapsed > 40, "nothing collapsed in {collapsed} rounds");
}
fn count_paths<W, F>(fst: &F) -> usize
where
W: Weight,
F: Fst<ArcTpl<W, i32, i32>>,
{
let Some(start) = fst.start() else {
return 0;
};
let mut stack = vec![start];
let mut paths = 0;
while let Some(state) = stack.pop() {
let final_weight = fst.final_weight(state);
if final_weight.is_member() && final_weight != W::zero() {
paths += 1;
}
for arc in fst.arcs(state) {
stack.push(arc.nextstate());
}
}
paths
}
#[test]
fn the_alignment_travels_with_the_word() {
let mut graph: StdVectorFst = VectorFst::new();
graph.add_state();
graph.set_start(0);
graph.set_final(0, TropicalWeight::one());
graph.add_arc(0, StdArc::new(1, 10, TropicalWeight::one(), 0));
graph.add_arc(0, StdArc::new(2, 0, TropicalWeight::one(), 0));
graph.add_arc(0, StdArc::new(3, 0, TropicalWeight::one(), 0));
graph.properties(K_FST_PROPERTIES, true);
let scores = [
0.0, 9.0, 9.0, 9.0, 0.0, 9.0, 9.0, 9.0, 0.0,
];
let lattice = decode(&graph, &scores, 3, 3).expect("a lattice");
let compact = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default()).unwrap();
let start = compact.start().expect("a start");
let arcs: Vec<_> = compact.arcs(start).collect();
assert_eq!(arcs.len(), 1, "one word, one arc");
assert_eq!(arcs[0].olabel(), 10);
let mut best: VectorFst<ArcTpl<CompactLatticeWeight<i32>, i32, i32>> = VectorFst::new();
sicada::algorithms::shortest_path::shortest_path(
&compact,
&mut best,
&sicada::algorithms::shortest_path::ShortestPathOptions::default(),
)
.expect("a best path");
let (words, weight) =
sicada::string::string_fst_to_output_labels(&best).expect("a single path");
assert_eq!(words, vec![10], "one word was said");
assert_eq!(
weight.alignment(),
&[1, 2, 3],
"the three frames the word spanned"
);
assert!(weight.weight().total().abs() < 1e-6, "{weight}");
}
#[test]
fn it_writes_and_reads_back_as_an_fst() {
use sicada::fst::{FstReadOptions, FstWriteOptions};
use std::io::Write as _;
let scores = [0.0, 1.0, 2.0, 0.5, 0.25, 3.0];
let mut rng = Rng(0x0011_2233_4455_6677);
let graph = random_graph(&mut rng, 3);
let Some(lattice) = decode(&graph, &scores, 2, 3) else {
return;
};
let compact = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default()).unwrap();
if compact.start().is_none() {
return;
}
assert_eq!(
<ArcTpl<CompactLatticeWeight<i32>, i32, i32> as Arc>::type_name().as_str(),
"compactlattice44",
"the name the header records"
);
let mut bytes = Vec::new();
compact
.write(&mut bytes, &FstWriteOptions::default())
.expect("written");
let directory = tempfile::tempdir().expect("a directory");
let path = directory.path().join("lattice.fst");
std::fs::File::create(&path)
.unwrap()
.write_all(&bytes)
.unwrap();
let mut file = std::fs::File::open(&path).unwrap();
let read: VectorFst<ArcTpl<CompactLatticeWeight<i32>, i32, i32>> =
VectorFst::read(&mut file, &FstReadOptions::default()).expect("read back");
assert_eq!(read.num_states(), compact.num_states());
assert_eq!(read.start(), compact.start());
for state in compact.states() {
assert_eq!(read.final_weight(state), compact.final_weight(state));
assert_eq!(
read.arcs(state).collect::<Vec<_>>(),
compact.arcs(state).collect::<Vec<_>>(),
"state {state}"
);
}
}
#[test]
fn it_narrows_the_beam_rather_than_giving_up() {
let symbols = 3;
let mut rng = Rng(0x0777_8888_9999_AAAA);
let graph = random_graph(&mut rng, symbols);
let frames = 4;
let scores: Vec<f32> = (0..frames * symbols).map(|_| rng.cost()).collect();
let Some(lattice) = decode(&graph, &scores, frames, symbols) else {
return;
};
let impossible = determinize_lattice_pruned(
&lattice,
&PrunedDeterminizeOptions {
determinize: DeterminizeLatticeOptions {
max_states: Some(1),
..DeterminizeLatticeOptions::default()
},
max_retries: 2,
..PrunedDeterminizeOptions::default()
},
);
assert!(impossible.is_err());
let fine = determinize_lattice_pruned(&lattice, &PrunedDeterminizeOptions::default())
.expect("a lattice");
assert_eq!(fine.narrowed, 0);
assert_eq!(fine.beam, PrunedDeterminizeOptions::default().beam);
}
#[test]
fn narrowing_never_loses_the_best_path() {
let symbols = 3;
let mut rng = Rng(0x00BB_CCDD_EEFF_0011);
let mut compared = 0;
for round in 0..100 {
let graph = random_graph(&mut rng, symbols);
let frames = 1 + rng.below(4);
let scores: Vec<f32> = (0..frames * symbols).map(|_| rng.cost()).collect();
let Some(lattice) = decode(&graph, &scores, frames, symbols) else {
continue;
};
let whole = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default())
.expect("a determinization");
let best = best_per_sequence(&whole)
.into_values()
.map(|weight| weight.weight().total())
.fold(f32::INFINITY, f32::min);
for beam in [8.0f32, 2.0, 0.5] {
let narrowed = determinize_lattice_pruned(
&lattice,
&PrunedDeterminizeOptions {
beam,
..PrunedDeterminizeOptions::default()
},
)
.expect("a lattice");
let after = best_per_sequence(&narrowed.lattice)
.into_values()
.map(|weight| weight.weight().total())
.fold(f32::INFINITY, f32::min);
assert!(
(after - best).abs() < 1e-3,
"round {round} at beam {beam}: best became {after}, was {best}"
);
}
compared += 1;
}
assert!(compared > 50, "only {compared} rounds produced a lattice");
}
#[test]
fn a_cap_on_the_states_is_reported_rather_than_truncating() {
let symbols = 3;
let mut rng = Rng(0x0123_4567_89AB_CDEF);
let graph = random_graph(&mut rng, symbols);
let frames = 4;
let scores: Vec<f32> = (0..frames * symbols).map(|_| rng.cost()).collect();
let Some(lattice) = decode(&graph, &scores, frames, symbols) else {
return;
};
let err = determinize_lattice(
&lattice,
&DeterminizeLatticeOptions {
max_states: Some(1),
..DeterminizeLatticeOptions::default()
},
);
assert!(err.is_err(), "a cap of one state should not be reachable");
}
}