sparse_ngrams/table.rs
1//! Bigram priority model.
2//!
3//! Assigns a frequency-based priority to each byte pair, used by the sparse n-gram
4//! extraction algorithm to decide where n-gram boundaries fall.
5//!
6//! Priorities used to be a full 256×256 `u16` table baked from a `bigrams.bin` frequency
7//! ranking (~64kB in memory). They are now reconstructed from a compact *factored* model
8//! (~8.5kB) trained offline against a bigram frequency ranking. The ascii bigram `(a, b)` is
9//! scored as
10//! `BIGRAM_H[a] + BIGRAM_H[b] + (code << BIGRAM_CODE_SHIFT) + 1`, where [`BIGRAM_H`] is a single
11//! shared per-byte weight and `code` is a 4-bit (`0..=15`) per-bigram correction. Bigrams absent
12//! from the training data carry `code == 0` and are not special-cased: they fall back to the bare
13//! factored score, i.e. the model extrapolates a priority for them. The byte index `idx` is folded
14//! into the low 16 bits so every bigram gets a *unique* priority, while a higher score still means
15//! a more frequent bigram (~1.4% inversions vs. the reference ranking).
16
17/// A casefolded indexable byte uses 7 bits for ascii characters; any non-ascii (unicode)
18/// character is expected to have its high bit set. Only ascii bigrams are ever present, so
19/// non-ascii characters always resolve to priority `0`. The model therefore only covers the 128
20/// ascii values per character.
21const BIGRAM_ALPHABET: usize = 128;
22
23/// The per-bigram 4-bit correction code is scaled by `1 << BIGRAM_CODE_SHIFT` and added to the
24/// shared per-byte weights. A plain shift replaces what used to be a lookup into a learned
25/// 16-entry offset table, at a negligible accuracy cost.
26const BIGRAM_CODE_SHIFT: u32 = 10;
27
28/// 4-bit correction code per ascii bigram, packed two codes per byte (the even index in the low
29/// nibble). Scaled by `1 << BIGRAM_CODE_SHIFT` and added to the shared per-byte weights.
30static BIGRAM_CODE: &[u8; BIGRAM_ALPHABET * BIGRAM_ALPHABET / 2] =
31 include_bytes!("bigram_code.bin");
32
33/// Shared per-byte weight. `BIGRAM_H[b]` contributes to the priority of every bigram containing
34/// byte `b`; `7272` is the filler weight for bytes absent from the training data.
35static BIGRAM_H: [u16; BIGRAM_ALPHABET] = [
36 1712, 811, 1084, 886, 596, 91, 132, 450, 724, 8461, 16403, 286, 1348, 6151, 140, 343, 333, 64,
37 162, 45, 103, 178, 27, 5, 35, 0, 161, 2323, 70, 125, 44, 101, 13403, 8259, 11824, 8316, 8468,
38 8305, 8173, 10620, 10671, 9834, 9082, 8674, 9982, 10753, 11148, 10895, 10788, 10958, 10844,
39 10474, 10331, 10220, 10066, 9935, 10004, 9859, 10243, 9293, 9469, 9631, 9908, 8272, 8073, 7272,
40 7272, 7272, 7272, 7272, 7272, 7272, 7272, 7272, 7272, 7272, 7272, 7272, 7272, 7272, 7272, 7272,
41 7272, 7272, 7272, 7272, 7272, 7272, 7272, 7272, 7272, 9286, 8956, 8593, 6514, 9968, 8679,
42 11923, 11057, 11570, 11787, 12294, 11116, 10949, 10849, 11401, 9455, 10261, 11463, 11238,
43 11595, 11281, 11422, 9364, 11668, 12007, 11945, 10678, 10380, 10412, 10197, 10534, 9280, 8890,
44 7898, 8767, 6372, 1475,
45];
46
47/// Reconstructs the priority of the ascii bigram `(a, b)`; see [`bigram_priority`]. This rolling
48/// variant avoids re-loading `BIGRAM_H[a]`: consecutive bigrams overlap by one byte, so the caller
49/// passes the `h_b` returned for the previous position as `h_a` and gets `BIGRAM_H[b]` back for the
50/// next one. The `H` value is only used for ascii bytes; for a non-ascii byte the bigram is absent,
51/// so the masked lookup (`b & 0x7f`) merely returns a value that the next step discards.
52#[inline]
53pub(crate) fn bigram_priority_rolling(a: u8, b: u8, h_a: u32) -> (u32, u32) {
54 let h_b = BIGRAM_H[(b & (BIGRAM_ALPHABET as u8 - 1)) as usize] as u32;
55 if (a | b) >= BIGRAM_ALPHABET as u8 {
56 return (0, h_b);
57 }
58 let idx = a as usize * BIGRAM_ALPHABET + b as usize;
59 let code = (BIGRAM_CODE[idx >> 1] >> ((idx & 1) * 4)) & 0xF;
60 // The 4-bit `code` is scaled by `1 << BIGRAM_CODE_SHIFT` (a plain shift in place of an offset
61 // table) and added to the shared per-byte weights. Absent bigrams carry `code == 0`, so they
62 // fall back to the bare factored score `h_a + h_b + 1`. A higher score still means a more
63 // frequent bigram; `base` fits in 16 bits, so the unique per-bigram `idx` in the low 16 bits
64 // keeps every priority unique.
65 let base = h_a + h_b + ((code as u32) << BIGRAM_CODE_SHIFT) + 1;
66 ((base << 16) | idx as u32, h_b)
67}
68
69/// The `BIGRAM_H` weight of a single byte, used to seed [`bigram_priority_rolling`].
70#[inline]
71pub(crate) fn bigram_h(byte: u8) -> u32 {
72 BIGRAM_H[(byte & (BIGRAM_ALPHABET as u8 - 1)) as usize] as u32
73}
74
75/// Reconstructs the frequency-ranking priority of the ascii bigram `(a, b)`. Absent or non-ascii
76/// bigrams resolve to `0`; present bigrams get a strictly positive, unique priority where a higher
77/// value means a more frequent bigram. This priority is used to split strings into smaller
78/// n-grams.
79pub fn bigram_priority(a: u8, b: u8) -> u32 {
80 bigram_priority_rolling(a, b, bigram_h(a)).0
81}
82
83#[cfg(test)]
84mod tests {
85 use super::*;
86
87 #[test]
88 fn non_ascii_is_zero() {
89 assert_eq!(bigram_priority(0x80, b'a'), 0);
90 assert_eq!(bigram_priority(b'a', 0x80), 0);
91 assert_eq!(bigram_priority(0xff, 0xff), 0);
92 }
93
94 #[test]
95 fn ascii_bigrams_are_positive_and_unique() {
96 // Distinct ascii bigrams get distinct, strictly positive priorities (the `idx` in the low
97 // 16 bits guarantees uniqueness).
98 assert!(bigram_priority(b'a', b'b') > 0);
99 assert_ne!(bigram_priority(b'a', b'b'), bigram_priority(b'b', b'a'));
100 assert_ne!(bigram_priority(b'a', b'b'), bigram_priority(b'a', b'c'));
101 }
102
103 #[test]
104 fn rolling_matches_direct() {
105 // The rolling variant must reproduce the standalone `bigram_priority` for every ascii pair,
106 // and hand back `BIGRAM_H[b]` for the next step.
107 for a in 0u8..128 {
108 for b in 0u8..128 {
109 let (p, h_b) = bigram_priority_rolling(a, b, bigram_h(a));
110 assert_eq!(p, bigram_priority(a, b), "mismatch at ({a}, {b})");
111 assert_eq!(h_b, bigram_h(b), "h_b mismatch at ({a}, {b})");
112 }
113 }
114 }
115}