khmer_tokenizer_core/
hmm.rs

1//! A lightweight Hidden Markov Model over Khmer Character Clusters, used only
2//! as a fallback for clusters no dictionary strategy could match at all (see
3//! [`KhmerTokenizer::with_hmm`](crate::KhmerTokenizer::with_hmm)). States are
4//! the classic **BMES** tag set (Begin/Middle/End/Single), decoded with
5//! Viterbi — the same jieba-style design as `Strategy::UnigramDp`'s DAG, but
6//! for clusters that have zero dictionary matches to score in the first
7//! place. See `docs/ROADMAP.md` Phase 4.
8
9use std::collections::HashMap;
10
11const NUM_STATES: usize = 4;
12const BEGIN: usize = 0;
13const MIDDLE: usize = 1;
14const END: usize = 2;
15const SINGLE: usize = 3;
16
17/// A trained BMES Hidden Markov Model for segmenting clusters the dictionary
18/// has no match for at all.
19///
20/// Ships with no built-in parameters: like
21/// [`with_frequencies`](crate::KhmerTokenizer::with_frequencies), training
22/// needs a segmented corpus, and no bundleable, commercially-clean one has
23/// been found (see `docs/ROADMAP.md` Phase 4). Build one with
24/// [`HmmModel::from_counts`] from BMES tag counts gathered over a corpus
25/// you're licensed to use, then attach it with
26/// [`KhmerTokenizer::with_hmm`](crate::KhmerTokenizer::with_hmm).
27#[derive(Clone)]
28pub struct HmmModel {
29    start: [f64; NUM_STATES],
30    trans: [[f64; NUM_STATES]; NUM_STATES],
31    emit: HashMap<String, [f64; NUM_STATES]>,
32}
33
34/// Add-one (Laplace) smoothed log-probabilities from raw counts.
35fn smoothed_log_probs(counts: &[u64; NUM_STATES]) -> [f64; NUM_STATES] {
36    let total: f64 = counts.iter().map(|&c| c as f64 + 1.0).sum();
37    let mut out = [0.0; NUM_STATES];
38    for (i, &c) in counts.iter().enumerate() {
39        out[i] = ((c as f64 + 1.0) / total).ln();
40    }
41    out
42}
43
44impl HmmModel {
45    /// Build a model from raw BMES tag counts: `start_counts[state]` is how
46    /// often a Khmer run began tagged `state`; `trans_counts[i][j]` is how
47    /// often a cluster tagged `i` was immediately followed by one tagged
48    /// `j`; `emit_counts[cluster][state]` is how often `cluster` was tagged
49    /// `state`. All three are add-one smoothed, so an unseen transition or
50    /// start state is merely unlikely, never impossible.
51    pub fn from_counts(
52        start_counts: [u64; NUM_STATES],
53        trans_counts: [[u64; NUM_STATES]; NUM_STATES],
54        emit_counts: HashMap<String, [u64; NUM_STATES]>,
55    ) -> Self {
56        let start = smoothed_log_probs(&start_counts);
57        let mut trans = [[0.0; NUM_STATES]; NUM_STATES];
58        for (i, row) in trans_counts.iter().enumerate() {
59            trans[i] = smoothed_log_probs(row);
60        }
61        let emit = emit_counts
62            .into_iter()
63            .map(|(cluster, counts)| (cluster, smoothed_log_probs(&counts)))
64            .collect();
65        Self { start, trans, emit }
66    }
67
68    /// Log-emission-probability of `cluster` under `state`. A cluster never
69    /// seen during training gets a uniform, uninformative floor, so decoding
70    /// falls back on transition structure alone rather than treating it as
71    /// impossible.
72    fn emit_log_prob(&self, cluster: &str, state: usize) -> f64 {
73        match self.emit.get(cluster) {
74            Some(probs) => probs[state],
75            None => (1.0 / NUM_STATES as f64).ln(),
76        }
77    }
78
79    /// Viterbi-decode the most likely BMES tag sequence for `clusters`.
80    /// `clusters` must be non-empty.
81    // Each state index `s`/`ps` indexes several parallel arrays at once
82    // (`score`, `back`, `self.trans`, plus a call into `emit_log_prob`) —
83    // an `.iter().enumerate()` rewrite wouldn't cover all of them and would
84    // read worse than the plain DP-style index loop.
85    #[allow(clippy::needless_range_loop)]
86    fn viterbi_tags(&self, clusters: &[String]) -> Vec<usize> {
87        let n = clusters.len();
88        let mut score = vec![[f64::NEG_INFINITY; NUM_STATES]; n];
89        let mut back = vec![[0usize; NUM_STATES]; n];
90
91        for s in 0..NUM_STATES {
92            score[0][s] = self.start[s] + self.emit_log_prob(&clusters[0], s);
93        }
94        for t in 1..n {
95            for s in 0..NUM_STATES {
96                let mut best_score = f64::NEG_INFINITY;
97                let mut best_prev = 0;
98                for ps in 0..NUM_STATES {
99                    let candidate = score[t - 1][ps] + self.trans[ps][s];
100                    if candidate > best_score {
101                        best_score = candidate;
102                        best_prev = ps;
103                    }
104                }
105                back[t][s] = best_prev;
106                score[t][s] = best_score + self.emit_log_prob(&clusters[t], s);
107            }
108        }
109
110        let mut best_final = 0;
111        for s in 1..NUM_STATES {
112            if score[n - 1][s] > score[n - 1][best_final] {
113                best_final = s;
114            }
115        }
116
117        let mut tags = vec![0usize; n];
118        tags[n - 1] = best_final;
119        for t in (1..n).rev() {
120            tags[t - 1] = back[t][tags[t]];
121        }
122        tags
123    }
124
125    /// Segment a run of clusters that a dictionary strategy matched nothing
126    /// in at all, placing boundaries from the Viterbi-decoded BMES tags.
127    pub(crate) fn segment_oov(&self, clusters: &[String]) -> Vec<Vec<String>> {
128        if clusters.is_empty() {
129            return Vec::new();
130        }
131        let tags = self.viterbi_tags(clusters);
132
133        let mut tokens = Vec::new();
134        let mut current: Vec<String> = Vec::new();
135        for (cluster, &tag) in clusters.iter().zip(&tags) {
136            match tag {
137                BEGIN => {
138                    if !current.is_empty() {
139                        tokens.push(std::mem::take(&mut current));
140                    }
141                    current.push(cluster.clone());
142                }
143                MIDDLE => current.push(cluster.clone()),
144                END => {
145                    current.push(cluster.clone());
146                    tokens.push(std::mem::take(&mut current));
147                }
148                SINGLE => {
149                    if !current.is_empty() {
150                        tokens.push(std::mem::take(&mut current));
151                    }
152                    tokens.push(vec![cluster.clone()]);
153                }
154                _ => unreachable!("tags are always in 0..NUM_STATES"),
155            }
156        }
157        if !current.is_empty() {
158            tokens.push(current);
159        }
160        tokens
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167
168    /// Build a model where "a" strongly emits Begin, "b" strongly emits End,
169    /// "c" strongly emits Single, Begin strongly transitions to End, and End
170    /// strongly transitions to Single — decisively decoding ["a","b","c"]
171    /// into the two words ["a","b"] and ["c"].
172    fn two_word_model() -> HmmModel {
173        let start = [50, 0, 0, 0]; // sequences overwhelmingly start Begin
174        let mut trans = [[0u64; NUM_STATES]; NUM_STATES];
175        trans[BEGIN][END] = 50;
176        trans[END][SINGLE] = 50;
177
178        let mut emit = HashMap::new();
179        emit.insert("a".to_string(), [50, 0, 0, 0]);
180        emit.insert("b".to_string(), [0, 0, 50, 0]);
181        emit.insert("c".to_string(), [0, 0, 0, 50]);
182
183        HmmModel::from_counts(start, trans, emit)
184    }
185
186    #[test]
187    fn decodes_a_two_word_run_via_bmes_tags() {
188        let model = two_word_model();
189        let clusters = vec!["a".to_string(), "b".to_string(), "c".to_string()];
190        assert_eq!(
191            model.segment_oov(&clusters),
192            vec![
193                vec!["a".to_string(), "b".to_string()],
194                vec!["c".to_string()]
195            ]
196        );
197    }
198
199    #[test]
200    fn single_cluster_run_is_its_own_token() {
201        let model = two_word_model();
202        let clusters = vec!["c".to_string()];
203        assert_eq!(model.segment_oov(&clusters), vec![vec!["c".to_string()]]);
204    }
205
206    #[test]
207    fn empty_run_yields_no_tokens() {
208        let model = two_word_model();
209        assert!(model.segment_oov(&[]).is_empty());
210    }
211
212    #[test]
213    fn unseen_cluster_falls_back_to_uniform_emission_without_panicking() {
214        let model = two_word_model();
215        let clusters = vec!["z".to_string()];
216        // No assertion on the exact tag — just confirm decoding an unseen
217        // cluster degrades gracefully instead of crashing.
218        assert_eq!(model.segment_oov(&clusters).concat(), vec!["z".to_string()]);
219    }
220}
khmer_tokenizer_core/hmm.rs

khmer_tokenizer_core/
hmm.rs
