bayesian/

lib.rs

use std::collections::{HashMap, HashSet};
use std::hash::Hash;
use std::io::{Read, Write};

use flate2::{Compression, read::GzDecoder, write::GzEncoder};
use serde::{Deserialize, Serialize};

const DEFAULT_PROB: f64 = 1e-11;

// ---------------------------------------------------------------------------
// Private data structures
// ---------------------------------------------------------------------------

#[derive(Serialize, Deserialize)]
struct ClassData {
    freqs: HashMap<String, f64>,
    total: usize,
    tf_acc: HashMap<String, TfAccumulator>,
}

impl ClassData {
    fn new() -> Self {
        Self {
            freqs: HashMap::new(),
            total: 0,
            tf_acc: HashMap::new(),
        }
    }

    fn word_prob(&self, word: &str) -> f64 {
        let vocab = self.freqs.len();
        if self.total == 0 || vocab == 0 {
            return DEFAULT_PROB;
        }
        let count = self.freqs.get(word).copied().unwrap_or(0.0);
        (count + 1.0) / (self.total as f64 + vocab as f64)
    }
}

#[derive(Serialize, Deserialize)]
struct TfAccumulator {
    count: usize,
    sum_ln1p_tf: f64,
}

impl Default for TfAccumulator {
    fn default() -> Self {
        Self {
            count: 0,
            sum_ln1p_tf: 0.0,
        }
    }
}

#[derive(Serialize, Deserialize)]
struct TfIdfClassData {
    weights: HashMap<String, f64>,
    total: f64,
}

impl TfIdfClassData {
    fn word_prob(&self, word: &str) -> f64 {
        let vocab = self.weights.len();
        if self.total == 0.0 || vocab == 0 {
            return DEFAULT_PROB;
        }
        let weight = self.weights.get(word).copied().unwrap_or(0.0);
        (weight + 1.0) / (self.total + vocab as f64)
    }
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// A Naive Bayes classifier, optionally upgraded with TF-IDF weights.
///
/// Plain Naive Bayes is always available via [`classify`], [`log_scores`], and
/// [`prob_scores`]. After calling [`build_tfidf`] the TF-IDF variants become
/// available. You may call [`build_tfidf`] again at any time after learning
/// more documents — it recomputes from scratch over all accumulated data.
#[derive(Serialize, Deserialize)]
#[serde(bound(
    serialize = "C: Serialize",
    deserialize = "C: Deserialize<'de> + Eq + Hash"
))]
pub struct Classifier<C> {
    classes: Vec<C>,
    data: HashMap<C, ClassData>,
    tfidf: Option<HashMap<C, TfIdfClassData>>,
    learned: usize,
}

impl<C: Eq + Hash + Clone> Classifier<C> {
    /// Creates a new classifier.
    ///
    /// Panics if fewer than two classes are given, or if they are not unique.
    pub fn new(classes: Vec<C>) -> Self {
        assert!(classes.len() >= 2, "provide at least two classes");

        let mut seen = HashSet::new();
        for c in &classes {
            assert!(seen.insert(c), "class labels must be unique");
        }

        let data = classes
            .iter()
            .map(|c| (c.clone(), ClassData::new()))
            .collect();

        Self {
            classes,
            data,
            tfidf: None,
            learned: 0,
        }
    }

    /// Trains the classifier on a document (a slice of words) for the given class.
    ///
    /// Accumulates both raw word counts (used by plain Naive Bayes) and
    /// compact TF accumulators (used by [`build_tfidf`]). The two sets of data
    /// are kept separate, so calling [`build_tfidf`] never destroys the raw
    /// counts and plain classification always remains available.
    pub fn learn<S: AsRef<str>>(&mut self, document: &[S], class: &C) {
        let entry = self.data.get_mut(class).expect("unknown class");

        // Plain NB: raw word counts.
        for word in document {
            *entry.freqs.entry(word.as_ref().to_owned()).or_default() += 1.0;
            entry.total += 1;
        }

        // TF-IDF: one TF sample per unique word in this document, but stored
        // compactly as accumulators instead of vectors.
        let doc_len = document.len() as f64;
        if doc_len > 0.0 {
            let mut counts: HashMap<&str, usize> = HashMap::new();
            for word in document {
                *counts.entry(word.as_ref()).or_default() += 1;
            }
            for (word, count) in counts {
                let tf = count as f64 / doc_len;
                let ln1p = tf.ln_1p(); // ln(1 + tf)
                let acc = entry.tf_acc.entry(word.to_owned()).or_default();
                acc.count += 1;
                acc.sum_ln1p_tf += ln1p;
            }
        }

        self.learned += 1;
    }

