Spam filter performance and accuracy improvements (part 8)

This commit is contained in:
mdecimus
2025-12-09 22:46:12 +01:00
parent 6c4d28a877
commit efd1d86ff3
21 changed files with 1780 additions and 661 deletions

View File

@@ -7,9 +7,12 @@
use super::server::tls::{build_self_signed_cert, parse_certificates};
use crate::{
CacheSwap, Caches, Data, DavResource, DavResources, MailboxCache, MessageStoreCache,
MessageUidCache, SpamClassifier, TlsConnectors,
MessageUidCache, TlsConnectors,
auth::{AccessToken, roles::RolePermissions},
config::smtp::resolver::{Policy, Tlsa},
config::{
smtp::resolver::{Policy, Tlsa},
spamfilter::SpamClassifier,
},
listener::blocked::BlockedIps,
manager::webadmin::WebAdminManager,
};

View File

@@ -7,7 +7,7 @@
use super::{Variable, functions::ResolveVariable, if_block::IfBlock, tokenizer::TokenMap};
use ahash::AHashSet;
use mail_auth::common::resolver::ToReverseName;
use nlp::classifier::sgd::TextClassifier;
use nlp::classifier::model::{CcfhClassifier, FhClassifier};
use std::{
net::{IpAddr, SocketAddr},
time::Duration,
@@ -20,12 +20,17 @@ use utils::{
};
#[derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize, Debug, Default)]
pub struct SpamClassifierModel {
pub classifier: TextClassifier,
pub ham_count: u64,
pub spam_count: u64,
pub last_sample_expiry: u64,
pub last_trained_at: u64,
pub enum SpamClassifier {
FhClassifier {
classifier: FhClassifier,
last_trained_at: u64,
},
CcfhClassifier {
classifier: CcfhClassifier,
last_trained_at: u64,
},
#[default]
Disabled,
}
#[derive(Debug, Clone, Default)]
@@ -75,10 +80,10 @@ pub enum SpamFilterAction<T> {
#[derive(Debug, Clone, Default)]
pub struct ClassifierConfig {
pub epochs: usize,
pub feature_hash_size: usize,
pub alpha: f32,
pub train_batch_size: usize,
pub w_params: FtrlParameters,
pub i_params: Option<FtrlParameters>,
pub num_epochs: usize,
pub reservoir_capacity: usize,
pub min_ham_samples: u64,
pub min_spam_samples: u64,
pub auto_learn_reply_ham: bool,
@@ -87,6 +92,15 @@ pub struct ClassifierConfig {
pub train_frequency: Option<u64>,
}
#[derive(Debug, Clone, Default)]
pub struct FtrlParameters {
pub feature_hash_size: usize,
pub alpha: f64,
pub beta: f64,
pub l1_ratio: f64,
pub l2_ratio: f64,
}
#[derive(Debug, Clone)]
pub struct PyzorConfig {
pub address: SocketAddr,
@@ -443,25 +457,29 @@ impl ClassifierConfig {
return None;
}
let feature_hash_size: usize = config
.property_or_default("spam-filter.classifier.parameters.features", "1048576")
.unwrap_or(1048576);
if !feature_hash_size.is_power_of_two() {
config.new_build_error(
"spam-filter.classifier.parameters.features",
"Feature size must be a power of two.",
);
}
let w_params = FtrlParameters::parse(config, "spam-filter.classifier.parameters", 20);
let i_params = if config
.property_or_default("spam-filter.classifier.ccfh.enable", "false")
.unwrap_or(false)
{
Some(FtrlParameters::parse(
config,
"spam-filter.classifier.ccfh.parameters",
w_params.feature_hash_size - 2,
))
} else {
None
};
ClassifierConfig {
feature_hash_size,
epochs: config
.property_or_default("spam-filter.classifier.parameters.epochs", "1000")
.unwrap_or(1000),
alpha: config
.property_or_default("spam-filter.classifier.parameters.alpha", "0.00001")
.unwrap_or(0.00001),
w_params,
i_params,
num_epochs: config
.property_or_default("spam-filter.classifier.training.epochs", "3")
.unwrap_or(3),
reservoir_capacity: config
.property_or_default("spam-filter.classifier.reservoir-capacity", "1024")
.unwrap_or(1024),
auto_learn_card_is_ham: config
.property_or_default("spam-filter.card-is-ham.learn", "true")
.unwrap_or(true),
@@ -478,9 +496,6 @@ impl ClassifierConfig {
min_spam_samples: config
.property_or_default("spam-filter.classifier.samples.min-spam", "10")
.unwrap_or(10),
train_batch_size: config
.property_or_default("spam-filter.classifier.training.batch-size", "100")
.unwrap_or(100),
train_frequency: config
.property_or_default::<Option<Duration>>(
"spam-filter.classifier.training.frequency",
@@ -493,6 +508,43 @@ impl ClassifierConfig {
}
}
impl FtrlParameters {
pub fn parse(config: &mut Config, prefix: &str, default_features: usize) -> Self {
let feature_hash_size: usize = config
.property((prefix, "features"))
.unwrap_or(default_features);
if !(16..=28).contains(&feature_hash_size) {
config.new_build_error(
(prefix, "features"),
"Feature size must be between 2^16 and 2^28.",
);
}
FtrlParameters {
feature_hash_size: 1 << feature_hash_size,
alpha: config
.property_or_default((prefix, "alpha"), "2.0")
.unwrap_or(2.0),
beta: config
.property_or_default((prefix, "beta"), "1.0")
.unwrap_or(1.0),
l1_ratio: config
.property_or_default((prefix, "l1"), "0.001")
.unwrap_or(0.001),
l2_ratio: config
.property_or_default((prefix, "l2"), "0.0001")
.unwrap_or(0.0001),
}
}
}
impl SpamClassifier {
pub fn is_active(&self) -> bool {
!matches!(self, SpamClassifier::Disabled)
}
}
impl SpamFilterScoreConfig {
pub fn parse(config: &mut Config) -> Self {
SpamFilterScoreConfig {

View File

@@ -5,7 +5,7 @@
*/
use crate::{
Inner, Server, SpamClassifier,
Inner, Server,
auth::{AccessToken, ResourceToken, TenantInfo},
config::{
smtp::{
@@ -15,9 +15,10 @@ use crate::{
QueueStrategy, RequireOptional, RoutingStrategy, TlsStrategy, VirtualQueue,
},
},
spamfilter::SpamClassifierModel,
spamfilter::SpamClassifier,
},
ipc::{BroadcastEvent, PushEvent, PushNotification},
manager::SPAM_CLASSIFIER_KEY,
};
use directory::{Directory, QueryParams, Type, backend::internal::manage::ManageDirectory};
use mail_auth::IpLookupStrategy;
@@ -41,7 +42,7 @@ use types::{
blob::{BlobClass, BlobId},
blob_hash::BlobHash,
collection::{Collection, SyncCollection},
field::{Field, PrincipalField},
field::Field,
type_state::{DataType, StateChange},
};
use utils::{map::bitmap::Bitmap, snowflake::SnowflakeIdGenerator};
@@ -1044,41 +1045,20 @@ impl Server {
}
pub async fn spam_model_reload(&self) -> trc::Result<()> {
if let Some(config) = &self.core.spam.classifier {
if self.core.spam.classifier.is_some() {
if let Some(model) = self
.store()
.get_value::<Archive<AlignedBytes>>(ValueKey::property(
u32::MAX,
Collection::Principal,
u32::MAX,
PrincipalField::SpamModel,
))
.blob_store()
.get_blob(SPAM_CLASSIFIER_KEY, 0..usize::MAX)
.await
.and_then(|archive| match archive {
Some(archive) => archive.deserialize::<SpamClassifierModel>().map(Some),
Some(archive) => <Archive<AlignedBytes> as Deserialize>::deserialize(&archive)
.and_then(|archive| archive.deserialize_untrusted::<SpamClassifier>())
.map(Some),
None => Ok(None),
})
.caused_by(trc::location!())?
{
if model.ham_count >= config.min_ham_samples
&& model.spam_count >= config.min_spam_samples
{
self.inner
.data
.spam_classifier
.store(Arc::new(SpamClassifier {
model: model.classifier,
last_trained_at: model.last_trained_at,
}));
} else {
trc::event!(
Spam(SpamEvent::ModelNotReady),
Details = vec![
trc::Value::from(model.ham_count),
trc::Value::from(model.spam_count)
],
);
}
self.inner.data.spam_classifier.store(Arc::new(model));
} else {
trc::event!(Spam(SpamEvent::ModelNotFound));
}

View File

@@ -28,7 +28,6 @@ use ipc::{BroadcastEvent, HousekeeperEvent, PushEvent, QueueEvent, ReportingEven
use listener::{asn::AsnGeoLookupData, blocked::Security, tls::AcmeProviders};
use mail_auth::{MX, Txt};
use manager::webadmin::{Resource, WebAdminManager};
use nlp::classifier::sgd::TextClassifier;
use parking_lot::{Mutex, RwLock};
use rustls::sign::CertifiedKey;
use std::{
@@ -73,6 +72,8 @@ pub mod enterprise;
pub use psl;
use crate::config::spamfilter::SpamClassifier;
pub static VERSION_PRIVATE: &str = env!("CARGO_PKG_VERSION");
pub static VERSION_PUBLIC: &str = "1.0.0";
@@ -131,12 +132,6 @@ pub struct Inner {
pub ipc: Ipc,
}
#[derive(Default)]
pub struct SpamClassifier {
pub model: TextClassifier,
pub last_trained_at: u64,
}
pub struct Data {
pub spam_classifier: ArcSwap<SpamClassifier>,

View File

@@ -21,6 +21,8 @@ pub mod webadmin;
const DEFAULT_SPAMFILTER_URL: &str =
"https://github.com/stalwartlabs/spam-filter/releases/latest/download/spam-filter.toml";
pub const WEBADMIN_KEY: &[u8] = "STALWART_WEBADMIN".as_bytes();
pub const SPAM_TRAINER_KEY: &[u8] = "STALWART_SPAM_TRAIN_DATA.lz4".as_bytes();
pub const SPAM_CLASSIFIER_KEY: &[u8] = "STALWART_SPAM_CLASSIFIER_MODEL.lz4".as_bytes();
// SPDX-SnippetBegin
// SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>