Skip to main content

sicada_decode/
ctc.rs

1//! The decoding graph for a CTC model.
2//!
3//! A CTC model emits one symbol per frame from an alphabet that includes a
4//! *blank*, and the transcript is recovered by collapsing runs of the same
5//! symbol and then deleting the blanks, so `_ a a _ a b` reads as `a a b`.
6//! Note where the blank matters: it is what separates the two `a`s. Without it
7//! the run would collapse to one.
8//!
9//! As an FST that rule is a graph of `V + 1` states, one per symbol the model
10//! could have emitted last, plus the state meaning "the last thing was a
11//! blank". [`ctc_topo`] builds it. Composing it with the acoustic matrix and
12//! taking the best path is CTC decoding; composing the two of them with a
13//! lexicon and a language model instead is the rest of a recogniser, and this
14//! is the piece it starts from.
15//!
16//! This is k2's `k2.ctc_topo(max_token, modified=False)`. The *modified*
17//! topology, which lets a frame be skipped, is not here.
18
19use sicada::arc::{Arc, ArcLabel, ArcStateId};
20use sicada::error::OpenFstError;
21use sicada::fst::{Fst, MutableFst};
22use sicada::fsts::vector_fst::VectorFst;
23use sicada::properties::K_FST_PROPERTIES;
24use sicada::weight::Weight;
25
26/// Builds the CTC topology for a model with `num_symbols` columns, blank first.
27///
28/// `label_offset` says where column 0 sits, and has to be the one the
29/// [`DenseFst`](crate::dense::DenseFst) was built with. It is 1 by default,
30/// because label 0 is epsilon to every FST algorithm and a blank is not one.
31///
32/// Input labels are the model's columns, offset. Output labels are the same
33/// symbols, so subtracting `label_offset` from an answer's labels gives columns
34/// back; blank and repeats emit nothing.
35///
36/// # Errors
37///
38/// A model with no symbols, or an alphabet that does not fit the arc's label
39/// type.
40pub fn ctc_topo<A: Arc>(num_symbols: usize, label_offset: i64) -> Result<VectorFst<A>, OpenFstError>
41where
42    A::Weight: Weight,
43{
44    if num_symbols < 2 {
45        return Err(OpenFstError::InvalidOperation(format!(
46            "ctc_topo: {num_symbols} symbols is not enough for a blank and something else"
47        )));
48    }
49    let label_of = |column: usize| -> Result<A::Label, OpenFstError> {
50        A::Label::from_i64(label_offset + column as i64).ok_or_else(|| {
51            OpenFstError::InvalidOperation(format!(
52                "ctc_topo: label {} does not fit the arc's label type",
53                label_offset + column as i64
54            ))
55        })
56    };
57    if label_offset < 1 {
58        return Err(OpenFstError::InvalidOperation(
59            "ctc_topo: column 0 would be epsilon, which consumes no frame".into(),
60        ));
61    }
62
63    // State 0 means "the last frame was a blank, or there has been none"; state
64    // `t` means "the last frame emitted symbol `t`". That is all the history
65    // the collapsing rule needs.
66    let mut fst: VectorFst<A> = VectorFst::new();
67    fst.reserve_states(num_symbols);
68    for _ in 0..num_symbols {
69        fst.add_state();
70    }
71    fst.set_start(A::StateId::from_usize(0));
72
73    let blank = label_of(0)?;
74    for last in 0..num_symbols {
75        let from = A::StateId::from_usize(last);
76        // Every state is final: the audio may end wherever it ends.
77        fst.set_final(from, A::Weight::one());
78
79        // A blank says nothing and resets what may repeat.
80        fst.add_arc(
81            from,
82            A::new(
83                blank,
84                A::Label::epsilon(),
85                A::Weight::one(),
86                A::StateId::from_usize(0),
87            ),
88        );
89
90        for symbol in 1..num_symbols {
91            let label = label_of(symbol)?;
92            // Emitting the same symbol again with no blank between is the same
93            // symbol held longer, so it says nothing more.
94            let says = if symbol == last {
95                A::Label::epsilon()
96            } else {
97                label
98            };
99            fst.add_arc(
100                from,
101                A::new(
102                    label,
103                    says,
104                    A::Weight::one(),
105                    A::StateId::from_usize(symbol),
106                ),
107            );
108        }
109    }
110
111    fst.properties(K_FST_PROPERTIES, true);
112    Ok(fst)
113}
114
115/// The CTC collapsing rule applied to a sequence of columns: runs of the same
116/// symbol become one, then the blanks go.
117///
118/// Column 0 is the blank. [`ctc_topo`] encodes this same rule; here it is
119/// written out directly, which is useful for checking an answer and for a caller
120/// who has a greedy argmax rather than a lattice.
121pub fn collapse(columns: &[usize]) -> Vec<usize> {
122    let mut out: Vec<usize> = Vec::with_capacity(columns.len());
123    let mut previous = usize::MAX;
124    for &column in columns {
125        if column != previous && column != 0 {
126            out.push(column);
127        }
128        previous = column;
129    }
130    out
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136    use sicada::arc::StdArc;
137    use sicada::fst::ExpandedFst;
138    use sicada::fsts::vector_fst::StdVectorFst;
139    use sicada::properties::{K_I_DETERMINISTIC, K_NO_I_EPSILONS};
140
141    use crate::compact::{DeterminizeLatticeOptions, determinize_lattice};
142    use crate::dense::DenseFst;
143    use crate::frontier::DecodeOptions;
144    use crate::lattice::{LatticeDecodeOptions, lattice_decode};
145    use crate::nbest::n_best;
146    use crate::viterbi::viterbi_decode;
147
148    const SYMBOLS: usize = 4;
149
150    fn topo() -> StdVectorFst {
151        ctc_topo(SYMBOLS, 1).expect("a topology")
152    }
153
154    fn certain(columns: &[usize]) -> Vec<f32> {
155        let mut scores = vec![10.0; columns.len() * SYMBOLS];
156        for (frame, &column) in columns.iter().enumerate() {
157            scores[frame * SYMBOLS + column] = 0.0;
158        }
159        scores
160    }
161
162    fn columns_of(labels: &[i32]) -> Vec<usize> {
163        labels.iter().map(|label| (label - 1) as usize).collect()
164    }
165
166    #[test]
167    fn it_is_deterministic_on_the_frames_it_reads() {
168        let fst = topo();
169        assert_eq!(fst.num_states(), SYMBOLS);
170        let props = fst.properties(K_I_DETERMINISTIC | K_NO_I_EPSILONS, true);
171        assert_ne!(props & K_I_DETERMINISTIC, 0, "two arcs read the same frame");
172        assert_ne!(props & K_NO_I_EPSILONS, 0, "an arc reads no frame");
173        for state in fst.states() {
174            assert_eq!(fst.num_arcs(state), SYMBOLS, "one arc per column");
175        }
176    }
177
178    #[test]
179    fn it_collapses_exactly_as_the_rule_says() {
180        let graph = topo();
181        for length in 1..=5usize {
182            let mut columns = vec![0usize; length];
183            loop {
184                let scores = certain(&columns);
185                let dense = DenseFst::<StdArc>::new(&scores, length, SYMBOLS).unwrap();
186                let decoded = viterbi_decode(&graph, &dense, &DecodeOptions::exhaustive())
187                    .unwrap()
188                    .expect("a path");
189                assert_eq!(
190                    columns_of(&decoded.labels),
191                    collapse(&columns),
192                    "for the alignment {columns:?}"
193                );
194
195                // Odometer over every alignment of this length.
196                let mut place = 0;
197                loop {
198                    if place == length {
199                        break;
200                    }
201                    columns[place] += 1;
202                    if columns[place] < SYMBOLS {
203                        break;
204                    }
205                    columns[place] = 0;
206                    place += 1;
207                }
208                if place == length {
209                    break;
210                }
211            }
212        }
213    }
214
215    #[test]
216    fn a_blank_is_what_lets_a_symbol_repeat() {
217        let graph = topo();
218
219        let held = certain(&[1, 1, 1]);
220        let dense = DenseFst::<StdArc>::new(&held, 3, SYMBOLS).unwrap();
221        let decoded = viterbi_decode(&graph, &dense, &DecodeOptions::exhaustive())
222            .unwrap()
223            .unwrap();
224        assert_eq!(columns_of(&decoded.labels), vec![1], "one long symbol");
225
226        let separated = certain(&[1, 0, 1]);
227        let dense = DenseFst::<StdArc>::new(&separated, 3, SYMBOLS).unwrap();
228        let decoded = viterbi_decode(&graph, &dense, &DecodeOptions::exhaustive())
229            .unwrap()
230            .unwrap();
231        assert_eq!(columns_of(&decoded.labels), vec![1, 1], "two of them");
232    }
233
234    #[test]
235    fn the_whole_pipeline_agrees_with_the_rule() {
236        let graph = topo();
237        // Frames 0 and 2 are sure of symbol 1. Frame 1 is torn between holding
238        // it, giving one long symbol, and a blank, which would make them two.
239        // It costs 1.0 more to blank, so holding wins and "1 1" is the
240        // runner-up.
241        let scores = [
242            9.0, 0.0, 9.0, 9.0, //
243            1.0, 0.0, 9.0, 9.0, //
244            9.0, 0.0, 9.0, 9.0,
245        ];
246        let dense = DenseFst::<StdArc>::new(&scores, 3, SYMBOLS).unwrap();
247        let lattice = lattice_decode(&graph, &dense, &LatticeDecodeOptions::exhaustive())
248            .unwrap()
249            .expect("a lattice");
250        let compact = determinize_lattice(&lattice, &DeterminizeLatticeOptions::default()).unwrap();
251        let best = n_best(&compact, 2).unwrap();
252        assert_eq!(best.len(), 2);
253
254        assert_eq!(columns_of(&best[0].words), vec![1], "one long symbol");
255        assert_eq!(columns_of(&best[1].words), vec![1, 1], "two of them");
256        assert!((best[0].cost() - 0.0).abs() < 1e-6, "{}", best[0].cost());
257        assert!((best[1].cost() - 1.0).abs() < 1e-6, "{}", best[1].cost());
258
259        // And each answer carries the frames that produced it.
260        assert_eq!(best[0].alignment().len(), 3, "one label per frame");
261        assert_eq!(columns_of(best[0].alignment()), vec![1, 1, 1]);
262        assert_eq!(columns_of(best[1].alignment()), vec![1, 0, 1]);
263    }
264
265    #[test]
266    fn an_alphabet_with_nothing_in_it_is_refused() {
267        assert!(ctc_topo::<StdArc>(1, 1).is_err());
268        assert!(ctc_topo::<StdArc>(4, 0).is_err(), "column 0 on epsilon");
269    }
270
271    #[test]
272    fn collapsing_is_runs_first_then_blanks() {
273        assert_eq!(collapse(&[0, 1, 1, 0, 1, 2]), vec![1, 1, 2]);
274        assert_eq!(collapse(&[]), Vec::<usize>::new());
275        assert_eq!(collapse(&[0, 0, 0]), Vec::<usize>::new());
276        assert_eq!(collapse(&[2, 2, 2]), vec![2]);
277    }
278}