    /// Computes TF-IDF weights from all documents learned so far and stores
    /// them internally, enabling the `_tfidf` family of classification methods.
    ///
    /// Safe to call multiple times — each call recomputes from scratch over the
    /// full accumulated training set, so you can learn more documents and call
    /// this again to refresh the weights without losing anything.
    pub fn build_tfidf(&mut self) {
        let total_docs = self.learned as f64;

        // Global document frequency: how many docs across all classes contain
        // each word. Each accumulator's `count` represents the number of docs
        // in that class that contain the word.
        let mut doc_freq: HashMap<&str, usize> = HashMap::new();
        for class_data in self.data.values() {
            for (word, acc) in &class_data.tf_acc {
                *doc_freq.entry(word.as_str()).or_default() += acc.count;
            }
        }

        let mut tfidf: HashMap<C, TfIdfClassData> = HashMap::new();
        for (class, class_data) in &self.data {
            let mut weights: HashMap<String, f64> = HashMap::new();
            let mut total = 0.0_f64;
            for (word, acc) in &class_data.tf_acc {
                let df = doc_freq[word.as_str()] as f64;
                let idf = (1.0 + total_docs / df).ln();
                // Weight for this class is idf * Σ ln(1 + tf) over docs in class
                let weight: f64 = acc.sum_ln1p_tf * idf;
                weights.insert(word.clone(), weight);
                total += weight;
            }
            tfidf.insert(class.clone(), TfIdfClassData { weights, total });
        }

        self.tfidf = Some(tfidf);
    }

    /// Returns `true` if [`build_tfidf`] has been called and TF-IDF weights
    /// are ready for classification.
    pub fn has_tfidf(&self) -> bool {
        self.tfidf.is_some()
    }

    /// Returns the total number of documents the classifier has been trained on.
    pub fn learned(&self) -> usize {
        self.learned
    }

    /// Returns the ordered slice of class labels used to build this classifier.
    pub fn classes(&self) -> &[C] {
        &self.classes
    }

    // --- Shared private helpers ---

    fn priors(&self) -> Vec<f64> {
        let n = self.classes.len() as f64;
        let total: f64 = self.data.values().map(|d| d.total as f64).sum();
        self.classes
            .iter()
            .map(|c| (self.data[c].total as f64 + 1.0) / (total + n))
            .collect()
    }

