Import upstream v0.16.22, stripped

Upstream commit: 474dd0229cb20cf513036619781ed97bd8073c3f
Enterprise-only files removed or emptied: 63
Enterprise-only snippets removed: 117 in 50 files
Dangling module declarations removed: 5
Cargo edits turning enterprise off: 14
Verification: clean
Enterprise feature gates left for rebuilt features: 19 in 18 files

Produced by tools/fork/strip.py. The full report is in docs/fork/strip-reports/ on main.

crates/nlp/src/classifier/adam.rs

@@ -0,0 +1,111 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use crate::classifier::{Optimizer, model::FhClassifier};
+
+pub struct Adam {
+    parameters: Vec<f32>,
+    bias: f32,
+    learning_rate: f32,
+    beta1: f32,
+    beta2: f32,
+    epsilon: f32,
+    t: f32,
+    m0: Vec<f32>,
+    v0: Vec<f32>,
+    m_bias: f32,
+    v_bias: f32,
+
+    // Step info
+    bias2_sqrt: f32,
+    alpha_t: f32,
+}
+
+impl Adam {
+    pub fn new(n_parameters: usize, learning_rate: f32) -> Self {
+        Adam {
+            parameters: vec![0.0; n_parameters],
+            learning_rate,
+            beta1: 0.9,
+            beta2: 0.999,
+            epsilon: 1e-8,
+            t: 0.0,
+            m0: vec![0.0; n_parameters],
+            v0: vec![0.0; n_parameters],
+            m_bias: 0.0,
+            v_bias: 0.0,
+            bias: 0.0,
+            bias2_sqrt: 0.0,
+            alpha_t: 0.0,
+        }
+    }
+
+    pub fn with_hyperparams(mut self, beta1: f32, beta2: f32, epsilon: f32) -> Self {
+        self.beta1 = beta1;
+        self.beta2 = beta2;
+        self.epsilon = epsilon;
+        self
+    }
+
+    pub fn with_initial_weights(self, value: f32) -> Self {
+        Adam {
+            parameters: vec![value; self.parameters.len()],
+            ..self
+        }
+    }
+}
+
+impl Optimizer for Adam {
+    #[inline(always)]
+    fn step(&mut self) {
+        self.t += 1.0;
+        let bias1 = 1.0 - self.beta1.powf(self.t);
+        self.bias2_sqrt = (1.0 - self.beta2.powf(self.t)).sqrt();
+        self.alpha_t = self.learning_rate / bias1;
+    }
+
+    #[inline(always)]
+    fn update_param(&mut self, i: usize, g: f32) {
+        self.m0[i] = self.beta1 * self.m0[i] + (1.0 - self.beta1) * g;
+        self.v0[i] = self.beta2 * self.v0[i] + (1.0 - self.beta2) * g * g;
+        self.parameters[i] -=
+            self.alpha_t * self.m0[i] / (self.v0[i].sqrt() / self.bias2_sqrt + self.epsilon);
+    }
+
+    #[inline(always)]
+    fn update_bias(&mut self, g: f32) {
+        self.m_bias = self.beta1 * self.m_bias + (1.0 - self.beta1) * g;
+        self.v_bias = self.beta2 * self.v_bias + (1.0 - self.beta2) * g * g;
+        self.bias -=
+            self.alpha_t * self.m_bias / (self.v_bias.sqrt() / self.bias2_sqrt + self.epsilon);
+    }
+
+    #[inline(always)]
+    fn get_param(&self, idx: usize) -> f32 {
+        self.parameters[idx]
+    }
+
+    #[inline(always)]
+    fn get_bias(&self) -> f32 {
+        self.bias
+    }
+
+    #[inline(always)]
+    fn get_param_mut(&mut self, idx: usize) -> &mut f32 {
+        &mut self.parameters[idx]
+    }
+
+    fn build_classifier(&self) -> FhClassifier {
+        FhClassifier {
+            parameters: self.parameters.clone(),
+            bias: self.bias,
+        }
+    }
+
+    fn num_parameters(&self) -> usize {
+        self.parameters.len()
+    }
+}

crates/nlp/src/classifier/feature.rs

