cufflink_model/
lib.rs

1//! # cufflink-model
2//!
3//! Run small trained text classifiers inside Cufflink services. Training
4//! happens offline (the Python toolkit in `ml/`); a service embeds the
5//! exported file and classifies with it, natively or in WASM, without
6//! calling out to anything.
7//!
8//! ```rust,ignore
9//! use cufflink_model::Model;
10//!
11//! static GENRES: &[u8] = include_bytes!("../models/genre.cflm");
12//!
13//! let model = Model::from_bytes(GENRES)?;
14//! let p = model.predict("The Dragon of Ember Hill");
15//! if p.confidence >= model.threshold().unwrap_or(0.9) {
16//!     // use p.label
17//! }
18//! ```
19//!
20//! A WASM guest is a fresh instance on every call, so [`Model::from_bytes`]
21//! reads the file in place: it checks the layout once and keeps borrowed
22//! slices, with no tables rebuilt. Every offset is bounds-checked at load,
23//! and nothing in this crate panics on a malformed file or unusual input.
24//!
25//! The file format is described in `docs/design-notes/models.md`.
26
27pub mod featurise;
28
29use std::cmp::Ordering;
30use std::collections::BTreeMap;
31use std::fmt;
32
33use serde_json::Value;
34
35/// The first four bytes of every model file.
36pub const MAGIC: &[u8; 4] = b"CFLM";
37/// The file layout this crate reads.
38pub const FORMAT_VERSION: u16 = 1;
39
40/// The kind of model a file holds.
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum Kind {
43    /// Character n-gram TF-IDF features scored by a linear layer and softmax.
44    NgramLinear,
45}
46
47impl Kind {
48    fn from_code(code: u16) -> Result<Self, ModelError> {
49        match code {
50            1 => Ok(Kind::NgramLinear),
51            // 2 = transformer and 3 = ensemble are reserved; see the design note.
52            other => Err(ModelError::UnsupportedKind(other)),
53        }
54    }
55}
56
57/// Why a model file couldn't be read.
58#[derive(Debug, Clone, PartialEq, Eq)]
59pub enum ModelError {
60    /// The file doesn't start with [`MAGIC`].
61    NotAModel,
62    /// The file uses a layout version this crate doesn't read.
63    UnsupportedVersion(u16),
64    /// The file holds a kind of model this crate doesn't run.
65    UnsupportedKind(u16),
66    /// The file is cut short or an offset points outside it.
67    Truncated(&'static str),
68    /// The file's contents contradict each other.
69    Invalid(String),
70}
71
72impl fmt::Display for ModelError {
73    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
74        match self {
75            ModelError::NotAModel => write!(f, "not a cufflink model file"),
76            ModelError::UnsupportedVersion(v) => write!(f, "unsupported model format version {v}"),
77            ModelError::UnsupportedKind(k) => write!(f, "unsupported model kind {k}"),
78            ModelError::Truncated(what) => write!(f, "model file truncated at {what}"),
79            ModelError::Invalid(why) => write!(f, "invalid model file: {why}"),
80        }
81    }
82}
83
84impl std::error::Error for ModelError {}
85
86/// A model's answer for one text.
87#[derive(Debug, Clone, PartialEq)]
88pub struct Prediction {
89    /// The most likely label.
90    pub label: String,
91    /// Its probability, between 0 and 1.
92    pub confidence: f64,
93    /// The most likely labels, best first, with their probabilities (up to five).
94    pub ranked: Vec<(String, f64)>,
95}
96
97/// A loaded model, borrowing the file it was read from.
98#[derive(Debug, Clone)]
99pub struct Model<'a> {
100    kind: Kind,
101    meta: Value,
102    labels: Vec<String>,
103    ngram_min: usize,
104    ngram_max: usize,
105    sublinear_tf: bool,
106    l2_norm: bool,
107    linear: Linear<'a>,
108}
109
110#[derive(Debug, Clone)]
111struct Linear<'a> {
112    n_features: usize,
113    n_classes: usize,
114    intercepts: &'a [u8],
115    idf: &'a [u8],
116    str_offsets: &'a [u8],
117    strings: &'a [u8],
118    w_offsets: &'a [u8],
119    entries: &'a [u8],
120}
121
122const ENTRY: usize = 6; // u16 class + f32 weight
123
124struct Reader<'a> {
125    bytes: &'a [u8],
126    pos: usize,
127}
128
129impl<'a> Reader<'a> {
130    fn take(&mut self, n: usize, what: &'static str) -> Result<&'a [u8], ModelError> {
131        let end = self.pos.checked_add(n).ok_or(ModelError::Truncated(what))?;
132        let slice = self
133            .bytes
134            .get(self.pos..end)
135            .ok_or(ModelError::Truncated(what))?;
136        self.pos = end;
137        Ok(slice)
138    }
139    fn u16(&mut self, what: &'static str) -> Result<u16, ModelError> {
140        let b = self.take(2, what)?;
141        Ok(u16::from_le_bytes([b[0], b[1]]))
142    }
143    fn u32(&mut self, what: &'static str) -> Result<u32, ModelError> {
144        let b = self.take(4, what)?;
145        Ok(u32::from_le_bytes([b[0], b[1], b[2], b[3]]))
146    }
147    fn array(
148        &mut self,
149        count: usize,
150        width: usize,
151        what: &'static str,
152    ) -> Result<&'a [u8], ModelError> {
153        let n = count
154            .checked_mul(width)
155            .ok_or(ModelError::Truncated(what))?;
156        self.take(n, what)
157    }
158}
159
160fn u32_at(slice: &[u8], i: usize) -> Option<u32> {
161    let b = slice.get(i * 4..i * 4 + 4)?;
162    Some(u32::from_le_bytes([b[0], b[1], b[2], b[3]]))
163}
164
165fn f32_at(slice: &[u8], i: usize) -> Option<f32> {
166    let b = slice.get(i * 4..i * 4 + 4)?;
167    Some(f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
168}
169
170/// Offsets must start at 0, never decrease, and end exactly at `len`.
171fn check_offsets(offsets: &[u8], count: usize, len: usize, what: &str) -> Result<(), ModelError> {
172    let mut prev = 0u32;
173    for i in 0..=count {
174        let o = u32_at(offsets, i)
175            .ok_or_else(|| ModelError::Invalid(format!("{what} offsets short")))?;
176        if (i == 0 && o != 0) || o < prev {
177            return Err(ModelError::Invalid(format!("{what} offsets out of order")));
178        }
179        prev = o;
180    }
181    if prev as usize != len {
182        return Err(ModelError::Invalid(format!(
183            "{what} offsets don't cover the data"
184        )));
185    }
186    Ok(())
187}
188
189impl<'a> Model<'a> {
190    /// Read a model file. Checks the whole layout; the returned model borrows `bytes`.
191    pub fn from_bytes(bytes: &'a [u8]) -> Result<Self, ModelError> {
192        let mut r = Reader { bytes, pos: 0 };
193        if r.take(4, "magic").map_err(|_| ModelError::NotAModel)? != MAGIC {
194            return Err(ModelError::NotAModel);
195        }
196        let version = r.u16("version")?;
197        if version != FORMAT_VERSION {
198            return Err(ModelError::UnsupportedVersion(version));
199        }
200        let kind = Kind::from_code(r.u16("kind")?)?;
201        let meta_len = r.u32("metadata length")? as usize;
202        let meta: Value = serde_json::from_slice(r.take(meta_len, "metadata")?)
203            .map_err(|e| ModelError::Invalid(format!("metadata: {e}")))?;
204
205        let labels: Vec<String> = meta["labels"]
206            .as_array()
207            .ok_or_else(|| ModelError::Invalid("metadata has no labels".into()))?
208            .iter()
209            .map(|l| l.as_str().map(String::from))
210            .collect::<Option<_>>()
211            .ok_or_else(|| ModelError::Invalid("labels must be strings".into()))?;
212        let feat = &meta["featuriser"];
213        if feat["kind"].as_str() != Some("char_wb") || feat["version"].as_u64() != Some(1) {
214            return Err(ModelError::Invalid(
215                "featuriser must be char_wb version 1".into(),
216            ));
217        }
218        let ngram_min = feat["ngram_min"].as_u64().unwrap_or(1) as usize;
219        let ngram_max = feat["ngram_max"].as_u64().unwrap_or(4) as usize;
220        if ngram_min == 0 || ngram_min > ngram_max || ngram_max > 16 {
221            return Err(ModelError::Invalid("bad n-gram range".into()));
222        }
223        let sublinear_tf = feat["sublinear_tf"].as_bool().unwrap_or(true);
224        let l2_norm = match feat["norm"].as_str() {
225            Some("l2") | None => true,
226            Some("none") => false,
227            Some(other) => return Err(ModelError::Invalid(format!("unknown norm {other}"))),
228        };
229
230        let n_features = r.u32("feature count")? as usize;
231        let n_classes = r.u32("class count")? as usize;
232        if n_classes != labels.len() || n_classes == 0 || n_classes > u16::MAX as usize {
233            return Err(ModelError::Invalid(
234                "class count doesn't match labels".into(),
235            ));
236        }
237        let intercepts = r.array(n_classes, 4, "intercepts")?;
238        let idf = r.array(n_features, 4, "idf")?;
239        let str_offsets = r.array(n_features + 1, 4, "feature offsets")?;
240        let str_len = u32_at(str_offsets, n_features)
241            .ok_or(ModelError::Truncated("feature offsets"))? as usize;
242        let strings = r.take(str_len, "feature strings")?;
243        let w_offsets = r.array(n_features + 1, 4, "weight offsets")?;
244        let n_entries =
245            u32_at(w_offsets, n_features).ok_or(ModelError::Truncated("weight offsets"))? as usize;
246        let entries = r.array(n_entries, ENTRY, "weights")?;
247        if r.pos != bytes.len() {
248            return Err(ModelError::Invalid("trailing bytes after the model".into()));
249        }
250
251        check_offsets(str_offsets, n_features, strings.len(), "feature")?;
252        check_offsets(w_offsets, n_features, n_entries, "weight")?;
253        if std::str::from_utf8(strings).is_err() {
254            return Err(ModelError::Invalid("feature strings aren't UTF-8".into()));
255        }
256        let linear = Linear {
257            n_features,
258            n_classes,
259            intercepts,
260            idf,
261            str_offsets,
262            strings,
263            w_offsets,
264            entries,
265        };
266        // Features must be strictly increasing by bytes for the binary search.
267        for i in 1..n_features {
268            if linear.feature(i - 1) >= linear.feature(i) {
269                return Err(ModelError::Invalid("features not sorted and unique".into()));
270            }
271        }
272        for e in 0..n_entries {
273            let class = u16::from_le_bytes([entries[e * ENTRY], entries[e * ENTRY + 1]]) as usize;
274            if class >= n_classes {
275                return Err(ModelError::Invalid(
276                    "weight for a class that doesn't exist".into(),
277                ));
278            }
279        }
280
281        Ok(Model {
282            kind,
283            meta,
284            labels,
285            ngram_min,
286            ngram_max,
287            sublinear_tf,
288            l2_norm,
289            linear,
290        })
291    }
292
293    /// The kind of model.
294    pub fn kind(&self) -> Kind {
295        self.kind
296    }
297
298    /// The labels the model chooses between, in output order.
299    pub fn labels(&self) -> &[String] {
300        &self.labels
301    }
302
303    /// Everything the trainer recorded about the model (name, version, threshold, ...).
304    pub fn meta(&self) -> &Value {
305        &self.meta
306    }
307
308    /// The model's own `name` and `version`, as `name@version`, for logs.
309    pub fn id(&self) -> String {
310        format!(
311            "{}@{}",
312            self.meta["name"].as_str().unwrap_or("model"),
313            self.meta["version"].as_str().unwrap_or("?")
314        )
315    }
316
317    /// The confidence the trainer suggests acting on, if it recorded one.
318    pub fn threshold(&self) -> Option<f64> {
319        self.meta["threshold"].as_f64()
320    }
321
322    /// The probability of every label for `text`, in [`Model::labels`] order.
323    pub fn probabilities(&self, text: &str) -> Vec<f64> {
324        let lin = &self.linear;
325        let mut counts: BTreeMap<usize, u32> = BTreeMap::new();
326        for gram in featurise::ngrams(text, self.ngram_min, self.ngram_max) {
327            if let Some(i) = lin.find(gram.as_bytes()) {
328                *counts.entry(i).or_insert(0) += 1;
329            }
330        }
331        let mut x: Vec<(usize, f64)> = counts
332            .into_iter()
333            .map(|(i, c)| {
334                let tf = if self.sublinear_tf {
335                    1.0 + (c as f64).ln()
336                } else {
337                    c as f64
338                };
339                (i, tf * f32_at(lin.idf, i).unwrap_or(0.0) as f64)
340            })
341            .collect();
342        if self.l2_norm {
343            let norm = x.iter().map(|(_, v)| v * v).sum::<f64>().sqrt();
344            if norm > 0.0 {
345                for (_, v) in &mut x {
346                    *v /= norm;
347                }
348            }
349        }
350        let mut scores: Vec<f64> = (0..lin.n_classes)
351            .map(|c| f32_at(lin.intercepts, c).unwrap_or(0.0) as f64)
352            .collect();
353        for (i, v) in x {
354            let (start, end) = (
355                u32_at(lin.w_offsets, i).unwrap_or(0) as usize,
356                u32_at(lin.w_offsets, i + 1).unwrap_or(0) as usize,
357            );
358            for e in start..end {
359                let Some(b) = lin.entries.get(e * ENTRY..e * ENTRY + ENTRY) else {
360                    break;
361                };
362                let class = u16::from_le_bytes([b[0], b[1]]) as usize;
363                let w = f32::from_le_bytes([b[2], b[3], b[4], b[5]]) as f64;
364                if let Some(s) = scores.get_mut(class) {
365                    *s += v * w;
366                }
367            }
368        }
369        softmax(&mut scores);
370        scores
371    }
372
373    /// The model's answer for `text`.
374    pub fn predict(&self, text: &str) -> Prediction {
375        let probs = self.probabilities(text);
376        let mut order: Vec<usize> = (0..probs.len()).collect();
377        order.sort_by(|&a, &b| {
378            probs[b]
379                .partial_cmp(&probs[a])
380                .unwrap_or(Ordering::Equal)
381                .then(a.cmp(&b))
382        });
383        let ranked: Vec<(String, f64)> = order
384            .iter()
385            .take(5)
386            .map(|&i| (self.labels[i].clone(), probs[i]))
387            .collect();
388        let (label, confidence) = ranked.first().cloned().unwrap_or_default();
389        Prediction {
390            label,
391            confidence,
392            ranked,
393        }
394    }
395}
396
397impl Linear<'_> {
398    fn feature(&self, i: usize) -> &[u8] {
399        let start = u32_at(self.str_offsets, i).unwrap_or(0) as usize;
400        let end = u32_at(self.str_offsets, i + 1).unwrap_or(0) as usize;
401        self.strings.get(start..end).unwrap_or(&[])
402    }
403
404    fn find(&self, gram: &[u8]) -> Option<usize> {
405        let (mut lo, mut hi) = (0usize, self.n_features);
406        while lo < hi {
407            let mid = lo + (hi - lo) / 2;
408            match self.feature(mid).cmp(gram) {
409                Ordering::Less => lo = mid + 1,
410                Ordering::Greater => hi = mid,
411                Ordering::Equal => return Some(mid),
412            }
413        }
414        None
415    }
416}
417
418fn softmax(scores: &mut [f64]) {
419    let max = scores.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
420    let mut sum = 0.0;
421    for s in scores.iter_mut() {
422        *s = (*s - max).exp();
423        sum += *s;
424    }
425    if sum > 0.0 {
426        for s in scores.iter_mut() {
427            *s /= sum;
428        }
429    }
430}
431
432#[cfg(test)]
433mod tests {
434    use super::*;
435
436    /// A hand-built two-label model: "x" votes for `b`, nothing else is known.
437    fn tiny() -> Vec<u8> {
438        let meta = serde_json::json!({
439            "name": "tiny", "version": "1", "labels": ["a", "b"], "threshold": 0.6,
440            "featuriser": {"kind": "char_wb", "version": 1, "ngram_min": 1, "ngram_max": 1,
441                           "sublinear_tf": true, "norm": "l2"}
442        })
443        .to_string();
444        let mut f = Vec::new();
445        f.extend_from_slice(MAGIC);
446        f.extend_from_slice(&1u16.to_le_bytes());
447        f.extend_from_slice(&1u16.to_le_bytes());
448        f.extend_from_slice(&(meta.len() as u32).to_le_bytes());
449        f.extend_from_slice(meta.as_bytes());
450        f.extend_from_slice(&1u32.to_le_bytes()); // features
451        f.extend_from_slice(&2u32.to_le_bytes()); // classes
452        for b in [0.5f32, 0.0] {
453            f.extend_from_slice(&b.to_le_bytes());
454        }
455        f.extend_from_slice(&1.0f32.to_le_bytes()); // idf of "x"
456        for o in [0u32, 1] {
457            f.extend_from_slice(&o.to_le_bytes());
458        }
459        f.extend_from_slice(b"x");
460        for o in [0u32, 1] {
461            f.extend_from_slice(&o.to_le_bytes());
462        }
463        f.extend_from_slice(&1u16.to_le_bytes());
464        f.extend_from_slice(&3.0f32.to_le_bytes());
465        f
466    }
467
468    #[test]
469    fn known_features_move_the_answer() {
470        let bytes = tiny();
471        let m = Model::from_bytes(&bytes).unwrap();
472        assert_eq!(m.id(), "tiny@1");
473        assert_eq!(m.threshold(), Some(0.6));
474        assert_eq!(m.predict("x").label, "b");
475        // Unknown text falls back to the intercepts, which favour "a".
476        let p = m.predict("zzz");
477        assert_eq!(p.label, "a");
478        let expected = 1.0 / (1.0 + (-0.5f64).exp());
479        assert!((p.confidence - expected).abs() < 1e-9);
480        assert_eq!(p.ranked.len(), 2);
481    }
482
483    #[test]
484    fn empty_text_is_answered_not_panicked() {
485        let bytes = tiny();
486        let m = Model::from_bytes(&bytes).unwrap();
487        assert_eq!(m.predict("").label, "a");
488        assert_eq!(m.predict("   ").label, "a");
489    }
490
491    #[test]
492    fn every_truncation_is_an_error_not_a_panic() {
493        let bytes = tiny();
494        for len in 0..bytes.len() {
495            assert!(
496                Model::from_bytes(&bytes[..len]).is_err(),
497                "accepted {len} bytes"
498            );
499        }
500        let mut longer = bytes.clone();
501        longer.push(0);
502        assert!(Model::from_bytes(&longer).is_err());
503    }
504
505    #[test]
506    fn corrupt_files_are_rejected() {
507        let bytes = tiny();
508        assert_eq!(
509            Model::from_bytes(b"nope").unwrap_err(),
510            ModelError::NotAModel
511        );
512
513        let mut v2 = bytes.clone();
514        v2[4] = 2;
515        assert_eq!(
516            Model::from_bytes(&v2).unwrap_err(),
517            ModelError::UnsupportedVersion(2)
518        );
519
520        let mut transformer = bytes.clone();
521        transformer[6] = 2;
522        assert_eq!(
523            Model::from_bytes(&transformer).unwrap_err(),
524            ModelError::UnsupportedKind(2)
525        );
526
527        let mut bad_class = bytes.clone();
528        let at = bad_class.len() - 6;
529        bad_class[at] = 9;
530        assert!(matches!(
531            Model::from_bytes(&bad_class),
532            Err(ModelError::Invalid(_))
533        ));
534    }
535}
cufflink_model/lib.rs

cufflink_model/
lib.rs