    fn find_max(scores: &[f64]) -> usize {
        scores
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.total_cmp(b))
            .map(|(i, _)| i)
            .expect("classifier has no classes")
    }

    fn log_to_probs(log_scores: &[f64]) -> Vec<f64> {
        // Log-sum-exp trick: shift by max before exponentiating to prevent
        // overflow/underflow for any document length.
        let max = log_scores.iter().copied().fold(f64::NEG_INFINITY, f64::max);
        let exps: Vec<f64> = log_scores.iter().map(|&s| (s - max).exp()).collect();
        let sum: f64 = exps.iter().sum();
        exps.iter().map(|&e| e / sum).collect()
    }

    /// Computes a log-score for each class given a document and a per-class
    /// word probability function. Shared by both NB and TF-IDF scoring paths.
    fn score_document<S, F>(&self, document: &[S], word_prob: F) -> Vec<f64>
    where
        S: AsRef<str>,
        F: Fn(&C, &str) -> f64,
    {
        self.priors()
            .into_iter()
            .zip(&self.classes)
            .map(|(prior, class)| {
                document.iter().fold(prior.ln(), |score, word| {
                    score + word_prob(class, word.as_ref()).ln()
                })
            })
            .collect()
    }

    // --- Plain Naive Bayes ---

    /// Returns the log-likelihood score for each class using raw word counts.
    /// Index `i` corresponds to `classes()[i]`.
    pub fn log_scores<S: AsRef<str>>(&self, document: &[S]) -> Vec<f64> {
        self.score_document(document, |class, word| self.data[class].word_prob(word))
    }

    /// Returns normalised probability scores for each class using raw word
    /// counts (sum to 1.0). Index `i` corresponds to `classes()[i]`.
    pub fn prob_scores<S: AsRef<str>>(&self, document: &[S]) -> Vec<f64> {
        Self::log_to_probs(&self.log_scores(document))
    }

    /// Returns the most likely class for the given document using plain Naive Bayes.
    pub fn classify<S: AsRef<str>>(&self, document: &[S]) -> &C {
        &self.classes[Self::find_max(&self.log_scores(document))]
    }

    // --- TF-IDF ---

    /// Returns the log-likelihood score for each class using TF-IDF weights.
    /// Index `i` corresponds to `classes()[i]`.
    ///
    /// Panics if [`build_tfidf`] has not been called yet.
    pub fn log_scores_tfidf<S: AsRef<str>>(&self, document: &[S]) -> Vec<f64> {
        let tfidf = self
            .tfidf
            .as_ref()
            .expect("call build_tfidf() before using TF-IDF classification");
        self.score_document(document, |class, word| tfidf[class].word_prob(word))
    }

    /// Returns normalised probability scores for each class using TF-IDF
    /// weights (sum to 1.0). Index `i` corresponds to `classes()[i]`.
    ///
    /// Panics if [`build_tfidf`] has not been called yet.
    pub fn prob_scores_tfidf<S: AsRef<str>>(&self, document: &[S]) -> Vec<f64> {
        Self::log_to_probs(&self.log_scores_tfidf(document))
    }

    /// Returns the most likely class for the given document using TF-IDF weights.
    ///
    /// Panics if [`build_tfidf`] has not been called yet.
    pub fn classify_tfidf<S: AsRef<str>>(&self, document: &[S]) -> &C {
        &self.classes[Self::find_max(&self.log_scores_tfidf(document))]
    }

    // --- Serialization ---

    /// Serializes the classifier (including any built TF-IDF weights) to a
    /// compressed binary blob (bincode + gzip).
    ///
    /// The result can be stored to disk, sent over the wire, etc. Pass it back
    /// to [`Classifier::from_data`] to reconstruct an identical classifier.
    pub fn serialize(&self) -> std::io::Result<Vec<u8>>
    where
        C: Serialize,
    {
        let bytes = bincode::serialize(self)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
        encoder.write_all(&bytes)?;
        encoder.finish()
    }

    /// Reconstructs a [`Classifier`] from bytes previously produced by
    /// [`Classifier::serialize`].
    pub fn from_data(data: impl AsRef<[u8]>) -> std::io::Result<Self>
    where
        C: for<'de> Deserialize<'de>,
    {
        let mut decoder = GzDecoder::new(data.as_ref());
        let mut bytes = Vec::new();
        decoder.read_to_end(&mut bytes)?;
        bincode::deserialize(&bytes)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
    }

    /// Writes the serialized classifier to a file at the given path.
    ///
    /// Creates the file if it does not exist, truncates it if it does.
    pub fn serialize_to_file(&self, path: impl AsRef<std::path::Path>) -> std::io::Result<()>
    where
        C: Serialize,
    {
        std::fs::write(path, self.serialize()?)
    }

    /// Reconstructs a [`Classifier`] from a file previously written by
    /// [`Classifier::serialize_to_file`].
    pub fn from_file(path: impl AsRef<std::path::Path>) -> std::io::Result<Self>
    where
        C: for<'de> Deserialize<'de>,
    {
        Self::from_data(std::fs::read(path)?)
    }
}

#[cfg(test)]
mod tests;