@@ -0,0 +1,175 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use std::collections::HashMap;
+use xxhash_rust::xxh3::xxh3_64_with_seed;
+
+#[derive(Debug)]
+pub struct Sample<T> {
+    pub features: Vec<T>,
+    pub class: f32,
+}
+
+pub struct FhFeatureBuilder {
+    pub(super) weight_mask: u64,
+}
+
+#[derive(Debug)]
+pub struct FhFeature {
+    pub idx: usize,
+    pub weight: f32,
+}
+
+#[derive(Debug)]
+pub struct CcfhFeature {
+    pub idx_w1: usize,
+    pub idx_w2: usize,
+    pub idx_i: usize,
+    pub weight: f32,
+}
+
+pub struct CcfhFeatureBuilder {
+    pub(super) weight_mask: u64,
+    pub(super) indicator_mask: u64,
+}
+
+pub trait FeatureWeight {
+    fn idx(&self) -> usize;
+    fn weight(&self) -> f32;
+    fn weight_mut(&mut self) -> &mut f32;
+}
+
+pub trait UnprocessedFeature {
+    fn prefix(&self) -> u16;
+    fn value(&self) -> &[u8];
+}
+
+impl FeatureWeight for FhFeature {
+    fn weight(&self) -> f32 {
+        self.weight
+    }
+
+    fn weight_mut(&mut self) -> &mut f32 {
+        &mut self.weight
+    }
+
+    fn idx(&self) -> usize {
+        self.idx
+    }
+}
+
+impl FeatureWeight for CcfhFeature {
+    fn weight(&self) -> f32 {
+        self.weight
+    }
+
+    fn weight_mut(&mut self) -> &mut f32 {
+        &mut self.weight
+    }
+
+    fn idx(&self) -> usize {
+        self.idx_w1
+    }
+}
+
+impl FeatureBuilder for FhFeatureBuilder {
+    type Feature = FhFeature;
+
+    fn build_feature(&self, bytes: &[u8], weight: f32) -> FhFeature {
+        let hash1 = xxh3_64_with_seed(bytes, 0);
+        let sign = if hash1 & (1 << 63) == 0 { 1.0 } else { -1.0 };
+
+        FhFeature {
+            idx: (hash1 & self.weight_mask) as usize,
+            weight: sign * weight,
+        }
+    }
+}
+
+impl FeatureBuilder for CcfhFeatureBuilder {
+    type Feature = CcfhFeature;
+
+    fn build_feature(&self, bytes: &[u8], weight: f32) -> CcfhFeature {
+        let hash1 = xxh3_64_with_seed(bytes, 0);
+        let hash2 = xxh3_64_with_seed(bytes, 0x9E3779B97F4A7C15);
+        let hash3 = xxh3_64_with_seed(bytes, 0x517CC1B727220A95);
+        let sign = if hash3 & (1 << 63) == 0 { 1.0 } else { -1.0 };
+
+        CcfhFeature {
+            idx_w1: (hash1 & self.weight_mask) as usize,
+            idx_w2: (hash2 & self.weight_mask) as usize,
+            idx_i: (hash3 & self.indicator_mask) as usize,
+            weight: sign * weight,
+        }
+    }
+}
+
+pub trait FeatureBuilder {
+    // Feature type associated type
+    type Feature: FeatureWeight;
+
+    fn build_feature(&self, bytes: &[u8], weight: f32) -> Self::Feature;
+
+    fn scale<I: UnprocessedFeature>(&self, features: &mut HashMap<I, f32>) {
+        // Log frequency scaling
+        for x in features.values_mut() {
+            *x = x.ln_1p();
+        }
+    }
+
+    fn build<I: UnprocessedFeature>(
+        &self,
+        features_in: &HashMap<I, f32>,
+        account_id: Option<u32>,
+        l2_normalize: bool,
+    ) -> Vec<Self::Feature> {
+        let mut features_out = Vec::with_capacity(features_in.len());
+        let mut buf = Vec::with_capacity(2 + 4 + 63);
+        for (feature, count) in features_in {
+            buf.extend_from_slice(&feature.prefix().to_be_bytes());
+            buf.extend_from_slice(feature.value());
+            features_out.push(self.build_feature(&buf, *count));
+
+            if let Some(account_id) = account_id {
+                buf.extend_from_slice(&account_id.to_be_bytes());
+                features_out.push(self.build_feature(&buf, *count));
+            }
+
+            buf.clear();
+        }
+
+        // L2 normalization
+        if l2_normalize {
+            let sum_of_squares = features_out
+                .iter()
+                .map(|f| f.weight() as f64 * f.weight() as f64)
+                .sum::<f64>();
+            if sum_of_squares > 0.0 {
+                let norm = sum_of_squares.sqrt() as f32;
+                for feature in &mut features_out {
+                    *feature.weight_mut() /= norm;
+                }
+            }
+        }
+
+        features_out
+    }
+}
+
+impl<T> Sample<T> {
+    pub fn new(features: Vec<T>, class: bool) -> Self {
+        Self {
+            features,
+            class: if class { 1.0 } else { 0.0 },
+        }
+    }
+}
+
+impl<T> AsRef<Sample<T>> for Sample<T> {
+    fn as_ref(&self) -> &Sample<T> {
+        self
+    }
+}

crates/nlp/src/classifier/ftrl.rs

