src/api.rs

//! Public library API for downstream Rust projects.

use std::io::{Cursor, Read};
use std::path::Path;

use crate::{EmotionClassifier, LinearMoodModel, ModelIoError, MoodAnalysis, MoodScore};

const DEFAULT_CHECKPOINT: &[u8] = include_bytes!("../models/goemotions-linear.dadoes");

/// Default score threshold used to report active moods.
pub const DEFAULT_ACTIVE_MOOD_THRESHOLD: f32 = 0.35;

/// Returns the embedded default DADOES checkpoint bytes.
#[must_use]
pub fn default_checkpoint_bytes() -> &'static [u8] {
    DEFAULT_CHECKPOINT
}

/// A ready-to-use DADOES mood classifier for other Rust projects.
#[derive(Debug)]
pub struct DadoesClassifier {
    model: LinearMoodModel,
    active_mood_threshold: f32,
}

impl DadoesClassifier {
    /// Loads the embedded default checkpoint.
    pub fn from_default_model() -> Result<Self, ModelIoError> {
        let mut reader = Cursor::new(DEFAULT_CHECKPOINT);
        Self::from_reader(&mut reader)
    }

    /// Loads a checkpoint from any byte reader.
    pub fn from_reader<R: Read>(reader: &mut R) -> Result<Self, ModelIoError> {
        let model = LinearMoodModel::load_from_reader(reader)?;
        Ok(Self::from_model(model))
    }

    /// Loads a checkpoint from a filesystem path.
    pub fn from_path(path: impl AsRef<Path>) -> Result<Self, ModelIoError> {
        let mut file = std::fs::File::open(path)?;
        Self::from_reader(&mut file)
    }

    /// Wraps an already-loaded linear model.
    #[must_use]
    pub fn from_model(model: LinearMoodModel) -> Self {
        Self {
            model,
            active_mood_threshold: DEFAULT_ACTIVE_MOOD_THRESHOLD,
        }
    }

    /// Returns the threshold used by [`Self::active_moods`].
    #[must_use]
    pub fn active_mood_threshold(&self) -> f32 {
        self.active_mood_threshold
    }

    /// Returns moods whose score is at least the classifier threshold.
    pub fn active_moods<'a>(
        &self,
        analysis: &'a MoodAnalysis,
    ) -> impl Iterator<Item = MoodScore> + 'a {
        let threshold = self.active_mood_threshold;
        analysis.active_moods(threshold)
    }

    /// Returns the underlying linear model.
    #[must_use]
    pub fn model(&self) -> &LinearMoodModel {
        &self.model
    }
}

impl EmotionClassifier for DadoesClassifier {
    fn classify(&self, text: &str) -> MoodAnalysis {
        self.model.classify(text)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn default_model_loads_and_scores_report() {
        let classifier = DadoesClassifier::from_default_model();
        assert!(classifier.is_ok());

        if let Ok(classifier) = classifier {
            let analysis = classifier.classify("I felt frustrated after another failed repair.");
            assert!(analysis.primary_mood().is_some());
        }
    }
}