@@ -0,0 +1,132 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use crate::classifier::{Optimizer, model::FhClassifier};
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug)]
+pub struct Ftrl {
+    alpha: f64,
+    beta: f64,
+    l1_ratio: f64,
+    l2_ratio: f64,
+    zn: Vec<Zn>,
+    zn_bias: Zn,
+}
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Clone, Copy, Debug, Default)]
+pub struct Zn {
+    z: f32,
+    n: f64,
+}
+
+impl Ftrl {
+    pub fn new(n_features: usize) -> Self {
+        Ftrl {
+            alpha: 2.0,
+            beta: 1.0,
+            l1_ratio: 0.001,
+            l2_ratio: 0.0001,
+            zn: vec![Zn::default(); n_features],
+            zn_bias: Zn::default(),
+        }
+    }
+
+    pub fn with_hyperparams(mut self, alpha: f64, beta: f64, l1_ratio: f64, l2_ratio: f64) -> Self {
+        self.alpha = alpha;
+        self.beta = beta;
+        self.l1_ratio = l1_ratio;
+        self.l2_ratio = l2_ratio;
+        self
+    }
+
+    pub fn set_hyperparams(&mut self, alpha: f64, beta: f64, l1_ratio: f64, l2_ratio: f64) {
+        self.alpha = alpha;
+        self.beta = beta;
+        self.l1_ratio = l1_ratio;
+        self.l2_ratio = l2_ratio;
+    }
+
+    pub fn with_initial_weights(self, value: f32) -> Self {
+        Ftrl {
+            zn: vec![Zn { z: value, n: 0.0 }; self.zn.len()],
+            ..self
+        }
+    }
+}
+
+impl Optimizer for Ftrl {
+    #[inline(always)]
+    fn update_param(&mut self, idx: usize, grad: f32) {
+        let zn = &mut self.zn[idx];
+        let current_w = if zn.z.abs() as f64 <= self.l1_ratio {
+            0.0
+        } else {
+            -(zn.z - zn.z.signum() * self.l1_ratio as f32)
+                / (self.l2_ratio + (self.beta + zn.n.sqrt()) / self.alpha) as f32
+        };
+        let grad = grad as f64;
+        let grad_sq = grad * grad;
+        let sigma = ((zn.n + grad_sq).sqrt() - zn.n.sqrt()) / self.alpha;
+        zn.z += (grad - sigma * current_w as f64) as f32;
+        zn.n += grad_sq;
+    }
+
+    #[inline(always)]
+    fn update_bias(&mut self, grad: f32) {
+        let current_bias = -self.zn_bias.z
+            / ((self.zn_bias.n.sqrt() + self.beta) / self.alpha + self.l2_ratio) as f32;
+        let grad = grad as f64;
+        let grad_sq = grad * grad;
+        let sigma = ((self.zn_bias.n + grad_sq).sqrt() - self.zn_bias.n.sqrt()) / self.alpha;
+        self.zn_bias.z += (grad - sigma * current_bias as f64) as f32;
+        self.zn_bias.n += grad_sq;
+    }
+
+    #[inline(always)]
+    fn get_param(&self, idx: usize) -> f32 {
+        let zn = self.zn[idx];
+        if zn.z.abs() as f64 <= self.l1_ratio {
+            0.0
+        } else {
+            -(zn.z - zn.z.signum() * self.l1_ratio as f32)
+                / (self.l2_ratio + (self.beta + zn.n.sqrt()) / self.alpha) as f32
+        }
+    }
+
+    #[inline(always)]
+    fn get_bias(&self) -> f32 {
+        -self.zn_bias.z / ((self.zn_bias.n.sqrt() + self.beta) / self.alpha + self.l2_ratio) as f32
+    }
+
+    fn step(&mut self) {}
+
+    #[inline(always)]
+    fn get_param_mut(&mut self, idx: usize) -> &mut f32 {
+        &mut self.zn[idx].z
+    }
+
+    fn build_classifier(&self) -> FhClassifier {
+        FhClassifier {
+            parameters: self
+                .zn
+                .iter()
+                .map(|zn| {
+                    if zn.z.abs() as f64 <= self.l1_ratio {
+                        0.0
+                    } else {
+                        -(zn.z - zn.z.signum() * self.l1_ratio as f32)
+                            / (self.l2_ratio + (self.beta + zn.n.sqrt()) / self.alpha) as f32
+                    }
+                })
+                .collect(),
+            bias: self.get_bias(),
+        }
+    }
+
+    fn num_parameters(&self) -> usize {
+        self.zn.len()
+    }
+}

crates/nlp/src/classifier/mod.rs

@@ -0,0 +1,49 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use crate::classifier::model::FhClassifier;
+
+pub mod adam;
+pub mod feature;
+pub mod ftrl;
+pub mod model;
+pub mod reservoir;
+pub mod sgd;
+pub mod train;
+
+const MAX_DLOSS: f32 = 1e4;
+
+pub trait Optimizer {
+    fn step(&mut self);
+    fn update_param(&mut self, i: usize, g: f32);
+    fn update_bias(&mut self, g: f32);
+    fn get_param(&self, idx: usize) -> f32;
+    fn get_param_mut(&mut self, idx: usize) -> &mut f32;
+    fn get_bias(&self) -> f32;
+    fn build_classifier(&self) -> FhClassifier;
+    fn num_parameters(&self) -> usize;
+}
+
+#[inline(always)]
+fn sigmoid(z: f32) -> f32 {
+    let z = z.clamp(-35.0, 35.0);
+    if z >= 0.0 {
+        1.0 / (1.0 + (-z).exp())
+    } else {
+        let exp_z = z.exp();
+        exp_z / (1.0 + exp_z)
+    }
+}
+
+#[inline(always)]
+fn gradient(y: f32, p: f32) -> f32 {
+    if p > -16.0 {
+        let exp_tmp = (-p).exp();
+        ((1.0 - y) - y * exp_tmp) / (1.0 + exp_tmp)
+    } else {
+        p.exp() - y
+    }
+}

crates/nlp/src/classifier/model.rs

@@ -0,0 +1,109 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use crate::classifier::{
+    feature::{CcfhFeature, CcfhFeatureBuilder, FhFeature, FhFeatureBuilder},
+    sigmoid,
+};
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
+pub struct FhClassifier {
+    pub(crate) parameters: Vec<f32>,
+    pub(crate) bias: f32,
+}
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
+pub struct CcfhClassifier {
+    pub(crate) parameters: Vec<f32>,
+    pub(crate) indicators: Vec<f32>,
+    pub(crate) bias: f32,
+}
+
+impl FhClassifier {
+    pub fn predict_proba_sample(&self, features: &[FhFeature]) -> f32 {
+        let mut z: f32 = 0.0;
+
+        for f in features {
+            z += self.parameters[f.idx] * f.weight;
+        }
+
+        sigmoid(z + self.bias)
+    }
+
+    pub fn predict(&self, features: &[FhFeature]) -> f32 {
+        if self.predict_proba_sample(features) > 0.7 {
+            1.0
+        } else {
+            0.0
+        }
+    }
+
+    pub fn predict_batch<I>(&self, test: I) -> Vec<f32>
+    where
+        I: IntoIterator,
+        I::Item: AsRef<Vec<FhFeature>>,
+    {
+        test.into_iter()
+            .map(|features| self.predict(features.as_ref()))
+            .collect()
+    }
+
+    pub fn feature_builder(&self) -> FhFeatureBuilder {
+        FhFeatureBuilder {
+            weight_mask: (self.parameters.len() - 1) as u64,
+        }
+    }
+
+    pub fn parameters(&self) -> &[f32] {
+        &self.parameters
+    }
+
+    pub fn bias(&self) -> f32 {
+        self.bias
+    }
+}
+
+impl CcfhClassifier {
+    pub fn predict_proba_sample(&self, features: &[CcfhFeature]) -> f32 {
+        let mut z: f32 = 0.0;
+        for f in features {
+            let q = self.indicators[f.idx_i];
+            let v1 = self.parameters[f.idx_w1];
+            let v2 = self.parameters[f.idx_w2];
+            z += (q * v1 + (1.0 - q) * v2) * f.weight;
+        }
+        sigmoid(z + self.bias)
+    }
+
+    pub fn predict(&self, features: &[CcfhFeature]) -> f32 {
+        if self.predict_proba_sample(features) >= 0.5 {
+            1.0
+        } else {
+            0.0
+        }
+    }
+
+    pub fn predict_batch<I>(&self, test: I) -> Vec<f32>
+    where
+        I: IntoIterator,
+        I::Item: AsRef<Vec<CcfhFeature>>,
+    {
+        test.into_iter()
+            .map(|features| self.predict(features.as_ref()))
+            .collect()
+    }
+
+    pub fn feature_builder(&self) -> CcfhFeatureBuilder {
+        CcfhFeatureBuilder {
+            weight_mask: (self.parameters.len() - 1) as u64,
+            indicator_mask: (self.indicators.len() - 1) as u64,
+        }
+    }
+
+    pub fn is_active(&self) -> bool {
+        !self.parameters.is_empty()
+    }
+}

crates/nlp/src/classifier/reservoir.rs

@@ -0,0 +1,91 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use rand::{RngExt, seq::IndexedRandom};
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug)]
+pub struct SampleReservoir<T> {
+    pub spam: SampleReservoirClass<T>,
+    pub ham: SampleReservoirClass<T>,
+}
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug)]
+pub struct SampleReservoirClass<T> {
+    pub buffer: Vec<T>,
+    pub total_seen: u64,
+}
+
+impl<T: Clone + Eq> SampleReservoir<T> {
+    pub fn update_reservoir(&mut self, item: &T, is_spam: bool, capacity: usize) {
+        let class = if is_spam {
+            &mut self.spam
+        } else {
+            &mut self.ham
+        };
+
+        class.total_seen += 1;
+
+        if class.buffer.len() < capacity {
+            class.buffer.push(item.clone());
+        } else if let Some(buf) = class
+            .buffer
+            .get_mut(rand::rng().random_range(0..class.total_seen as usize))
+        {
+            *buf = item.clone();
+        }
+    }
+
+    pub fn update_counts(&mut self, is_spam: bool) {
+        let class = if is_spam {
+            &mut self.spam
+        } else {
+            &mut self.ham
+        };
+
+        class.total_seen += 1;
+    }
+
+    pub fn replay_samples(
+        &mut self,
+        count_needed: usize,
+        is_spam: bool,
+    ) -> impl Iterator<Item = &T> {
+        (if is_spam {
+            &mut self.spam
+        } else {
+            &mut self.ham
+        })
+        .buffer
+        .sample(&mut rand::rng(), count_needed)
+    }
+
+    pub fn remove_sample(&mut self, item: &T, is_spam: bool) {
+        let class = if is_spam {
+            &mut self.spam
+        } else {
+            &mut self.ham
+        };
+
+        if let Some(pos) = class.buffer.iter().position(|x| x == item) {
+            class.buffer.swap_remove(pos);
+        }
+    }
+}
+
+impl<T> Default for SampleReservoir<T> {
+    fn default() -> Self {
+        SampleReservoir {
+            spam: SampleReservoirClass {
+                buffer: Vec::new(),
+                total_seen: 0,
+            },
+            ham: SampleReservoirClass {
+                buffer: Vec::new(),
+                total_seen: 0,
+            },
+        }
+    }
+}

crates/nlp/src/classifier/sgd.rs

@@ -0,0 +1,451 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use crate::classifier::{Optimizer, gradient, model::FhClassifier};
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
+pub struct Sgd {
+    parameters: Vec<f32>,
+    bias: f32,
+    alpha: f64,
+    l1_ratio: f64,
+    l2_ratio: f64,
+    t: f64,
+    w_scale: f32,
+    optimal_init: f64,
+    eta: f32,
+    u: f32,
+    q: Vec<f32>,
+}
+
+impl Sgd {
+    pub fn new(n_features: usize, alpha: f64, l1_ratio: f64, l2_ratio: f64) -> Self {
+        let typw = (1.0 / alpha.sqrt()).sqrt();
+        let initial_eta0 = typw / 1.0_f64.max(gradient(1.0, -typw as f32) as f64);
+        let optimal_init = 1.0 / (initial_eta0 * alpha);
+
+        Sgd {
+            parameters: vec![0.0; n_features],
+            bias: 0.0,
+            alpha,
+            l1_ratio,
+            l2_ratio,
+            t: 0.0,
+            w_scale: 1.0,
+            optimal_init,
+            eta: initial_eta0 as f32,
+            u: 0.0,
+            q: vec![0.0; n_features],
+        }
+    }
+
+    pub fn with_initial_parameters(self, value: f32) -> Self {
+        Sgd {
+            parameters: vec![value; self.parameters.len()],
+            ..self
+        }
+    }
+
+    fn maybe_rescale(&mut self) {
+        if !(1e-6..=1e6).contains(&self.w_scale) {
+            for w in &mut self.parameters {
+                *w *= self.w_scale;
+            }
+            self.w_scale = 1.0;
+        }
+    }
+
+    #[inline(always)]
+    fn apply_l1_penalty(&mut self) {
+        if self.l1_ratio > 0.0 {
+            for (z, q) in self.parameters.iter_mut().zip(self.q.iter_mut()) {
+                let z_orig = *z;
+                let scaled_z = *z * self.w_scale;
+                if scaled_z > 0.0 {
+                    *z = (*z - (self.u + *q) / self.w_scale).max(0.0);
+                } else if scaled_z < 0.0 {
+                    *z = (*z + (self.u - *q) / self.w_scale).min(0.0);
+                }
+                *q += self.w_scale * (z_orig - *z);
+            }
+        }
+    }
+}
+
+impl Optimizer for Sgd {
+    fn step(&mut self) {
+        self.t += 1.0;
+        self.eta = (1.0 / ((self.alpha) * (self.optimal_init + self.t - 1.0))) as f32;
+        self.w_scale *= 1.0 - ((1.0 - self.l1_ratio) as f32 * self.eta * self.l2_ratio as f32);
+        self.u += self.eta * self.l1_ratio as f32 * self.alpha as f32;
+    }
+
+    fn update_param(&mut self, i: usize, g: f32) {
+        self.parameters[i] += (-self.eta * g) / self.w_scale;
+    }
+
+    fn update_bias(&mut self, g: f32) {
+        self.bias += -self.eta * g;
+        self.maybe_rescale();
+        self.apply_l1_penalty();
+    }
+
+    #[inline(always)]
+    fn get_param(&self, idx: usize) -> f32 {
+        self.parameters[idx] * self.w_scale
+    }
+
+    #[inline(always)]
+    fn get_bias(&self) -> f32 {
+        self.bias
+    }
+
+    #[inline(always)]
+    fn get_param_mut(&mut self, idx: usize) -> &mut f32 {
+        &mut self.parameters[idx]
+    }
+
+    fn build_classifier(&self) -> FhClassifier {
+        FhClassifier {
+            parameters: self.parameters.iter().map(|w| w * self.w_scale).collect(),
+            bias: self.bias,
+        }
+    }
+
+    fn num_parameters(&self) -> usize {
+        self.parameters.len()
+    }
+}
+
+#[cfg(test)]
+pub mod tests {
+    use crate::classifier::{
+        Optimizer,
+        adam::Adam,
+        feature::{
+            CcfhFeature, CcfhFeatureBuilder, FeatureBuilder, FhFeature, FhFeatureBuilder, Sample,
+            UnprocessedFeature,
+        },
+        ftrl::Ftrl,
+        train::{CcfhTrainer, FhTrainer},
+    };
+    use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};
+    use std::{
+        collections::HashMap,
+        fs::File,
+        io::{BufRead, BufReader},
+        time::Instant,
+    };
+
+    #[ignore]
+    #[test]
+    fn text_classifier() {
+        let reader = BufReader::new(
+            File::open("/Users/me/code/playground/phishing_email.csv")
+                .expect("Could not open file"),
+        );
+        let mut samples = Vec::with_capacity(1024);
+
+        let time = Instant::now();
+
+        for line in reader.lines().skip(1) {
+            let line = line.unwrap();
+            let (text, class) = line.trim().rsplit_once(',').unwrap();
+            //let (class, text) = line.trim().split_once(',').unwrap();
+            let text = text.trim_start_matches('"').trim_end_matches('"');
+
+            samples.push((text.to_string(), class == "1"));
+        }
+
+        println!("Loaded {} samples in {:?}", samples.len(), time.elapsed());
+
+        samples.shuffle(&mut StdRng::seed_from_u64(42));
+
+        let (train_samples, test_samples) = train_test_split(&samples, 0.2);
+
+        println!(
+            "Training samples: {}, Testing samples: {}",
+            train_samples.len(),
+            test_samples.len()
+        );
+
+        const FH_SIZE: usize = 16;
+        const CCFH_SIZE: usize = FH_SIZE - 2;
+        let mut rng = StdRng::seed_from_u64(42);
+
+        let fh_builder = FhFeatureBuilder {
+            weight_mask: (1 << FH_SIZE) - 1,
+        };
+        let mut fh_train_samples = build_fh_samples(train_samples.as_slice(), &fh_builder);
+        fh_train_samples.shuffle(&mut rng);
+        let fh_test_samples = build_fh_samples(test_samples.as_slice(), &fh_builder);
+        let ccfh_builder = CcfhFeatureBuilder {
+            weight_mask: (1 << FH_SIZE) - 1,
+            indicator_mask: (1 << CCFH_SIZE) - 1,
+        };
+        let mut ccfh_train_samples = build_ccfh_samples(train_samples.as_slice(), &ccfh_builder);
+        ccfh_train_samples.shuffle(&mut rng);
+        let ccfh_test_samples = build_ccfh_samples(test_samples.as_slice(), &ccfh_builder);
+
+        fh_model_stats(
+            "FTRL",
+            FhTrainer::new(Ftrl::new(1 << FH_SIZE)),
+            &fh_train_samples,
+            &fh_test_samples,
+        );
+
+        ccfh_model_stats(
+            "FTRL + FTRL",
+            CcfhTrainer::new(
+                Ftrl::new(1 << FH_SIZE),
+                Ftrl::new(1 << CCFH_SIZE).with_initial_weights(0.5),
+            ),
+            &ccfh_train_samples,
+            &ccfh_test_samples,
+        );
+
+        fh_model_stats(
+            "Adam",
+            FhTrainer::new(Adam::new(1 << FH_SIZE, 0.01)),
+            &fh_train_samples,
+            &fh_test_samples,
+        );
+
+        ccfh_model_stats(
+            "Adam + Adam",
+            CcfhTrainer::new(
+                Adam::new(1 << FH_SIZE, 0.01),
+                Adam::new(1 << CCFH_SIZE, 0.01).with_initial_weights(0.5),
+            ),
+            &ccfh_train_samples,
+            &ccfh_test_samples,
+        );
+
+        /*fh_model_stats(
+            "SGD",
+            FhTrainer::new(Sgd::new(1 << FH_SIZE, 0.0001, 0.0, 0.0001)),
+            &fh_train_samples,
+            &fh_test_samples,
+        );
+
+        ccfh_model_stats(
+            "FTRL + SGD",
+            CcfhTrainer::new(
+                Ftrl::new(1 << FH_SIZE),
+                Sgd::new(1 << CCFH_SIZE, 0.0001, 0.0, 0.0001).with_initial_parameters(0.5),
+            ),
+            &ccfh_train_samples,
+            &ccfh_test_samples,
+        );*/
+    }
+
+    fn fh_model_stats(
+        name: &str,
+        mut model: FhTrainer<impl Optimizer>,
+        train_samples: &[Sample<FhFeature>],
+        test_samples: &[Sample<FhFeature>],
+    ) {
+        print!("⏳ Training {}... ", name);
+        let time = Instant::now();
+        let mut batch = Vec::new();
+        for sample in train_samples {
+            batch.push(sample);
+            if batch.len() == 128 {
+                model.fit(&mut batch, 5);
+                batch.clear();
+            }
+        }
+        if !batch.is_empty() {
+            model.fit(&mut batch, 5);
+        }
+        println!(" trained in {:?}", time.elapsed());
+        let y_pred = model
+            .build_classifier()
+            .predict_batch(test_samples.iter().map(|s| &s.features));
+        let y_train: Vec<f32> = test_samples.iter().map(|s| s.class).collect();
+        println!("Accuracy: {:.4}", accuracy_score(&y_train, &y_pred));
+        println!("Precision: {:.4}", precision_score(&y_train, &y_pred, 1.0));
+        println!("Recall: {:.4}", recall_score(&y_train, &y_pred, 1.0));
+        println!("F1 Score: {:.4}", f1_score(&y_train, &y_pred, 1.0));
+    }
+
+    fn ccfh_model_stats(
+        name: &str,
+        mut model: CcfhTrainer<impl Optimizer, impl Optimizer>,
+        train_samples: &[Sample<CcfhFeature>],
+        test_samples: &[Sample<CcfhFeature>],
+    ) {
+        print!("⏳ Training {}... ", name);
+        let time = Instant::now();
+        let mut batch = Vec::new();
+        for sample in train_samples {
+            batch.push(sample);
+            if batch.len() == 128 {
+                model.fit(&mut batch, 5);
+                batch.clear();
+            }
+        }
+        if !batch.is_empty() {
+            model.fit(&mut batch, 5);
+        }
+        println!(" trained in {:?}", time.elapsed());
+        let y_pred = model
+            .build_classifier()
+            .predict_batch(test_samples.iter().map(|s| &s.features));
+        let y_train: Vec<f32> = test_samples.iter().map(|s| s.class).collect();
+        println!("Accuracy: {:.4}", accuracy_score(&y_train, &y_pred));
+        println!("Precision: {:.4}", precision_score(&y_train, &y_pred, 1.0));
+        println!("Recall: {:.4}", recall_score(&y_train, &y_pred, 1.0));
+        println!("F1 Score: {:.4}", f1_score(&y_train, &y_pred, 1.0));
+    }
+
+    fn accuracy_score(y_true: &[f32], y_pred: &[f32]) -> f32 {
+        y_true
+            .iter()
+            .zip(y_pred.iter())
+            .filter(|(true_val, pred_val)| **true_val == **pred_val)
+            .count() as f32
+            / y_true.len() as f32
+    }
+
+    fn precision_score(y_true: &[f32], y_pred: &[f32], positive_class: f32) -> f32 {
+        let true_positives = y_true
+            .iter()
+            .zip(y_pred.iter())
+            .filter(|(true_val, pred_val)| {
+                **pred_val == positive_class && **true_val == positive_class
+            })
+            .count() as f32;
+
+        let predicted_positives = y_pred
+            .iter()
+            .filter(|pred_val| **pred_val == positive_class)
+            .count() as f32;
+
+        if predicted_positives == 0.0 {
+            0.0
+        } else {
+            true_positives / predicted_positives
+        }
+    }
+
+    fn recall_score(y_true: &[f32], y_pred: &[f32], positive_class: f32) -> f32 {
+        let true_positives = y_true
+            .iter()
+            .zip(y_pred.iter())
+            .filter(|(true_val, pred_val)| {
+                **pred_val == positive_class && **true_val == positive_class
+            })
+            .count() as f32;
+
+        let actual_positives = y_true
+            .iter()
+            .filter(|true_val| **true_val == positive_class)
+            .count() as f32;
+
+        if actual_positives == 0.0 {
+            0.0
+        } else {
+            true_positives / actual_positives
+        }
+    }
+
+    fn f1_score(y_true: &[f32], y_pred: &[f32], positive_class: f32) -> f32 {
+        let precision = precision_score(y_true, y_pred, positive_class);
+        let recall = recall_score(y_true, y_pred, positive_class);
+
+        if precision + recall == 0.0 {
+            0.0
+        } else {
+            2.0 * (precision * recall) / (precision + recall)
+        }
+    }
+
+    #[allow(clippy::type_complexity)]
+    pub fn train_test_split(
+        data: &[(String, bool)],
+        test_size: f32,
+    ) -> (Vec<(&String, bool)>, Vec<(&String, bool)>) {
+        let mut class_0: Vec<(&String, bool)> = Vec::new();
+        let mut class_1: Vec<(&String, bool)> = Vec::new();
+
+        for (sample, class) in data {
+            if !*class {
+                class_0.push((sample, *class));
+            } else {
+                class_1.push((sample, *class));
+            }
+        }
+
+        let test_count_0 = (class_0.len() as f32 * test_size).round() as usize;
+        let test_count_1 = (class_1.len() as f32 * test_size).round() as usize;
+
+        let (test_0, train_0) = class_0.split_at(test_count_0);
+        let (test_1, train_1) = class_1.split_at(test_count_1);
+
+        let mut train = Vec::new();
+        let mut test = Vec::new();
+
+        train.extend_from_slice(train_0);
+        train.extend_from_slice(train_1);
+        test.extend_from_slice(test_0);
+        test.extend_from_slice(test_1);
+
+        (train, test)
+    }
+
+    pub fn build_fh_samples(
+        data: &[(&String, bool)],
+        builder: &FhFeatureBuilder,
+    ) -> Vec<Sample<FhFeature>> {
+        let mut samples = Vec::with_capacity(data.len());
+
+        for (text, class) in data {
+            let mut sample: HashMap<String, f32> = HashMap::new();
+            for word in text.split_whitespace() {
+                *sample.entry(word.to_string()).or_default() += 1.0;
+            }
+            builder.scale(&mut sample);
+            samples.push(Sample {
+                features: builder.build(&sample, 12345.into(), true),
+                class: if *class { 1.0 } else { 0.0 },
+            });
+        }
+
+        samples
+    }
+
+    pub fn build_ccfh_samples(
+        data: &[(&String, bool)],
+        builder: &CcfhFeatureBuilder,
+    ) -> Vec<Sample<CcfhFeature>> {
+        let mut samples = Vec::with_capacity(data.len());
+
+        for (text, class) in data {
+            let mut sample: HashMap<String, f32> = HashMap::new();
+            for word in text.split_whitespace() {
+                *sample.entry(word.to_string()).or_default() += 1.0;
+            }
+            builder.scale(&mut sample);
+            samples.push(Sample {
+                features: builder.build(&sample, 12345.into(), true),
+                class: if *class { 1.0 } else { 0.0 },
+            });
+        }
+
+        samples
+    }
+
+    impl UnprocessedFeature for String {
+        fn prefix(&self) -> u16 {
+            0
+        }
+
+        fn value(&self) -> &[u8] {
+            self.as_bytes()
+        }
+    }
+}

crates/nlp/src/classifier/train.rs

@@ -0,0 +1,157 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use crate::classifier::{
+    MAX_DLOSS, Optimizer,
+    feature::{CcfhFeature, CcfhFeatureBuilder, FhFeature, FhFeatureBuilder, Sample},
+    gradient,
+    model::{CcfhClassifier, FhClassifier},
+};
+use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
+pub struct FhTrainer<T: Optimizer> {
+    pub optimizer: T,
+}
+
+#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
+pub struct CcfhTrainer<W: Optimizer, I: Optimizer> {
+    pub w_optimizer: W,
+    pub i_optimizer: I,
+}
+
+impl<T: Optimizer> FhTrainer<T> {
+    pub fn new(optimizer: T) -> Self {
+        FhTrainer { optimizer }
+    }
+
+    pub fn fit(&mut self, samples: &mut [impl AsRef<Sample<FhFeature>>], num_epochs: usize) {
+        for _ in 0..num_epochs {
+            samples.shuffle(&mut StdRng::seed_from_u64(42));
+
+            for sample in samples.iter() {
+                let sample = sample.as_ref();
+                let mut dot: f32 = 0.0;
+                for f in &sample.features {
+                    dot += self.optimizer.get_param(f.idx) * f.weight;
+                }
+                let p = dot + self.optimizer.get_bias();
+                let dloss = gradient(sample.class, p).clamp(-MAX_DLOSS, MAX_DLOSS);
+
+                self.optimizer.step();
+
+                for f in &sample.features {
+                    self.optimizer.update_param(f.idx, dloss * f.weight);
+                }
+
+                self.optimizer.update_bias(dloss);
+            }
+        }
+    }
+
+    pub fn feature_builder(&self) -> FhFeatureBuilder {
+        FhFeatureBuilder {
+            weight_mask: (self.optimizer.num_parameters() - 1) as u64,
+        }
+    }
+
+    pub fn build_classifier(&self) -> FhClassifier {
+        self.optimizer.build_classifier()
+    }
+
+    pub fn optimizer(&self) -> &T {
+        &self.optimizer
+    }
+
+    pub fn optimizer_mut(&mut self) -> &mut T {
+        &mut self.optimizer
+    }
+}
+
+impl<W: Optimizer, I: Optimizer> CcfhTrainer<W, I> {
+    pub fn new(w_optimizer: W, i_optimizer: I) -> Self {
+        CcfhTrainer {
+            w_optimizer,
+            i_optimizer,
+        }
+    }
+
+    pub fn fit(&mut self, samples: &mut [impl AsRef<Sample<CcfhFeature>>], num_epochs: usize) {
+        for _ in 0..num_epochs {
+            samples.shuffle(&mut StdRng::seed_from_u64(42));
+
+            for sample in samples.iter() {
+                let sample = sample.as_ref();
+                let mut dot: f32 = 0.0;
+                for f in &sample.features {
+                    let q = self.i_optimizer.get_param(f.idx_i);
+                    let v1 = self.w_optimizer.get_param(f.idx_w1);
+                    let v2 = self.w_optimizer.get_param(f.idx_w2);
+                    dot += (q * v1 + (1.0 - q) * v2) * f.weight;
+                }
+                let p = dot + self.w_optimizer.get_bias();
+                let dloss = gradient(sample.class, p).clamp(-MAX_DLOSS, MAX_DLOSS);
+
+                self.w_optimizer.step();
+                self.i_optimizer.step();
+
+                for f in &sample.features {
+                    let q = self.i_optimizer.get_param(f.idx_i);
+                    let v1 = self.w_optimizer.get_param(f.idx_w1);
+                    let v2 = self.w_optimizer.get_param(f.idx_w2);
+
+                    // Update weights
+                    let d_v1 = f.weight * q;
+                    let d_v2 = f.weight * (1.0 - q);
+                    self.w_optimizer.update_param(f.idx_w1, dloss * d_v1);
+                    self.w_optimizer.update_param(f.idx_w2, dloss * d_v2);
+
+                    // Update indicator
+                    let d_q = (v1 - v2) * f.weight;
+                    self.i_optimizer.update_param(f.idx_i, dloss * d_q);
+                    let fi = self.i_optimizer.get_param_mut(f.idx_i);
+                    *fi = fi.clamp(0.0, 1.0);
+                }
+
+                self.w_optimizer.update_bias(dloss);
+            }
+        }
+    }
+
+    pub fn feature_builder(&self) -> CcfhFeatureBuilder {
+        CcfhFeatureBuilder {
+            weight_mask: (self.w_optimizer.num_parameters() - 1) as u64,
+            indicator_mask: (self.i_optimizer.num_parameters() - 1) as u64,
+        }
+    }
+
+    pub fn build_classifier(&self) -> CcfhClassifier {
+        let w_classifier = self.w_optimizer.build_classifier();
+        let i_classifier = self.i_optimizer.build_classifier();
+
+        CcfhClassifier {
+            parameters: w_classifier.parameters,
+            indicators: i_classifier.parameters,
+            bias: w_classifier.bias,
+        }
+    }
+
+    pub fn w_optimizer(&self) -> &W {
+        &self.w_optimizer
+    }
+
+    pub fn w_optimizer_mut(&mut self) -> &mut W {
+        &mut self.w_optimizer
+    }
+
+    pub fn i_optimizer(&self) -> &I {
+        &self.i_optimizer
+    }
+
+    pub fn i_optimizer_mut(&mut self) -> &mut I {
+        &mut self.i_optimizer
+    }
+}