1use crate::corpus::{self, KEEP_FREE, Place, Record};
58use crate::detail::{self, Detail, Origin, Tested};
59use crate::eval::{self, Numbers};
60use crate::external;
61use crate::labels::{self, Copy, Evidence, Label};
62use crate::platt::{self, Platt};
63use crate::spamcore::{self as spam, N, Table};
64use crate::svm::{self, HAM, Problem, SPAM};
65use crate::{Cancel, Dirs, LearnError, Progress, Stage, io_error, now};
66use serde::{Deserialize, Serialize};
67use sioul_core::card::Card;
68use sioul_core::config::Config;
69use sioul_core::spam::tokenize::Lexicon;
70use sioul_core::words::SpamWords;
71use std::collections::{BTreeMap, HashMap, HashSet};
72use std::io::{BufReader, BufWriter, Seek, Write};
73use std::path::{Path, PathBuf};
74
75#[derive(Debug, Clone, PartialEq)]
77pub struct Options {
78 pub dim: u32,
80 pub epochs: u32,
83 pub min_count: u32,
85 pub bucket: u32,
87 pub minn: u32,
89 pub maxn: u32,
90 pub threads: u32,
92 pub seed: u64,
94 pub ham_weight: f64,
96 pub costs: Vec<f64>,
98 pub threshold_spam: f64,
100 pub threshold_unsure: f64,
101 pub least_of_each: u64,
103 pub replace: bool,
106 pub model: Model,
108 pub lexicon: SpamWords,
112}
113
114#[derive(Debug, Clone, Copy, PartialEq)]
118pub enum Model {
119 Supervised { word_ngrams: u32, lr: f64 },
123 Centroid,
124}
125
126impl Default for Model {
127 fn default() -> Self {
128 Model::Supervised { word_ngrams: 1, lr: 0.1 }
129 }
130}
131
132impl Model {
133 pub fn epochs(&self) -> u32 {
135 match self {
136 Model::Supervised { .. } => 25,
137 Model::Centroid => 5,
138 }
139 }
140
141 pub fn name(&self) -> &'static str {
143 match self {
144 Model::Supervised { .. } => "supervised",
145 Model::Centroid => "centroid",
146 }
147 }
148}
149
150impl Default for Options {
151 fn default() -> Self {
152 Options {
153 dim: 100,
154 epochs: Model::default().epochs(),
155 min_count: 5,
156 bucket: 200_000,
157 minn: 3,
158 maxn: 6,
159 threads: std::thread::available_parallelism().map_or(4, |n| n.get() as u32),
160 seed: 1,
161 ham_weight: 5.0,
162 costs: vec![0.01, 0.1, 1.0, 10.0],
163 threshold_spam: 0.95,
164 threshold_unsure: 0.5,
165 least_of_each: 20,
166 replace: true,
167 model: Model::default(),
168 lexicon: sioul_core::words::Words::builtin_ref().spam.clone(),
169 }
170 }
171}
172
173impl Options {
174 pub fn of(config: &Config) -> Options {
177 let (spam, unsure) = config.spam.thresholds();
178 let lexicon = sioul_core::words::Words::of(config).spam.clone();
179 Options { threshold_spam: f64::from(spam), threshold_unsure: f64::from(unsure), lexicon, ..Options::default() }
180 }
181}
182
183#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
185pub struct Summary {
186 pub trained_at: i64,
188 pub seconds: f64,
189 pub tokenizer: u32,
190 pub features: u32,
191 pub replaced: bool,
193 pub reason: String,
197 pub threshold_spam: f64,
198 pub threshold_unsure: f64,
199 pub labels: labels::Summary,
200 pub split: Split,
201 pub accounts: BTreeMap<String, Counts>,
203 pub model: ModelSummary,
204 pub test: Numbers,
207 #[serde(default, skip_serializing_if = "Option::is_none")]
210 pub test_settled: Option<Numbers>,
211 #[serde(default, skip_serializing_if = "Option::is_none")]
214 pub current: Option<Numbers>,
215 #[serde(default, skip_serializing_if = "Option::is_none")]
217 pub compared: Option<Compared>,
218 #[serde(default, skip_serializing_if = "Option::is_none")]
220 pub refit: Option<Refit>,
221 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
223 pub outside: BTreeMap<String, Outside>,
224 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
228 pub trial: bool,
229}
230
231#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
234pub struct Outside {
235 pub train_ham: u64,
236 pub train_spam: u64,
237 pub held_ham: u64,
238 pub held_spam: u64,
239 #[serde(default, skip_serializing_if = "Option::is_none")]
242 pub baseline: Option<Numbers>,
243}
244
245#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
248pub struct Compared {
249 pub since: i64,
251 pub ham: u64,
252 pub spam: u64,
253 pub new_ham_lost: u64,
254 pub current_ham_lost: u64,
255}
256
257#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
260pub struct Refit {
261 pub ham: u64,
262 pub spam: u64,
263 pub platt_a: f64,
266 pub platt_b: f64,
267 pub calibration_scores: u64,
268 pub fold_error: f64,
270}
271
272#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
273pub struct Split {
274 pub train_ham: u64,
275 pub train_spam: u64,
276 pub test_ham: u64,
277 pub test_spam: u64,
278 pub test_from: i64,
280}
281
282#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
283pub struct Counts {
284 pub train_ham: u64,
285 pub train_spam: u64,
286 pub test_ham: u64,
287 pub test_spam: u64,
288}
289
290#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
293pub struct ModelSummary {
294 #[serde(default, skip_serializing_if = "String::is_empty")]
296 pub kind: String,
297 #[serde(default, skip_serializing_if = "is_zero")]
299 pub word_ngrams: u32,
300 #[serde(default, skip_serializing_if = "is_zero_f")]
301 pub lr: f64,
302 pub dim: u32,
303 pub vocabulary: u64,
305 pub bucket: u32,
306 pub minn: u32,
307 pub maxn: u32,
308 pub epochs: u32,
309 #[serde(default)]
311 pub min_count: u32,
312 #[serde(default)]
314 pub threads: u32,
315 #[serde(default, skip_serializing_if = "Vec::is_empty")]
317 pub costs: Vec<f64>,
318 pub c: f64,
320 pub validation_auc: BTreeMap<String, f64>,
321 pub ham_weight: f64,
322 pub platt_a: f64,
323 pub platt_b: f64,
324 pub calibration_scores: u64,
326 pub fold_error: f64,
328 pub table_bytes: u64,
330 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
333 pub header_weights: BTreeMap<String, f64>,
334}
335
336fn is_zero(n: &u32) -> bool {
337 *n == 0
338}
339
340fn is_zero_f(x: &f64) -> bool {
341 *x == 0.0
342}
343
344pub fn last(dirs: &Dirs) -> Option<Summary> {
346 std::fs::read_to_string(dirs.trained()).ok().and_then(|t| toml::from_str(&t).ok())
347}
348
349pub fn run(config: &Config, dirs: &Dirs, progress: &mut dyn FnMut(&Progress), cancel: &Cancel) -> Result<(corpus::Update, Summary), LearnError> {
354 let update = corpus::update(&config.accounts, &|a| sioul_sync::secret::password(a), dirs, &corpus::free_space, progress, cancel);
355 if update.cancelled {
356 return Err(LearnError::Cancelled);
357 }
358 let summary = train(dirs, &trusted_ids(config), &Options::of(config), progress, cancel)?;
359 Ok((update, summary))
360}
361
362pub fn trusted_ids(config: &Config) -> BTreeMap<String, Vec<String>> {
365 config.mail_sources().into_iter().filter_map(|s| Some((s.account?, s.trusted_ids))).collect()
366}
367
368pub(crate) struct Message {
370 pub(crate) label: Label,
371 pub(crate) date: i64,
372 pub(crate) account: String,
373 pub(crate) place: Option<Place>,
375 pub(crate) evidence: Option<Evidence>,
376 pub(crate) source: Option<usize>,
378 pub(crate) ordinal: u64,
380 pub(crate) unsettled: bool,
382 pub(crate) known: bool,
385 pub(crate) features: [f32; N],
386 pub(crate) offset: u64,
388 pub(crate) length: u64,
389 pub(crate) words: u64,
390}
391
392struct TokenFile(PathBuf);
394
395impl Drop for TokenFile {
396 fn drop(&mut self) {
397 let _ = std::fs::remove_file(&self.0);
398 }
399}
400
401const LINE_WORDS: usize = 1000;
403
404pub fn train(dirs: &Dirs, trusted: &BTreeMap<String, Vec<String>>, options: &Options, progress: &mut dyn FnMut(&Progress), cancel: &Cancel) -> Result<Summary, LearnError> {
407 train_with(dirs, trusted, options, 0, progress, cancel).map(|(summary, _)| summary)
408}
409
410pub fn train_with(dirs: &Dirs, trusted: &BTreeMap<String, Vec<String>>, options: &Options, errors: usize, progress: &mut dyn FnMut(&Progress), cancel: &Cancel) -> Result<(Summary, Detail), LearnError> {
415 let started = std::time::Instant::now();
416 let started_at = now();
417 let check = || if cancel.cancelled() { Err(LearnError::Cancelled) } else { Ok(()) };
418
419 let mut copies = Vec::new();
421 let mut texts = 0u64;
422 corpus::read_all(dirs, |record| {
423 texts += record.plain.as_ref().map_or(0, |t| t.text.len() as u64) + record.html.as_ref().map_or(0, |t| t.text.len() as u64) + record.header.len() as u64;
424 copies.push(Copy::of(&record));
425 if copies.len() % 1000 == 0 {
426 progress(&Progress { stage: Stage::Labels, done: copies.len() as u64, total: 0, detail: String::new() });
427 }
428 })?;
429 let log = labels::read_log(dirs);
430 let (labeled, label_summary) = labels::decide(copies, &log, &labels::read_moved(dirs), &labels::read_flagged(dirs));
431 let own_keys: HashSet<String> = labeled.iter().map(|l| l.key.clone()).collect();
433 let yours = |message: &external::Message| message.message_id().is_some_and(|id| own_keys.contains(&id));
434 let mut outside: BTreeMap<String, Outside> = BTreeMap::new();
435 external::read_all(dirs, |source, message| {
436 if yours(&message) {
437 return;
438 }
439 texts += (message.subject.len() + message.text.len() + message.html.as_ref().map_or(0, String::len)) as u64;
440 let counts = outside.entry(source.to_string()).or_default();
441 match message.label {
442 Label::Ham => counts.train_ham += 1,
443 Label::Spam => counts.train_spam += 1,
444 }
445 })?;
446 let sources: Vec<String> = outside.keys().cloned().collect();
447 progress(&Progress { stage: Stage::Labels, done: label_summary.records, total: label_summary.records, detail: String::new() });
448 let (ham, spam) = outside.values().fold((label_summary.ham, label_summary.spam), |(h, s), o| (h + o.train_ham, s + o.train_spam));
449 if ham < options.least_of_each || spam < options.least_of_each {
450 return Err(LearnError::TooFew { ham, spam });
451 }
452 check()?;
453
454 let model_bytes = (u64::from(options.bucket) + 2 * (label_summary.records + ham + spam).min(500_000)) * u64::from(options.dim) * 4;
456 std::fs::create_dir_all(&dirs.cache).map_err(|e| io_error(&dirs.cache, e))?;
457 if let Some(free) = corpus::free_space(&dirs.cache)
458 && free < KEEP_FREE + texts + model_bytes
459 {
460 return Err(LearnError::Disk(format!("{} MB free, {} MB needed above the 1 GB kept free", free >> 20, (texts + model_bytes) >> 20)));
461 }
462
463 let chosen: HashMap<Place, usize> = labeled.iter().enumerate().map(|(i, l)| (l.place.clone(), i)).collect();
465 let token_path = dirs.cache.join(format!("tokens-{}.txt", std::process::id()));
466 let token_file = TokenFile(token_path.clone());
467 let mut messages: Vec<Option<Message>> = (0..labeled.len()).map(|_| None).collect();
468 let mut outside_messages: Vec<Message> = Vec::new();
469 let tokens_total = labeled.len() as u64 + (ham + spam) - (label_summary.ham + label_summary.spam);
471 let lexicon = Lexicon::cached(&options.lexicon);
472 {
473 let file = std::fs::File::create(&token_path).map_err(|e| io_error(&token_path, e))?;
474 let mut out = BufWriter::new(file);
475 let mut offset = 0u64;
476 let mut done = 0u64;
477 let mut failure = None;
478 corpus::read_all(dirs, |record| {
479 if failure.is_some() {
480 return;
481 }
482 let Some(&index) = chosen.get(&record.place()) else { return };
483 if messages[index].is_some() {
484 return;
485 }
486 let (words, features) = read_message(&record, trusted, &lexicon);
487 let mut length = 0u64;
488 for line in words.chunks(LINE_WORDS) {
489 let text = line.join(" ");
490 if let Err(e) = out.write_all(text.as_bytes()).and_then(|()| out.write_all(b"\n")) {
491 failure = Some(io_error(&token_path, e));
492 return;
493 }
494 length += text.len() as u64 + 1;
495 }
496 let l = &labeled[index];
497 messages[index] = Some(Message {
498 label: l.label,
499 date: l.date,
500 account: l.place.account.clone(),
501 place: Some(l.place.clone()),
502 evidence: Some(l.evidence),
503 source: None,
504 ordinal: 0,
505 unsettled: detail::unsettled(l.label, l.role, l.date, started_at),
506 known: true,
507 features,
508 offset,
509 length,
510 words: words.len() as u64,
511 });
512 offset += length;
513 done += 1;
514 if done % 500 == 0 {
515 progress(&Progress { stage: Stage::Tokens, done, total: tokens_total, detail: String::new() });
516 }
517 })?;
518 let mut ranks: HashMap<String, u64> = HashMap::new();
520 external::read_all(dirs, |source, message| {
521 let rank = ranks.entry(source.to_string()).or_default();
522 let ordinal = *rank;
523 *rank += 1;
524 if failure.is_some() || yours(&message) {
525 return;
526 }
527 let words = message.words(&lexicon);
528 let mut length = 0u64;
529 for line in words.chunks(LINE_WORDS) {
530 let text = line.join(" ");
531 if let Err(e) = out.write_all(text.as_bytes()).and_then(|()| out.write_all(b"\n")) {
532 failure = Some(io_error(&token_path, e));
533 return;
534 }
535 length += text.len() as u64 + 1;
536 }
537 let source = sources.iter().position(|s| s == source);
538 outside_messages.push(Message {
539 label: message.label,
540 date: message.date,
541 account: String::new(),
542 place: None,
543 evidence: None,
544 source,
545 ordinal,
546 unsettled: false,
547 known: false,
548 features: [0.0; N],
549 offset,
550 length,
551 words: words.len() as u64,
552 });
553 offset += length;
554 done += 1;
555 if done % 500 == 0 {
556 progress(&Progress { stage: Stage::Tokens, done, total: tokens_total, detail: String::new() });
557 }
558 })?;
559 if let Some(e) = failure {
560 return Err(e);
561 }
562 out.flush().map_err(|e| io_error(&token_path, e))?;
563 }
564 let mut messages: Vec<Message> = messages.into_iter().flatten().collect();
566 messages.append(&mut outside_messages);
567 progress(&Progress { stage: Stage::Tokens, done: messages.len() as u64, total: messages.len() as u64, detail: String::new() });
568 check()?;
569
570 if let Model::Supervised { word_ngrams, lr } = options.model {
572 let words = words_of(&token_path, &messages, &(0..messages.len()).collect::<Vec<_>>())?;
573 drop(token_file);
574 let read = Read { messages, words, sources, outside, labels: label_summary, started };
575 return crate::supervised::train(dirs, read, options, word_ngrams, lr, errors, progress, cancel);
576 }
577
578 progress(&Progress { stage: Stage::Language, done: 0, total: 0, detail: String::new() });
580 let args = fasttext::args::Args {
581 input: token_path.clone(),
582 model: fasttext::args::ModelName::SkipGram,
583 loss: fasttext::args::LossName::NegativeSampling,
584 dim: options.dim as i32,
585 epoch: options.epochs as i32,
586 min_count: options.min_count as i32,
587 bucket: options.bucket as i32,
588 minn: options.minn as i32,
589 maxn: options.maxn as i32,
590 thread: options.threads.max(1) as i32,
591 seed: options.seed as i32,
592 verbose: 0,
593 ..fasttext::args::Args::default()
594 };
595 let language = fasttext::FastText::train_with_abort(args, cancel.flag()).map_err(|e| LearnError::Language(e.to_string()))?;
596 check()?;
597 progress(&Progress { stage: Stage::Language, done: 1, total: 1, detail: String::new() });
598
599 let dim = options.dim as usize;
601 let vectors = message_vectors(&language, &token_path, &messages, options.threads.max(1) as usize, progress, cancel)?;
602 check()?;
603
604 let Parts { order, train: train_ids, test: test_ids, held } = parts(&messages, sources.len());
609 let (train_ids, test_ids) = (train_ids.as_slice(), test_ids.as_slice());
610 let test_words = words_of(&token_path, &messages, test_ids)?;
611 let held_words = held.iter().map(|ids| words_of(&token_path, &messages, ids)).collect::<Result<Vec<_>, _>>()?;
612 drop(token_file);
613 let rows = Rows { vectors: &vectors, messages: &messages, dim };
614 let width = rows.width();
615
616 let (mean, deviation) = rows.standardization(train_ids);
619 let fits = (options.costs.len() + 5) as u64;
620 let mut fitted = 0u64;
621 let (c, validation_auc, model, sigmoid, calibration_scores) = {
622 let train_x = rows.standardized(train_ids, &mean, &deviation);
623 let train_y = rows.labels(train_ids);
624 let n_train = train_ids.len();
625 let fit_end = (n_train as f64 * 0.8).round() as usize;
626 let mut validation_auc = BTreeMap::new();
627 let mut best: Option<(f64, f64)> = None;
628 for &c in &options.costs {
629 check()?;
630 let model = fit_prefix(&train_x, &train_y, width, fit_end, c, options, cancel)?;
631 let scores: Vec<(f64, bool)> = (fit_end..n_train).map(|r| (model.decision(&train_x[r * width..(r + 1) * width]), train_y[r] == SPAM)).collect();
632 let quality = eval::auc(&scores).unwrap_or_else(|| -hinge_loss(&scores, options.ham_weight));
634 validation_auc.insert(format!("{c}"), quality);
635 if best.is_none_or(|(_, q)| quality > q + 1e-9) {
636 best = Some((c, quality));
637 }
638 fitted += 1;
639 progress(&Progress { stage: Stage::Classifier, done: fitted, total: fits, detail: format!("C = {c}") });
640 }
641 let c = best.map_or(0.1, |(c, _)| c);
642 let (model, sigmoid, calibration_scores) = fit_with_platt(&train_x, &train_y, width, c, options, cancel, &mut |detail| {
643 fitted += 1;
644 progress(&Progress { stage: Stage::Classifier, done: fitted, total: fits, detail: detail.to_string() });
645 })?;
646 (c, validation_auc, model, sigmoid, calibration_scores)
647 };
648 check()?;
649
650 progress(&Progress { stage: Stage::Evaluation, done: 0, total: test_ids.len() as u64, detail: String::new() });
652 let (split, accounts) = counted(&messages, train_ids, test_ids);
653 let table = Table::from_bytes(&fold(&language, &model, &mean, &deviation, sigmoid, options, &split).to_bytes()).map_err(LearnError::Fold)?;
655 let test_x = rows.standardized(test_ids, &mean, &deviation);
656 let fold_error = fold_error(&table, &model, &test_x, width, &test_words, test_ids, &messages)?;
657 let new_scored: Vec<(f64, bool)> = test_ids.iter().enumerate().map(|(row, &i)| (f64::from(table.score(&test_words[row], &messages[i].features).p), messages[i].label == Label::Spam)).collect();
658 drop(test_x);
659 let test = eval::evaluate(&new_scored, options.threshold_spam, options.threshold_unsure);
660 let settled: Vec<(f64, bool)> = test_ids.iter().zip(&new_scored).filter(|&(&i, _)| !messages[i].unsettled).map(|(_, &s)| s).collect();
662 let test_settled = Some(eval::evaluate(&settled, options.threshold_spam, options.threshold_unsure));
663 let lacking = lacking(&table);
665 let mut tested: Vec<Tested> = test_ids.iter().zip(&new_scored).map(|(&i, &(p, _))| tested_of(&messages[i], p, &sources)).collect();
666 for (s, name) in sources.iter().enumerate() {
667 let scored: Vec<(f64, bool)> = held[s].iter().zip(&held_words[s]).map(|(&i, words)| (f64::from(table.score(words, &lacking).p), messages[i].label == Label::Spam)).collect();
668 tested.extend(held[s].iter().zip(&scored).map(|(&i, &(p, _))| tested_of(&messages[i], p, &sources)));
669 held_out(outside.entry(name.clone()).or_default(), &scored, options);
670 }
671 let detail = detail::of(dirs, &tested, options.threshold_spam, options.threshold_unsure, errors)?;
673 drop(tested);
674 progress(&Progress { stage: Stage::Evaluation, done: test_ids.len() as u64, total: test_ids.len() as u64, detail: String::new() });
675
676 let (current_numbers, compared, replaced, reason) = against_current(dirs, test_ids, &test_words, &messages, &new_scored, options)?;
678 check()?;
679
680 let replaced = replaced && options.replace;
683 let mut table_bytes = table.to_bytes().len() as u64;
684 let refit = if replaced {
685 let refits = 6u64;
686 let mut done = 0u64;
687 let mut step = |detail: &str| {
688 done += 1;
689 progress(&Progress { stage: Stage::Export, done, total: refits, detail: detail.to_string() });
690 };
691 let (mut written, refit) = refit_all(&language, &rows, &order, test_ids, &test_words, c, options, cancel, &mut step)?;
692 written.meta.metrics = metrics(&test);
693 check()?;
694 written.write(&dirs.table()).map_err(|e| io_error(&dirs.table(), e))?;
696 table_bytes = std::fs::metadata(dirs.table()).map_or(table_bytes, |m| m.len());
697 let fresh = dirs.language().with_extension("bin.new");
698 language.save_model(&fresh).map_err(|e| io_error(&fresh, e))?;
699 std::fs::rename(&fresh, dirs.language()).map_err(|e| io_error(&dirs.language(), e))?;
700 step("");
701 Some(refit)
702 } else {
703 progress(&Progress { stage: Stage::Export, done: 1, total: 1, detail: String::new() });
704 None
705 };
706 let summary = Summary {
707 trained_at: now(),
708 seconds: started.elapsed().as_secs_f64(),
709 tokenizer: spam::TOKENIZER,
710 features: spam::FEATURES,
711 replaced,
712 reason: reason.to_string(),
713 threshold_spam: options.threshold_spam,
714 threshold_unsure: options.threshold_unsure,
715 labels: label_summary,
716 split,
717 accounts,
718 model: ModelSummary {
719 kind: "centroid".to_string(),
720 word_ngrams: 0,
721 lr: 0.0,
722 dim: options.dim,
723 vocabulary: language.dict().nwords().max(1) as u64 - 1,
724 bucket: options.bucket,
725 minn: options.minn,
726 maxn: options.maxn,
727 epochs: options.epochs,
728 min_count: options.min_count,
729 threads: options.threads,
730 costs: options.costs.clone(),
731 c,
732 validation_auc,
733 ham_weight: options.ham_weight,
734 platt_a: sigmoid.a,
735 platt_b: sigmoid.b,
736 calibration_scores,
737 fold_error,
738 table_bytes,
739 header_weights: (0..N).map(|h| (spam::NAMES[h].to_string(), model.w[dim + h])).collect(),
740 },
741 test,
742 test_settled,
743 current: current_numbers,
744 compared,
745 refit,
746 outside,
747 trial: !options.replace,
748 };
749 if options.replace {
751 let text = toml::to_string(&summary).map_err(|e| io_error(&dirs.trained(), e))?;
752 corpus::write_whole(&dirs.trained(), text.as_bytes())?;
753 }
754 Ok((summary, detail))
755}
756
757pub(crate) struct Read {
759 pub(crate) messages: Vec<Message>,
760 pub(crate) words: Vec<Vec<String>>,
762 pub(crate) sources: Vec<String>,
763 pub(crate) outside: BTreeMap<String, Outside>,
764 pub(crate) labels: labels::Summary,
765 pub(crate) started: std::time::Instant,
766}
767
768pub(crate) struct Parts {
772 pub(crate) order: Vec<usize>,
773 pub(crate) train: Vec<usize>,
774 pub(crate) test: Vec<usize>,
775 pub(crate) held: Vec<Vec<usize>>,
776}
777
778pub(crate) fn parts(messages: &[Message], sources: usize) -> Parts {
782 let by_date = |ids: &mut Vec<usize>| ids.sort_by_key(|&i| (messages[i].date, i));
783 let mut order: Vec<usize> = (0..messages.len()).collect();
784 by_date(&mut order);
785 let newest_fifth = |source: Option<usize>| -> (Vec<usize>, Vec<usize>) {
786 let ids: Vec<usize> = order.iter().copied().filter(|&i| messages[i].source == source).collect();
787 let (mut learned, mut tested) = newest_fifths(&ids, |i| (messages[i].account.as_str(), messages[i].label == Label::Spam));
788 by_date(&mut learned);
789 by_date(&mut tested);
790 (learned, tested)
791 };
792 let (mut train, test) = newest_fifth(None);
793 let mut held: Vec<Vec<usize>> = Vec::with_capacity(sources);
794 for s in 0..sources {
795 let (learned, held_out) = newest_fifth(Some(s));
796 train.extend(learned);
797 held.push(held_out);
798 }
799 by_date(&mut train);
800 Parts { order, train, test, held }
801}
802
803pub(crate) fn counted(messages: &[Message], train: &[usize], test: &[usize]) -> (Split, BTreeMap<String, Counts>) {
805 let mut split = Split::default();
806 let mut accounts: BTreeMap<String, Counts> = BTreeMap::new();
807 for &i in train {
808 if messages[i].source.is_some() {
809 continue;
810 }
811 let counts = accounts.entry(messages[i].account.clone()).or_default();
812 match messages[i].label {
813 Label::Ham => (split.train_ham += 1, counts.train_ham += 1),
814 Label::Spam => (split.train_spam += 1, counts.train_spam += 1),
815 };
816 }
817 for &i in test {
818 let counts = accounts.entry(messages[i].account.clone()).or_default();
819 match messages[i].label {
820 Label::Ham => (split.test_ham += 1, counts.test_ham += 1),
821 Label::Spam => (split.test_spam += 1, counts.test_spam += 1),
822 };
823 }
824 split.test_from = test.first().map_or(0, |&i| messages[i].date);
825 (split, accounts)
826}
827
828pub(crate) fn held_out(counts: &mut Outside, scored: &[(f64, bool)], options: &Options) {
830 counts.held_ham = scored.iter().filter(|(_, spam)| !spam).count() as u64;
831 counts.held_spam = scored.len() as u64 - counts.held_ham;
832 counts.train_ham -= counts.held_ham;
833 counts.train_spam -= counts.held_spam;
834 counts.baseline = (!scored.is_empty()).then(|| eval::evaluate(scored, options.threshold_spam, options.threshold_unsure));
835}
836
837pub(crate) fn tested_of(message: &Message, p: f64, sources: &[String]) -> Tested {
839 let origin = match (&message.place, message.source) {
840 (Some(place), _) => Origin::Own { place: place.clone(), evidence: message.evidence.unwrap_or(Evidence::Folder) },
841 (None, source) => Origin::Outside { source: source.and_then(|s| sources.get(s)).cloned().unwrap_or_default(), ordinal: message.ordinal },
842 };
843 Tested { p, label: message.label, date: message.date, origin, unsettled: message.unsettled }
844}
845
846pub(crate) fn newest_fifths<'a>(ids: &[usize], group: impl Fn(usize) -> (&'a str, bool)) -> (Vec<usize>, Vec<usize>) {
851 let mut groups: BTreeMap<(&str, bool), Vec<usize>> = BTreeMap::new();
852 for &i in ids {
853 groups.entry(group(i)).or_default().push(i);
854 }
855 let (mut learned, mut tested) = (Vec::new(), Vec::new());
856 for (_, mut ids) in groups {
857 let newest = ids.split_off(ids.len() - (ids.len() as f64 * 0.2).round() as usize);
858 learned.extend(ids);
859 tested.extend(newest);
860 }
861 (learned, tested)
862}
863
864pub(crate) fn metrics(test: &Numbers) -> BTreeMap<String, f64> {
866 let mut metrics = BTreeMap::from([
867 ("threshold_spam".to_string(), test.at_spam.threshold),
868 ("threshold_unsure".to_string(), test.at_unsure.threshold),
869 ("ham_called_spam".to_string(), test.at_spam.ham_called_spam.fraction()),
870 ("spam_caught".to_string(), test.at_spam.spam_caught.fraction()),
871 ("ham_called_unsure".to_string(), test.at_unsure.ham_called_spam.fraction()),
872 ("spam_caught_unsure".to_string(), test.at_unsure.spam_caught.fraction()),
873 ("unsure".to_string(), test.unsure.fraction()),
874 ]);
875 if let Some(auc) = test.auc {
876 metrics.insert("auc".to_string(), auc);
877 }
878 metrics
879}
880
881pub(crate) fn against_current(dirs: &Dirs, test_ids: &[usize], test_words: &[Vec<String>], messages: &[Message], new_scored: &[(f64, bool)], options: &Options) -> Result<(Option<Numbers>, Option<Compared>, bool, &'static str), LearnError> {
893 let current = match Table::read(&dirs.table()) {
894 Ok(current) => Ok(Some(current)),
895 Err(_) if !dirs.table().exists() => Ok(None),
896 Err(e) => Err(e),
897 };
898 let (current_numbers, compared) = match ¤t {
899 Ok(Some(current)) => {
900 let since = current.meta.trained_at;
901 let mut unseen: Vec<usize> = (0..test_ids.len()).filter(|&row| messages[test_ids[row]].date > since).collect();
902 let lexicon = current.lexicon();
903 let own_words: Option<HashMap<usize, Vec<String>>> = if lexicon.words == options.lexicon || unseen.is_empty() {
904 None
905 } else {
906 let wanted: HashMap<Place, usize> = unseen.iter().filter_map(|&row| Some((messages[test_ids[row]].place.clone()?, row))).collect();
907 let mut read: HashMap<usize, Vec<String>> = HashMap::new();
908 corpus::read_all(dirs, |record| {
909 if let Some(&row) = wanted.get(&record.place()) {
910 read.entry(row).or_insert_with(|| words_of_record(&record, &lexicon));
911 }
912 })?;
913 unseen.retain(|row| read.contains_key(row));
914 Some(read)
915 };
916 let words = |row: usize| own_words.as_ref().and_then(|own| own.get(&row)).unwrap_or(&test_words[row]);
917 let label = |row: usize| messages[test_ids[row]].label == Label::Spam;
918 let theirs: Vec<(f64, bool)> = unseen.iter().map(|&row| (f64::from(current.score(words(row), &messages[test_ids[row]].features).p), label(row))).collect();
919 let ours: Vec<(f64, bool)> = unseen.iter().map(|&row| new_scored[row]).collect();
920 let compared = Compared { since, ham: theirs.iter().filter(|(_, s)| !s).count() as u64, spam: theirs.iter().filter(|(_, s)| *s).count() as u64, new_ham_lost: eval::evaluate(&ours, options.threshold_spam, options.threshold_unsure).at_spam.ham_called_spam.count, current_ham_lost: 0 };
921 let numbers = eval::evaluate(&theirs, options.threshold_spam, options.threshold_unsure);
922 (Some(numbers.clone()), Some(Compared { current_ham_lost: numbers.at_spam.ham_called_spam.count, ..compared }))
923 }
924 _ => (None, None),
925 };
926 let (replaced, reason) = match (¤t, &compared) {
927 (Err(_), _) => (true, "unreadable"),
928 (Ok(None), _) => (true, "no-table"),
929 (Ok(Some(_)), Some(c)) => replacement(c.new_ham_lost, c.current_ham_lost, c.ham + c.spam),
930 (Ok(Some(_)), None) => (true, "no-table"),
931 };
932 Ok((current_numbers, compared, replaced, reason))
933}
934
935pub(crate) fn replacement(new_ham_lost: u64, current_ham_lost: u64, compared: u64) -> (bool, &'static str) {
940 if compared == 0 {
941 (true, "nothing-new")
942 } else if new_ham_lost <= current_ham_lost {
943 (true, "no-worse")
944 } else {
945 (false, "worse")
946 }
947}
948
949pub(crate) fn read_message(record: &Record, trusted: &BTreeMap<String, Vec<String>>, lexicon: &Lexicon) -> (Vec<String>, [f32; N]) {
956 let Some(card) = spam::card_of(record) else { return (Vec::new(), [0.0; N]) };
957 let own = record.checked.as_ref().and_then(|c| c.reading(record.date, sioul_core::spam::features::signed(&card.headers)));
958 features_of(&card, trusted.get(&record.account).map_or(&[][..], Vec::as_slice), (record.date > 0).then_some(record.date), own, lexicon)
959}
960
961pub(crate) fn features_of(card: &Card, trusted: &[String], internal_date: Option<i64>, own: Option<sioul_core::trust::AuthResults>, lexicon: &Lexicon) -> (Vec<String>, [f32; N]) {
966 let tokens = spam::tokens_with(lexicon, &card.subject, &card.excerpt);
967 let auth = own.or_else(|| sioul_core::spam::features::auth_results(&card.headers, trusted));
968 let features = spam::features(card, auth.as_ref(), internal_date, &tokens.links);
969 (kept(tokens.words), features)
970}
971
972fn words_of_record(record: &Record, lexicon: &Lexicon) -> Vec<String> {
976 spam::card_of(record).map(|card| kept(spam::tokens_with(lexicon, &card.subject, &card.excerpt).words)).unwrap_or_default()
977}
978
979fn kept(words: Vec<String>) -> Vec<String> {
981 words.into_iter().filter(|w| !w.is_empty() && !w.contains(char::is_whitespace)).collect()
982}
983
984struct Rows<'a> {
986 vectors: &'a [f64],
987 messages: &'a [Message],
988 dim: usize,
989}
990
991impl Rows<'_> {
992 fn width(&self) -> usize {
993 self.dim + N
994 }
995
996 fn raw(&self, i: usize) -> impl Iterator<Item = f64> + '_ {
997 self.vectors[i * self.dim..(i + 1) * self.dim].iter().copied().chain(self.messages[i].features.iter().map(|&x| f64::from(x)))
998 }
999
1000 fn standardization(&self, ids: &[usize]) -> (Vec<f64>, Vec<f64>) {
1004 standardization(|k| self.raw(ids[k]).collect(), |k| self.messages[ids[k]].known, ids.len(), self.width(), self.dim)
1005 }
1006
1007 fn standardized(&self, ids: &[usize], mean: &[f64], deviation: &[f64]) -> Vec<f32> {
1011 ids.iter()
1012 .flat_map(|&i| {
1013 let known = self.messages[i].known;
1014 self.raw(i)
1015 .zip(mean)
1016 .zip(deviation)
1017 .enumerate()
1018 .map(move |(c, ((x, m), s))| if *s > 0.0 && (known || c < self.dim) && !x.is_nan() { ((x - m) / s) as f32 } else { 0.0 })
1019 .collect::<Vec<f32>>()
1020 })
1021 .collect()
1022 }
1023
1024 fn labels(&self, ids: &[usize]) -> Vec<i8> {
1025 ids.iter().map(|&i| if self.messages[i].label == Label::Spam { SPAM } else { HAM }).collect()
1026 }
1027}
1028
1029fn standardization(row: impl Fn(usize) -> Vec<f64>, known: impl Fn(usize) -> bool, n: usize, width: usize, dim: usize) -> (Vec<f64>, Vec<f64>) {
1034 let counted = |column: usize, k: usize, x: f64| (column < dim || known(k)) && !x.is_nan();
1035 let mut counts = vec![0usize; width];
1036 let mut mean = vec![0.0; width];
1037 for k in 0..n {
1038 for (c, ((m, x), count)) in mean.iter_mut().zip(row(k)).zip(counts.iter_mut()).enumerate() {
1039 if counted(c, k, x) {
1040 *m += x;
1041 *count += 1;
1042 }
1043 }
1044 }
1045 let counts: Vec<f64> = counts.iter().map(|&c| c.max(1) as f64).collect();
1046 mean.iter_mut().zip(&counts).for_each(|(m, n)| *m /= n);
1047 let mut deviation = vec![0.0; width];
1048 for k in 0..n {
1049 for (c, ((d, x), m)) in deviation.iter_mut().zip(row(k)).zip(&mean).enumerate() {
1050 if counted(c, k, x) {
1051 *d += (x - m) * (x - m);
1052 }
1053 }
1054 }
1055 deviation.iter_mut().zip(&counts).for_each(|(d, n)| {
1056 *d = (*d / n).sqrt();
1057 if *d < 1e-12 {
1058 *d = 0.0;
1059 }
1060 });
1061 (mean, deviation)
1062}
1063
1064fn lacking(_table: &Table) -> [f32; N] {
1068 [f32::NAN; N]
1069}
1070
1071fn fit_prefix(x: &[f32], y: &[i8], width: usize, end: usize, c: f64, options: &Options, cancel: &Cancel) -> Result<svm::Model, LearnError> {
1073 let params = svm::Params { c, ham_weight: options.ham_weight, seed: options.seed, ..svm::Params::default() };
1074 svm::train(&Problem { x: &x[..end * width], dim: width, y: &y[..end] }, ¶ms, &|| cancel.cancelled()).map_err(|_| LearnError::Cancelled)
1075}
1076
1077fn fit_with_platt(x: &[f32], y: &[i8], width: usize, c: f64, options: &Options, cancel: &Cancel, step: &mut dyn FnMut(&str)) -> Result<(svm::Model, Platt, u64), LearnError> {
1083 let n = y.len();
1084 let scores = |model: &svm::Model, rows: std::ops::Range<usize>| -> Vec<(f64, bool)> { rows.map(|r| (model.decision(&x[r * width..(r + 1) * width]), y[r] == SPAM)).collect() };
1085 let mut calibration = Vec::new();
1086 for k in 1..5 {
1087 let (start, end) = (n * k / 5, n * (k + 1) / 5);
1088 if y[..start].contains(&SPAM) && y[..start].contains(&HAM) && end > start {
1089 let model = fit_prefix(x, y, width, start, c, options, cancel)?;
1090 calibration.extend(scores(&model, start..end));
1091 }
1092 step("Platt");
1093 }
1094 let model = fit_prefix(x, y, width, n, c, options, cancel)?;
1095 step(&format!("C = {c}"));
1096 let calibration = if platt::fit(&calibration).is_some() { calibration } else { scores(&model, 0..n) };
1097 let sigmoid = platt::fit(&calibration).unwrap_or(Platt { a: -1.0, b: 0.0 });
1098 Ok((model, sigmoid, calibration.len() as u64))
1099}
1100
1101fn fold_error(table: &Table, model: &svm::Model, x: &[f32], width: usize, words: &[Vec<String>], ids: &[usize], messages: &[Message]) -> Result<f64, LearnError> {
1104 let mut error = 0f64;
1105 for (row, &i) in ids.iter().enumerate() {
1106 let full = model.decision(&x[row * width..(row + 1) * width]);
1107 let score = table.score(&words[row], &messages[i].features);
1108 error = error.max((f64::from(score.f) - full).abs());
1109 }
1110 if error > 1e-4 {
1111 return Err(LearnError::Fold(format!("the table's score differs from the model's by {error:e}")));
1112 }
1113 Ok(error)
1114}
1115
1116#[allow(clippy::too_many_arguments)]
1122fn refit_all(
1123 language: &fasttext::FastText,
1124 rows: &Rows,
1125 order: &[usize],
1126 test_ids: &[usize],
1127 test_words: &[Vec<String>],
1128 c: f64,
1129 options: &Options,
1130 cancel: &Cancel,
1131 step: &mut dyn FnMut(&str),
1132) -> Result<(Table, Refit), LearnError> {
1133 let width = rows.width();
1134 let (mean, deviation) = rows.standardization(order);
1135 let (model, sigmoid, calibration_scores) = {
1136 let x = rows.standardized(order, &mean, &deviation);
1137 let y = rows.labels(order);
1138 fit_with_platt(&x, &y, width, c, options, cancel, step)?
1139 };
1140 let ham_of = |ids: &[usize]| ids.iter().filter(|&&i| rows.messages[i].label == Label::Ham).count() as u64;
1142 let all = Split {
1143 train_ham: ham_of(order),
1144 train_spam: order.len() as u64 - ham_of(order),
1145 test_ham: ham_of(test_ids),
1146 test_spam: test_ids.len() as u64 - ham_of(test_ids),
1147 test_from: test_ids.first().map_or(0, |&i| rows.messages[i].date),
1148 };
1149 let table = Table::from_bytes(&fold(language, &model, &mean, &deviation, sigmoid, options, &all).to_bytes()).map_err(LearnError::Fold)?;
1150 let test_x = rows.standardized(test_ids, &mean, &deviation);
1151 let fold_error = fold_error(&table, &model, &test_x, width, test_words, test_ids, rows.messages)?;
1152 Ok((table, Refit { ham: all.train_ham, spam: all.train_spam, platt_a: sigmoid.a, platt_b: sigmoid.b, calibration_scores, fold_error }))
1153}
1154
1155fn hinge_loss(scores: &[(f64, bool)], ham_weight: f64) -> f64 {
1157 let total: f64 = scores
1158 .iter()
1159 .map(|&(f, spam)| {
1160 let (y, weight) = if spam { (1.0, 1.0) } else { (-1.0, ham_weight) };
1161 weight * (1.0 - y * f).max(0.0).powi(2)
1162 })
1163 .sum();
1164 total / scores.len().max(1) as f64
1165}
1166
1167fn message_vectors(language: &fasttext::FastText, path: &Path, messages: &[Message], threads: usize, progress: &mut dyn FnMut(&Progress), cancel: &Cancel) -> Result<Vec<f64>, LearnError> {
1170 let dim = language.get_dimension() as usize;
1171 let nwords = language.dict().nwords() as usize;
1173 let mut known = vec![0f32; nwords * dim];
1174 std::thread::scope(|s| {
1175 for (t, chunk) in known.chunks_mut(nwords.div_ceil(threads).max(1) * dim).enumerate() {
1176 s.spawn(move || {
1177 let first = t * nwords.div_ceil(threads).max(1);
1178 for (k, out) in chunk.chunks_mut(dim).enumerate() {
1179 language.get_word_vector_into(language.dict().get_word((first + k) as i32), out);
1180 }
1181 });
1182 }
1183 });
1184 let mut vectors = vec![0f64; messages.len() * dim];
1185 let per = messages.len().div_ceil(threads).max(1);
1186 let done = std::sync::atomic::AtomicU64::new(0);
1187 let failed: std::sync::Mutex<Option<LearnError>> = std::sync::Mutex::new(None);
1188 std::thread::scope(|s| {
1189 let mut workers = Vec::new();
1190 for (t, out) in vectors.chunks_mut(per * dim).enumerate() {
1191 let (known, done, failed) = (&known, &done, &failed);
1192 workers.push(s.spawn(move || {
1193 let first = t * per;
1194 let mut file = match std::fs::File::open(path) {
1195 Ok(file) => file,
1196 Err(e) => {
1197 *failed.lock().unwrap_or_else(std::sync::PoisonError::into_inner) = Some(io_error(path, e));
1198 return;
1199 }
1200 };
1201 let mut word = vec![0f32; dim];
1202 for (k, vector) in out.chunks_mut(dim).enumerate() {
1203 if cancel.cancelled() {
1204 return;
1205 }
1206 let message = &messages[first + k];
1207 if message.words == 0 {
1208 continue;
1209 }
1210 let mut text = vec![0u8; message.length as usize];
1211 if let Err(e) = file.seek(std::io::SeekFrom::Start(message.offset)).and_then(|_| std::io::Read::read_exact(&mut file, &mut text)) {
1212 *failed.lock().unwrap_or_else(std::sync::PoisonError::into_inner) = Some(io_error(path, e));
1213 return;
1214 }
1215 let mut count = 0u64;
1216 for token in String::from_utf8_lossy(&text).split_ascii_whitespace() {
1217 let row = match language.get_word_id(token) {
1218 Some(id) => &known[id as usize * dim..(id as usize + 1) * dim],
1219 None => {
1220 language.get_word_vector_into(token, &mut word);
1221 &word[..]
1222 }
1223 };
1224 for (v, &x) in vector.iter_mut().zip(row) {
1225 *v += f64::from(x);
1226 }
1227 count += 1;
1228 }
1229 if count > 0 {
1230 vector.iter_mut().for_each(|v| *v /= count as f64);
1231 }
1232 done.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1233 }
1234 }));
1235 }
1236 loop {
1238 let finished = workers.iter().all(|w| w.is_finished());
1239 progress(&Progress { stage: Stage::Features, done: done.load(std::sync::atomic::Ordering::Relaxed), total: messages.len() as u64, detail: String::new() });
1240 if finished {
1241 break;
1242 }
1243 std::thread::sleep(std::time::Duration::from_millis(100));
1244 }
1245 });
1246 if let Some(e) = failed.into_inner().unwrap_or_else(std::sync::PoisonError::into_inner) {
1247 return Err(e);
1248 }
1249 if cancel.cancelled() {
1250 return Err(LearnError::Cancelled);
1251 }
1252 Ok(vectors)
1253}
1254
1255pub(crate) fn words_of(path: &Path, messages: &[Message], ids: &[usize]) -> Result<Vec<Vec<String>>, LearnError> {
1257 let file = std::fs::File::open(path).map_err(|e| io_error(path, e))?;
1258 let mut reader = BufReader::new(file);
1259 ids.iter()
1260 .map(|&i| {
1261 let message = &messages[i];
1262 let mut text = vec![0u8; message.length as usize];
1263 reader.seek(std::io::SeekFrom::Start(message.offset)).and_then(|_| std::io::Read::read_exact(&mut reader, &mut text)).map_err(|e| io_error(path, e))?;
1264 Ok(String::from_utf8_lossy(&text).split_ascii_whitespace().map(str::to_string).collect())
1265 })
1266 .collect()
1267}
1268
1269fn fold(language: &fasttext::FastText, model: &svm::Model, mean: &[f64], deviation: &[f64], sigmoid: Platt, options: &Options, split: &Split) -> Table {
1274 let dim = options.dim as usize;
1275 let u: Vec<f64> = (0..dim).map(|d| if deviation[d] > 0.0 { model.w[d] / deviation[d] } else { 0.0 }).collect();
1276 let dot = |row: &[f32]| -> f64 { u.iter().zip(row).map(|(a, &b)| a * f64::from(b)).sum() };
1277 let (vocabulary, _) = language.get_vocab();
1278 let mut vector = vec![0f32; dim];
1279 let words: Vec<(u64, f32)> = vocabulary
1280 .iter()
1281 .filter(|w| w.as_str() != fasttext::dictionary::EOS)
1282 .map(|w| {
1283 language.get_word_vector_into(w, &mut vector);
1284 (spam::word_hash(w), dot(&vector) as f32)
1285 })
1286 .collect();
1287 let nwords = i64::from(language.dict().nwords());
1288 let input = language.input_matrix();
1289 let buckets: Vec<f32> = (0..i64::from(options.bucket)).map(|b| dot(input.row(nwords + b)) as f32).collect();
1290 let text_mean = u.iter().zip(&mean[..dim]).map(|(a, m)| a * m).sum::<f64>();
1291 let weights: Vec<f32> = (0..N).map(|h| if deviation[dim + h] > 0.0 { (model.w[dim + h] / deviation[dim + h]) as f32 } else { 0.0 }).collect();
1292 let means: Vec<f32> = (0..N).map(|h| mean[dim + h] as f32).collect();
1293 Table {
1294 tokenizer: spam::TOKENIZER,
1295 features: spam::FEATURES,
1296 dim: options.dim,
1297 minn: options.minn,
1298 maxn: options.maxn,
1299 bucket: options.bucket,
1300 words,
1301 buckets,
1302 weights,
1303 means,
1304 bias: model.b as f32,
1305 text_mean: text_mean as f32,
1306 platt_a: sigmoid.a as f32,
1307 platt_b: sigmoid.b as f32,
1308 meta: spam::Meta {
1309 trained_at: now(),
1310 ham: split.train_ham,
1311 spam: split.train_spam,
1312 test_ham: split.test_ham,
1313 test_spam: split.test_spam,
1314 metrics: BTreeMap::new(),
1315 device: sioul_sync::lease::host_name(),
1316 lexicon: Some(options.lexicon.clone()),
1317 },
1318 }
1319}
1320
1321#[derive(Debug, Clone, PartialEq)]
1323pub struct Evaluation {
1324 pub labels: labels::Summary,
1325 pub split: Split,
1327 pub numbers: Numbers,
1328 pub table: spam::Meta,
1330 pub outside: BTreeMap<String, Numbers>,
1332 pub unseen: Option<Unseen>,
1337}
1338
1339#[derive(Debug, Clone, PartialEq)]
1341pub struct Unseen {
1342 pub since: i64,
1344 pub numbers: Numbers,
1345}
1346
1347pub fn evaluate(dirs: &Dirs, trusted: &BTreeMap<String, Vec<String>>, options: &Options, progress: &mut dyn FnMut(&Progress), cancel: &Cancel) -> Result<Evaluation, LearnError> {
1351 evaluate_with(dirs, trusted, options, 0, progress, cancel).map(|(evaluation, _)| evaluation)
1352}
1353
1354pub fn evaluate_with(dirs: &Dirs, trusted: &BTreeMap<String, Vec<String>>, options: &Options, errors: usize, progress: &mut dyn FnMut(&Progress), cancel: &Cancel) -> Result<(Evaluation, Detail), LearnError> {
1357 let table = Table::read(&dirs.table()).map_err(LearnError::NoTable)?;
1358 let lexicon = table.lexicon();
1360 let mut copies = Vec::new();
1361 corpus::read_all(dirs, |record| copies.push(Copy::of(&record)))?;
1362 let (mut labeled, label_summary) = labels::decide(copies, &labels::read_log(dirs), &labels::read_moved(dirs), &labels::read_flagged(dirs));
1363 progress(&Progress { stage: Stage::Labels, done: label_summary.records, total: label_summary.records, detail: String::new() });
1364 if cancel.cancelled() {
1365 return Err(LearnError::Cancelled);
1366 }
1367 labeled.sort_by(|a, b| (a.date, &a.key).cmp(&(b.date, &b.key)));
1369 let all: Vec<usize> = (0..labeled.len()).collect();
1370 let (_, newest) = newest_fifths(&all, |i| (labeled[i].place.account.as_str(), labeled[i].label == Label::Spam));
1371 let newest: HashSet<usize> = newest.into_iter().collect();
1372 let (tested_labels, labeled): (Vec<(usize, labels::Labeled)>, Vec<(usize, labels::Labeled)>) = labeled.into_iter().enumerate().partition(|(i, _)| newest.contains(i));
1373 let (tested_labels, labeled): (Vec<labels::Labeled>, Vec<labels::Labeled>) = (tested_labels.into_iter().map(|(_, l)| l).collect(), labeled.into_iter().map(|(_, l)| l).collect());
1374 let wanted: HashMap<Place, usize> = tested_labels.iter().enumerate().map(|(i, l)| (l.place.clone(), i)).collect();
1375 let mut scored: Vec<Option<(f64, bool)>> = vec![None; tested_labels.len()];
1376 let mut done = 0u64;
1377 corpus::read_all(dirs, |record| {
1378 let Some(&i) = wanted.get(&record.place()) else { return };
1379 if scored[i].is_some() || cancel.cancelled() {
1380 return;
1381 }
1382 let (words, features) = read_message(&record, trusted, &lexicon);
1383 scored[i] = Some((f64::from(table.score(&words, &features).p), tested_labels[i].label == Label::Spam));
1384 done += 1;
1385 if done % 500 == 0 {
1386 progress(&Progress { stage: Stage::Evaluation, done, total: tested_labels.len() as u64, detail: String::new() });
1387 }
1388 })?;
1389 if cancel.cancelled() {
1390 return Err(LearnError::Cancelled);
1391 }
1392 let since = table.meta.trained_at;
1394 let mut tested: Vec<Tested> = tested_labels
1395 .iter()
1396 .zip(&scored)
1397 .filter_map(|(l, scored)| scored.map(|(p, _)| Tested { p, label: l.label, date: l.date, origin: Origin::Own { place: l.place.clone(), evidence: l.evidence }, unsettled: detail::unsettled(l.label, l.role, l.date, now()) }))
1398 .collect();
1399 let unseen: Vec<(f64, bool)> = tested.iter().filter(|t| t.date > since).map(|t| (t.p, t.label == Label::Spam)).collect();
1400 let unseen = (!unseen.is_empty()).then(|| Unseen { since, numbers: eval::evaluate(&unseen, options.threshold_spam, options.threshold_unsure) });
1401 let scored: Vec<(f64, bool)> = scored.into_iter().flatten().collect();
1402 let split = Split {
1403 test_ham: scored.iter().filter(|(_, s)| !s).count() as u64,
1404 test_spam: scored.iter().filter(|(_, s)| *s).count() as u64,
1405 test_from: tested_labels.first().map_or(0, |l| l.date),
1406 ..Split::default()
1407 };
1408 let own_keys: HashSet<String> = labeled.iter().chain(&tested_labels).map(|l| l.key.clone()).collect();
1410 let mut sources: BTreeMap<String, Vec<(u64, external::Message)>> = BTreeMap::new();
1411 let mut ranks: HashMap<String, u64> = HashMap::new();
1412 external::read_all(dirs, |source, message| {
1413 let rank = ranks.entry(source.to_string()).or_default();
1414 let ordinal = *rank;
1415 *rank += 1;
1416 if !message.message_id().is_some_and(|id| own_keys.contains(&id)) {
1417 sources.entry(source.to_string()).or_default().push((ordinal, message));
1418 }
1419 })?;
1420 let lacking = lacking(&table);
1421 let mut outside = BTreeMap::new();
1422 for (name, mut messages) in sources {
1423 messages.sort_by_key(|(_, m)| m.date);
1425 let all: Vec<usize> = (0..messages.len()).collect();
1426 let (_, newest) = newest_fifths(&all, |i| ("", messages[i].1.label == Label::Spam));
1427 let newest: HashSet<usize> = newest.into_iter().collect();
1428 let held: Vec<(u64, external::Message)> = messages.into_iter().enumerate().filter(|(i, _)| newest.contains(i)).map(|(_, m)| m).collect();
1429 let scored: Vec<(f64, bool)> = held.iter().map(|(_, m)| (f64::from(table.score(&m.words(&lexicon), &lacking).p), m.label == Label::Spam)).collect();
1430 tested.extend(held.iter().zip(&scored).map(|((ordinal, m), &(p, _))| Tested { p, label: m.label, date: m.date, origin: Origin::Outside { source: name.clone(), ordinal: *ordinal }, unsettled: false }));
1431 outside.insert(name, eval::evaluate(&scored, options.threshold_spam, options.threshold_unsure));
1432 }
1433 let detail = detail::of(dirs, &tested, options.threshold_spam, options.threshold_unsure, errors)?;
1434 let numbers = eval::evaluate(&scored, options.threshold_spam, options.threshold_unsure);
1435 Ok((Evaluation { labels: label_summary, split, numbers, table: table.meta, outside, unseen }, detail))
1436}
1437
1438#[derive(Debug, Clone, PartialEq)]
1441pub struct Why {
1442 pub p: f64,
1444 pub f: f64,
1445 pub words: Vec<(String, f64)>,
1447 pub headers: Vec<(String, f64)>,
1449}
1450
1451pub fn why(dirs: &Dirs, raw: &[u8], trusted: &[String], internal_date: Option<i64>) -> Result<Why, String> {
1455 let table = Table::read(&dirs.table())?;
1456 let card = Card::from_bytes(raw).ok_or_else(|| "not a message".to_string())?;
1457 let (words, features) = features_of(&card, trusted, internal_date, None, &table.lexicon());
1458 let score = table.score(&words, &features);
1459 let words = score.words.iter().filter(|(_, s)| *s > 0.0).take(5).map(|(w, s)| (w.clone(), f64::from(*s))).collect();
1461 let headers = score.signs.iter().filter(|(_, s)| *s > 0.0).take(3).map(|&(h, s)| (spam::NAMES[h].to_string(), f64::from(s))).collect();
1462 Ok(Why { p: f64::from(score.p), f: f64::from(score.f), words, headers })
1463}
1464
1465#[cfg(test)]
1466mod tests {
1467 use super::*;
1468 use crate::synthetic;
1469
1470 fn scratch(name: &str) -> PathBuf {
1471 let dir = std::env::temp_dir().join(format!("sioul-learn-train-{name}-{}", std::process::id()));
1472 let _ = std::fs::remove_dir_all(&dir);
1473 std::fs::create_dir_all(&dir).unwrap();
1474 dir
1475 }
1476
1477 fn small() -> Options {
1480 Options { dim: 16, min_count: 2, bucket: 2000, threads: 1, seed: 7, ..Options::default() }
1481 }
1482
1483 fn centroid() -> Options {
1485 Options { epochs: Model::Centroid.epochs(), model: Model::Centroid, ..small() }
1486 }
1487
1488 #[test]
1490 fn options_read_your_words() {
1491 let config: Config = toml::from_str("[words]\nlanguages = [\"en\"]\n[words.spam.provider_tags]\nadd = [\"verdächtig\"]\n").unwrap();
1492 let lexicon = Options::of(&config).lexicon;
1493 assert!(lexicon.provider_tags.iter().any(|t| t == "verdächtig"), "{lexicon:?}");
1494 assert!(lexicon.elisions.is_empty(), "no French elisions in English alone");
1495 assert_eq!(Options::default().lexicon, sioul_core::words::Words::builtin_ref().spam);
1496 }
1497
1498 fn trusted() -> BTreeMap<String, Vec<String>> {
1499 ["home", "work"].into_iter().map(|a| (a.to_string(), vec!["mx.example.net".to_string()])).collect()
1500 }
1501
1502 fn records(mail: &[synthetic::Mail]) -> Vec<Record> {
1503 mail.iter().enumerate().map(|(i, m)| synthetic::record_of(m, 1, i as u32 + 1)).collect()
1504 }
1505
1506 fn compared_with_current(dirs: &Dirs, test_ids: &[usize], test_words: &[Vec<String>], messages: &[Message], new_scored: &[(f64, bool)], options: &Options) -> (Option<Numbers>, Option<Compared>, bool, &'static str) {
1508 against_current(dirs, test_ids, test_words, messages, new_scored, options).unwrap()
1509 }
1510
1511 #[test]
1516 fn the_table_in_place_reads_with_its_own_words() {
1517 let root = scratch("own-words");
1518 let dirs = Dirs::under(&root);
1519 let mail = |i: u32| {
1521 let spam = i % 2 == 0;
1522 let (subject, body) = if spam { ("lottery jackpot", "You won the lottery jackpot, claim the jackpot prize now") } else { ("rehearsal notes", "The rehearsal notes for the concert, see you at the rehearsal") };
1523 let date = 1_790_000_000 + i64::from(i) * 3600;
1524 let stamp = jiff::Timestamp::from_second(date).unwrap().strftime("%a, %d %b %Y %H:%M:%S +0000").to_string();
1525 let raw = format!("From: Someone <s{i}@example.org>\r\nTo: owner@example.org\r\nSubject: {subject}\r\nDate: {stamp}\r\nMessage-ID: <{i}@example.org>\r\nMIME-Version: 1.0\r\nContent-Type: text/plain; charset=utf-8\r\n\r\n{body}\r\n");
1526 let (folder, role) = if spam { ("Junk", sioul_core::folders::Role::Junk) } else { ("INBOX", sioul_core::folders::Role::Inbox) };
1527 synthetic::Mail { account: "home", folder, role, flags: Vec::new(), date, raw: raw.into_bytes() }
1528 };
1529 let mails: Vec<synthetic::Mail> = (0..20).map(mail).collect();
1530 let records = records(&mails);
1531 corpus::store(&dirs, &records).unwrap();
1532 let builtin = Lexicon::builtin();
1534 let word = |w: &str| spam::word_hash(&spam::tokens_with(&builtin, "", w).words[0]);
1535 let mut table = Table {
1536 tokenizer: spam::TOKENIZER,
1537 features: spam::FEATURES,
1538 dim: 16,
1539 minn: 3,
1540 maxn: 6,
1541 bucket: 8,
1542 words: vec![(word("lottery"), 400.0), (word("jackpot"), 400.0), (word("rehearsal"), -400.0)],
1543 buckets: vec![0.0; 8],
1544 weights: Vec::new(),
1545 means: Vec::new(),
1546 bias: 0.0,
1547 text_mean: 0.0,
1548 platt_a: -1.0,
1549 platt_b: 0.0,
1550 meta: spam::Meta { trained_at: 1_780_000_000, ..spam::Meta::default() },
1551 };
1552 table.sort();
1553 table.write(&dirs.table()).unwrap();
1554 let mut lexicon = Options::default().lexicon;
1556 lexicon.stop_words.extend(["lottery".to_string(), "jackpot".to_string()]);
1557 let options = Options { lexicon: lexicon.clone(), ..small() };
1558 let ours = Lexicon::cached(&lexicon);
1559 let trusted = BTreeMap::new();
1560 let messages: Vec<Message> = records
1561 .iter()
1562 .zip(&mails)
1563 .map(|(record, m)| Message {
1564 label: if m.role == sioul_core::folders::Role::Junk { Label::Spam } else { Label::Ham },
1565 date: m.date,
1566 account: record.account.clone(),
1567 place: Some(record.place()),
1568 evidence: None,
1569 source: None,
1570 ordinal: 0,
1571 unsettled: false,
1572 known: true,
1573 features: read_message(record, &trusted, &builtin).1,
1574 offset: 0,
1575 length: 0,
1576 words: 0,
1577 })
1578 .collect();
1579 let test_ids: Vec<usize> = (0..messages.len()).collect();
1580 let test_words: Vec<Vec<String>> = records.iter().map(|r| read_message(r, &trusted, &ours).0).collect();
1581 let new_scored: Vec<(f64, bool)> = messages.iter().map(|m| (0.0, m.label == Label::Spam)).collect();
1582 let read_with = |lexicon: &Lexicon| {
1584 let scored: Vec<(f64, bool)> = records.iter().zip(&messages).map(|(r, m)| (f64::from(table.score(&read_message(r, &trusted, lexicon).0, &m.features).p), m.label == Label::Spam)).collect();
1585 eval::evaluate(&scored, options.threshold_spam, options.threshold_unsure)
1586 };
1587 let (own, misread) = (read_with(&builtin), read_with(&ours));
1588 assert_eq!(own.at_spam.spam_caught.count, 10, "its own words: every spam caught");
1589 assert_ne!(own, misread, "the new words would make it look otherwise");
1590 let (current, compared, _, _) = compared_with_current(&dirs, &test_ids, &test_words, &messages, &new_scored, &options);
1591 assert_eq!(current.as_ref(), Some(&own), "the table in place read with its own words");
1592 assert_eq!(compared.map(|c| (c.ham, c.spam)), Some((10, 10)));
1593 let _ = std::fs::remove_dir_all(&root);
1594 }
1595
1596 #[test]
1600 fn learns_invented_spam_apart_from_ham() {
1601 learns_apart(&small(), "classifier", &[Stage::Labels, Stage::Tokens, Stage::Language, Stage::Evaluation, Stage::Export], 100);
1602 learns_apart(¢roid(), "centroid", &[Stage::Labels, Stage::Tokens, Stage::Language, Stage::Features, Stage::Classifier, Stage::Evaluation, Stage::Export], 600);
1603 }
1604
1605 fn learns_apart(options: &Options, kind: &str, steps: &[Stage], refit_calibration: u64) {
1606 let root = scratch(&format!("learns-{kind}"));
1607 let dirs = Dirs::under(&root);
1608 let mail = synthetic::mailbox(11, 500, 300);
1609 corpus::store(&dirs, &records(&mail)).unwrap();
1610 let mut stages = Vec::new();
1611 let summary = train(&dirs, &trusted(), options, &mut |p| stages.push(p.stage), &Cancel::new()).unwrap();
1612 assert_eq!(summary.model.kind, kind);
1613 assert_eq!((summary.labels.ham, summary.labels.spam), (500, 300), "{:?}", summary.labels);
1614 assert_eq!(summary.split.train_ham + summary.split.train_spam + summary.split.test_ham + summary.split.test_spam, 800);
1615 assert_eq!(summary.split.test_ham + summary.split.test_spam, 160, "the newest fifth");
1616 let test = &summary.test;
1617 assert!(test.auc.is_some_and(|a| a > 0.97), "{test:?}");
1618 assert!(test.at_unsure.spam_caught.fraction() > 0.9, "{test:?}");
1619 assert!(test.at_spam.ham_called_spam.fraction() < 0.02, "{test:?}");
1620 assert!(summary.model.platt_a < 0.0, "Platt rises with the score: {:?}", summary.model);
1621 assert!(summary.model.fold_error < 1e-4, "{:?}", summary.model);
1622 assert!(summary.model.vocabulary > 50 && summary.model.calibration_scores >= 100, "{:?}", summary.model);
1623 assert!(summary.replaced && summary.reason == "no-table");
1624 assert!(dirs.table().exists() && dirs.language().exists() && !dirs.previous_table().exists());
1625 assert_eq!(last(&dirs).as_ref(), Some(&summary), "trained.toml says it all");
1626 let refit = summary.refit.as_ref().expect("replaced: refit");
1628 assert_eq!(refit.ham + refit.spam, 800);
1629 assert_eq!(summary.split.train_ham + summary.split.train_spam, 640, "the numbers' model learned from the oldest 80 %");
1630 assert!(refit.fold_error < 1e-4 && refit.platt_a < 0.0 && refit.calibration_scores >= refit_calibration, "{refit:?}");
1631 let written = Table::read(&dirs.table()).unwrap();
1632 assert_eq!(written.kind(), if kind == "classifier" { sioul_core::spam::table::Kind::Rows } else { sioul_core::spam::table::Kind::Centroid });
1633 assert_eq!((written.meta.ham + written.meta.spam, written.meta.test_ham + written.meta.test_spam), (800, 160));
1634 assert_eq!(written.meta.metrics.get("auc").copied(), summary.test.auc);
1635 assert!(!written.meta.device.is_empty());
1636 assert_eq!(written.meta.lexicon.as_ref(), Some(&options.lexicon), "the table keeps the words it read with");
1637 let firsts: Vec<Stage> = stages.iter().fold(Vec::new(), |mut seen, s| {
1639 if seen.last() != Some(s) {
1640 seen.push(*s);
1641 }
1642 seen
1643 });
1644 assert_eq!(firsts, steps);
1645 assert_eq!(std::fs::read_dir(&dirs.cache).unwrap().count(), 0);
1646 let spam = synthetic::message(&mut crate::Rng::new(99), true, 9999, 1_790_000_000);
1648 let why = why(&dirs, &spam, &[], None).unwrap();
1649 assert!(why.p > 0.5 && why.f > 0.0, "{why:?}");
1650 assert_eq!(why.words.len(), 5);
1651 assert!(why.words.iter().chain(&why.headers).all(|(_, share)| *share > 0.0) && why.headers.len() <= 3, "{why:?}");
1652 let _ = std::fs::remove_dir_all(root);
1653 }
1654
1655 #[test]
1659 fn a_worse_table_never_replaces_a_better_one() {
1660 let root = scratch("replace");
1661 let dirs = Dirs::under(&root);
1662 let mut mail = synthetic::mailbox(5, 400, 250);
1663 let mut rng = crate::Rng::new(5);
1665 let newest = mail.iter().map(|m| m.date).max().unwrap();
1666 for k in 0..15 {
1667 let date = newest + 3600 * (k + 1);
1668 mail.push(synthetic::Mail { account: "home", folder: "INBOX", role: sioul_core::folders::Role::Inbox, flags: vec![], date, raw: synthetic::message(&mut rng, true, 50_000 + k as u32, date) });
1669 }
1670 corpus::store(&dirs, &records(&mail)).unwrap();
1671 let first = train(&dirs, &trusted(), &small(), &mut |_| {}, &Cancel::new()).unwrap();
1672 assert!(first.replaced && first.reason == "no-table" && first.refit.is_some());
1673 assert!(first.test.at_spam.ham_called_spam.count > 0, "the spam-like ham is lost: {:?}", first.test);
1674 let first_table = std::fs::read(dirs.table()).unwrap();
1675 let second = train(&dirs, &trusted(), &small(), &mut |_| {}, &Cancel::new()).unwrap();
1677 assert!(second.replaced && second.reason == "nothing-new", "{} {:?}", second.reason, second.compared);
1678 assert_eq!(second.compared.map(|c| c.ham + c.spam), Some(0));
1679 assert_eq!(std::fs::read(dirs.previous_table()).unwrap(), first_table);
1680 let midway = mail.iter().map(|m| m.date).max().unwrap() - 3600 * 40;
1683 let mut never = Table::read(&dirs.table()).unwrap();
1684 never.platt_b = 1000.0;
1685 never.meta.trained_at = midway;
1686 never.write(&dirs.table()).unwrap();
1687 let kept = std::fs::read(dirs.table()).unwrap();
1688 let third = train(&dirs, &trusted(), &small(), &mut |_| {}, &Cancel::new()).unwrap();
1689 assert!(!third.replaced && third.reason == "worse" && third.refit.is_none(), "{} {:?}", third.reason, third.compared);
1690 let compared = third.compared.unwrap();
1691 assert_eq!(compared.since, midway);
1692 assert!(compared.ham >= 15 && compared.current_ham_lost == 0 && compared.new_ham_lost > 0, "{compared:?}");
1693 assert_eq!(third.current.as_ref().map(|c| c.ham + c.spam), Some(compared.ham + compared.spam), "the same messages");
1694 assert_eq!(std::fs::read(dirs.table()).unwrap(), kept, "the table in place stays");
1695 assert_eq!(last(&dirs).map(|s| s.reason), Some("worse".to_string()));
1696 let mut always = Table::read(&dirs.table()).unwrap();
1698 always.platt_b = -1000.0;
1699 always.write(&dirs.table()).unwrap();
1700 let fourth = train(&dirs, &trusted(), &small(), &mut |_| {}, &Cancel::new()).unwrap();
1701 assert!(fourth.replaced && fourth.reason == "no-worse", "{} {:?}", fourth.reason, fourth.compared);
1702 std::fs::write(dirs.table(), b"not a table").unwrap();
1704 let fifth = train(&dirs, &trusted(), &small(), &mut |_| {}, &Cancel::new()).unwrap();
1705 assert!(fifth.replaced && fifth.reason == "unreadable");
1706 let _ = std::fs::remove_dir_all(root);
1707 }
1708
1709 #[test]
1712 fn the_newest_month_is_in_the_table_that_runs() {
1713 let root = scratch("newest");
1714 let dirs = Dirs::under(&root);
1715 let mut mail = synthetic::mailbox(17, 600, 300);
1716 let newest = mail.iter().map(|m| m.date).max().unwrap();
1717 let mut rng = crate::Rng::new(17);
1718 for k in 0..90u32 {
1719 let date = newest + 3600 * (i64::from(k) + 1);
1720 mail.push(synthetic::Mail { account: "work", folder: "Junk", role: sioul_core::folders::Role::Junk, flags: vec![], date, raw: synthetic::campaign(&mut rng, 60_000 + k, date) });
1721 }
1722 corpus::store(&dirs, &records(&mail)).unwrap();
1723 let summary = train(&dirs, &trusted(), &small(), &mut |_| {}, &Cancel::new()).unwrap();
1724 let refit = summary.refit.expect("replaced: refit");
1725 assert_eq!(refit.ham + refit.spam, 990, "every message, the campaign included");
1726 let fresh = synthetic::campaign(&mut crate::Rng::new(99), 70_000, newest + 86_400 * 3);
1728 let why = why(&dirs, &fresh, &[], None).unwrap();
1729 assert!(why.p >= 0.95, "{why:?}");
1730 let _ = std::fs::remove_dir_all(root);
1731 }
1732
1733 #[test]
1740 fn outside_material_teaches_words_and_gives_a_baseline() {
1741 let root = scratch("outside");
1742 let dirs = Dirs::under(&root);
1743 corpus::store(&dirs, &records(&synthetic::mailbox(23, 300, 150))).unwrap();
1744 let mut rng = crate::Rng::new(23);
1746 let lines: Vec<String> = (0..500)
1747 .map(|k| {
1748 let spam = k % 3 == 0;
1749 let date = 1_640_995_200 + 3600 * i64::from(k);
1750 let card = Card::from_bytes(&synthetic::message(&mut rng, spam, 90_000 + k, date)).unwrap();
1751 serde_json::json!({ "label": if spam { "spam" } else { "ham" }, "date": date, "subject": card.subject, "text": card.excerpt }).to_string()
1752 })
1753 .collect();
1754 let file = root.join("old-filter.jsonl");
1755 std::fs::write(&file, lines.join("\n")).unwrap();
1756 let imported = external::import(&dirs, &file, None).unwrap();
1757 assert_eq!((imported.sources["old-filter"].ham, imported.sources["old-filter"].spam), (333, 167));
1758 let summary = train(&dirs, &trusted(), &small(), &mut |_| {}, &Cancel::new()).unwrap();
1759 assert_eq!(summary.split.test_ham + summary.split.test_spam, 90);
1761 assert_eq!(summary.split.train_ham + summary.split.train_spam, 360, "your own mail's oldest 80 %");
1762 let outside = &summary.outside["old-filter"];
1764 assert_eq!((outside.train_ham + outside.train_spam, outside.held_ham + outside.held_spam), (400, 100), "{outside:?}");
1765 let baseline = outside.baseline.as_ref().expect("held out: measured");
1766 assert!(baseline.auc.is_some_and(|a| a > 0.9) && baseline.ham + baseline.spam == 100, "{baseline:?}");
1767 assert_eq!(summary.refit.as_ref().map(|r| r.ham + r.spam), Some(450 + 500));
1769 let evaluation = evaluate(&dirs, &trusted(), &small(), &mut |_| {}, &Cancel::new()).unwrap();
1771 assert_eq!(evaluation.outside["old-filter"].ham + evaluation.outside["old-filter"].spam, 100);
1772 external::remove(&dirs, "old-filter").unwrap();
1774 let without = train(&dirs, &trusted(), &small(), &mut |_| {}, &Cancel::new()).unwrap();
1775 assert!(without.outside.is_empty() && without.refit.map(|r| r.ham + r.spam) == Some(450));
1776 let _ = std::fs::remove_dir_all(root);
1777 }
1778
1779 #[test]
1784 fn a_trial_changes_nothing_and_names_its_errors() {
1785 let root = scratch("trial");
1786 let dirs = Dirs::under(&root);
1787 corpus::store(&dirs, &records(&synthetic::mailbox(31, 300, 200))).unwrap();
1788 let trial = Options { replace: false, ..small() };
1790 let (summary, detail) = train_with(&dirs, &trusted(), &trial, 5, &mut |_| {}, &Cancel::new()).unwrap();
1791 assert!(summary.trial && !summary.replaced && summary.reason == "no-table" && summary.refit.is_none(), "{} {}", summary.reason, summary.replaced);
1792 assert!(!dirs.table().exists() && !dirs.language().exists() && !dirs.trained().exists() && last(&dirs).is_none());
1793 assert_eq!(std::fs::read_dir(&dirs.cache).unwrap().count(), 0, "the tokenized corpus is gone");
1794 assert_eq!((summary.model.threads, summary.model.min_count, summary.model.epochs, summary.model.word_ngrams), (1, 2, 25, 1), "how it was asked to learn");
1795 assert!(summary.model.costs.is_empty() && summary.model.kind == "classifier");
1796 let tested: u64 = detail.by_account.values().map(|c| c.ham.total() + c.spam.total()).sum();
1797 assert_eq!(tested, summary.split.test_ham + summary.split.test_spam);
1798 assert_eq!(detail.grid.len(), detail::GRID.len());
1799 assert!(detail.errors.ham_called_spam.len() <= 5 && detail.errors.spam_missed.len() <= 5);
1800 for w in detail.errors.ham_called_spam.iter().chain(&detail.errors.spam_missed) {
1801 assert!(w.card.is_some() && !w.outside && !w.account.is_empty(), "{w:?}");
1802 }
1803 assert!(detail.errors.ham_called_spam.windows(2).all(|w| w[0].p >= w[1].p), "the surest first");
1804 assert!(detail.errors.spam_missed.windows(2).all(|w| w[0].p <= w[1].p), "the least sure first");
1805 let kept = train(&dirs, &trusted(), &small(), &mut |_| {}, &Cancel::new()).unwrap();
1807 assert!(kept.replaced && !kept.trial && last(&dirs).is_some_and(|l| !l.trial));
1808 let in_place = std::fs::read(dirs.table()).unwrap();
1809 let (evaluation, detail) = evaluate_with(&dirs, &trusted(), &small(), 3, &mut |_| {}, &Cancel::new()).unwrap();
1811 assert!(evaluation.unseen.is_none(), "{:?}", evaluation.unseen);
1812 assert!(detail.errors.ham_called_spam.len() <= 3 && detail.errors.spam_missed.len() <= 3);
1813 let mut rng = crate::Rng::new(31);
1815 let soon = now() + 3600;
1816 let newer: Vec<Record> = (0..40u32)
1817 .map(|k| {
1818 let spam = k % 2 == 0;
1819 let date = soon + i64::from(k) * 60;
1820 let (folder, role) = if spam { ("Junk", sioul_core::folders::Role::Junk) } else { ("INBOX", sioul_core::folders::Role::Inbox) };
1821 let mail = synthetic::Mail { account: "home", folder, role, flags: vec![], date, raw: synthetic::message(&mut rng, spam, 80_000 + k, date) };
1822 synthetic::record_of(&mail, 1, 20_000 + k)
1823 })
1824 .collect();
1825 corpus::store(&dirs, &newer).unwrap();
1826 let (evaluation, _) = evaluate_with(&dirs, &trusted(), &small(), 0, &mut |_| {}, &Cancel::new()).unwrap();
1827 let unseen = evaluation.unseen.expect("mail since the training");
1828 assert_eq!(unseen.since, Table::read(&dirs.table()).unwrap().meta.trained_at);
1829 assert_eq!(unseen.numbers.ham + unseen.numbers.spam, 40);
1830 let (second, _) = train_with(&dirs, &trusted(), &trial, 0, &mut |_| {}, &Cancel::new()).unwrap();
1832 assert!(second.trial && !second.replaced && second.compared.is_some(), "{second:?}");
1833 assert_eq!(std::fs::read(dirs.table()).unwrap(), in_place, "the table in place stays");
1834 assert!(!dirs.previous_table().exists());
1835 assert_eq!(last(&dirs).map(|l| l.trained_at), Some(kept.trained_at), "trained.toml stays the last training's");
1836 let _ = std::fs::remove_dir_all(root);
1837 }
1838
1839 #[test]
1843 fn what_outside_material_lacks_weighs_nothing() {
1844 let message = |known: bool, x: f32| Message { label: Label::Ham, date: 0, account: String::new(), place: None, evidence: None, source: (!known).then_some(0), ordinal: 0, unsettled: false, known, features: [x; N], offset: 0, length: 0, words: 0 };
1845 let messages = vec![message(true, 1.0), message(true, 3.0), message(false, 1000.0)];
1846 let vectors = vec![1.0, 2.0, 3.0];
1847 let rows = Rows { vectors: &vectors, messages: &messages, dim: 1 };
1848 let ids = [0, 1, 2];
1849 let (mean, deviation) = rows.standardization(&ids);
1850 assert_eq!((mean[1], deviation[1]), (2.0, 1.0), "your own mail's alone");
1851 assert_eq!((mean[0], deviation[0]), (2.0, (2.0f64 / 3.0).sqrt()), "the words' over all");
1852 let x = rows.standardized(&ids, &mean, &deviation);
1853 let width = rows.width();
1854 assert!(x[2 * width + 1..3 * width].iter().all(|v| *v == 0.0), "outside: 0 for every header feature");
1855 assert_ne!(x[2 * width], 0.0, "its words count");
1856 assert_eq!((x[width + 1], x[1]), (1.0, -1.0));
1857 }
1858
1859 #[test]
1860 fn too_few_or_stopped() {
1861 let root = scratch("few");
1862 let dirs = Dirs::under(&root);
1863 corpus::store(&dirs, &records(&synthetic::mailbox(3, 100, 5))).unwrap();
1864 assert!(matches!(train(&dirs, &trusted(), &small(), &mut |_| {}, &Cancel::new()), Err(LearnError::TooFew { spam: 5, .. })));
1865 corpus::store(&dirs, &records(&synthetic::mailbox(4, 60, 60))).unwrap();
1866 let cancel = Cancel::new();
1867 cancel.cancel();
1868 assert_eq!(train(&dirs, &trusted(), &small(), &mut |_| {}, &cancel).err(), Some(LearnError::Cancelled));
1869 assert!(!dirs.table().exists());
1870 let _ = std::fs::remove_dir_all(root);
1871 }
1872
1873 #[test]
1878 fn the_phones_hashing_is_fasttexts() {
1879 let root = scratch("hashing");
1880 let text = root.join("text.txt");
1881 let mut rng = crate::Rng::new(1);
1882 let lines: Vec<String> = (0..300)
1883 .map(|_| {
1884 let spam = rng.unit() < 0.5;
1885 String::from_utf8(synthetic::message(&mut rng, spam, 1, 1_750_000_000)).unwrap().replace("\r\n", " ")
1886 })
1887 .collect();
1888 std::fs::write(&text, lines.join("\n") + "\n réunion naïve été\n").unwrap();
1889 let (minn, maxn, bucket) = (3, 6, 1000);
1890 let args = fasttext::args::Args {
1891 input: text.clone(),
1892 dim: 8,
1893 epoch: 1,
1894 min_count: 1,
1895 bucket: bucket as i32,
1896 minn: minn as i32,
1897 maxn: maxn as i32,
1898 thread: 1,
1899 ..fasttext::args::Args::default()
1900 };
1901 let model = fasttext::FastText::train(args).unwrap();
1902 let nwords = model.dict().nwords();
1903 let mut checked = (0, 0);
1904 let long = "z".repeat(40);
1905 let words = ["lottery", "bonjour", "réunion", "naïve", "zzqxv", "loterrie", "é", "a", "ab", "_URL_", "élève", "🙂ok", &long];
1906 for word in words {
1907 let ids = model.dict().get_subwords_for_string(word);
1908 let expected: Vec<i32> = spam::ngram_buckets(word, minn, maxn, bucket).iter().map(|&b| nwords + b as i32).collect();
1909 match model.get_word_id(word) {
1910 Some(id) => {
1911 assert_eq!(ids[0], id, "{word}");
1912 assert_eq!(ids[1..], expected[..], "{word}");
1913 checked.0 += 1;
1914 }
1915 None => {
1916 assert_eq!(ids, expected, "{word}");
1917 let vector = model.get_word_vector(word);
1918 let input = model.input_matrix();
1919 for d in 0..8 {
1920 let mean = expected.iter().map(|&id| f64::from(input.row(i64::from(id))[d])).sum::<f64>() / expected.len() as f64;
1921 assert!((mean - f64::from(vector[d])).abs() < 1e-6, "{word}[{d}]");
1922 }
1923 checked.1 += 1;
1924 }
1925 }
1926 }
1927 assert!(checked.0 >= 4 && checked.1 >= 5, "in and out of the vocabulary: {checked:?}");
1928 let _ = std::fs::remove_dir_all(root);
1929 }
1930
1931 #[test]
1935 #[ignore]
1936 fn measure_the_spec_sizes() {
1937 let root = scratch("measure");
1938 let dirs = Dirs::under(&root);
1939 let mail = synthetic::mailbox(21, 6000, 3000);
1940 corpus::store(&dirs, &records(&mail)).unwrap();
1941 let started = std::time::Instant::now();
1942 let summary = train(&dirs, &trusted(), &Options::default(), &mut |_| {}, &Cancel::new()).unwrap();
1943 let seconds = started.elapsed().as_secs_f64();
1944 let bytes = std::fs::metadata(dirs.table()).unwrap().len();
1945 let language = std::fs::metadata(dirs.language()).unwrap().len();
1946 let corpus: u64 = std::fs::read_dir(dirs.corpus().join("home")).unwrap().chain(std::fs::read_dir(dirs.corpus().join("work")).unwrap()).filter_map(Result::ok).map(|e| e.metadata().unwrap().len()).sum();
1947 println!(
1948 "messages {}, vocabulary {}, table {bytes} bytes, language model {language} bytes, corpus {corpus} bytes, {seconds:.1} s, fold error {:e}, AUC {:?}, C {}",
1949 summary.labels.ham + summary.labels.spam,
1950 summary.model.vocabulary,
1951 summary.model.fold_error,
1952 summary.test.auc,
1953 summary.model.c
1954 );
1955 assert!(summary.model.fold_error < 1e-4);
1956 let _ = std::fs::remove_dir_all(root);
1957 }
1958
1959 #[test]
1966 #[ignore]
1967 fn write_an_invented_profile() {
1968 let Some(root) = std::env::var_os("SIOUL_LEARN_DEMO").map(PathBuf::from) else { return };
1969 let config = "[[account]]\nid = \"home\"\nkind = \"imap\"\naddress = \"owner@example.org\"\nhost = \"imap.example.org\"\ntrusted_authserv_ids = [\"mx.example.org\"]\n\n[[account]]\nid = \"work\"\nkind = \"imap\"\naddress = \"owner@example.com\"\nhost = \"imap.example.com\"\n\n[spam]\nthreshold_spam = 0.9\nthreshold_unsure = 0.5\n";
1970 std::fs::create_dir_all(root.join("config/sioul")).unwrap();
1971 std::fs::write(root.join("config/sioul/config.toml"), config).unwrap();
1972 let dirs = Dirs { data: root.join("data/sioul/spam"), state: root.join("state/sioul/spam"), cache: root.join("cache/sioul/spam") };
1973 corpus::store(&dirs, &records(&synthetic::mailbox(13, 3000, 1500))).unwrap();
1974 std::fs::write(root.join("spam.eml"), synthetic::message(&mut crate::Rng::new(77), true, 77_777, 1_790_000_000)).unwrap();
1975 }
1976
1977 #[test]
1983 fn training_and_sorting_read_the_same_checks() {
1984 let mut mail = synthetic::mailbox(41, 3, 3).swap_remove(0);
1985 mail.raw = [b"DKIM-Signature: v=1; a=rsa-sha256; d=example.org; s=s1; bh=x; b=y\r\n".as_slice(), &mail.raw].concat();
1987 let mut record = synthetic::record_of(&mail, 1, 1);
1988 let stamp = "Authentication-Results: sioul-0123456789ab.invalid;\r\n\tdkim=pass header.d=example.org header.s=s1 header.b=AbCdEf;\r\n\tspf=fail (sioul-0123456789ab.invalid: domain of example.org does not designate 192.0.2.7 as permitted sender) smtp.mailfrom=example.org;\r\n\tdmarc=pass header.from=example.org policy.dmarc=reject\r\n";
1989 record.checked = Some(corpus::Checked { at: record.date + 3600, results: stamp.into(), whole: true, gone: false });
1990 let ids = vec!["sioul-0123456789ab.invalid".to_string(), "mx.example.net".to_string()];
1991 let trusted: BTreeMap<String, Vec<String>> = [(record.account.clone(), ids.clone())].into();
1992 let lexicon = Lexicon::builtin();
1993 let (words, trained) = read_message(&record, &trusted, &lexicon);
1994 let stored = Card::from_bytes(&[stamp.as_bytes(), &mail.raw].concat()).unwrap();
1995 let (sorted_words, sorted) = features_of(&stored, &ids, Some(record.date), None, &lexicon);
1996 assert_eq!(words, sorted_words, "the same words");
1997 assert_eq!(trained.map(f32::to_bits), sorted.map(f32::to_bits), "the same header features");
1998 let value = |x: &[f32; N], name: &str| x[spam::NAMES.iter().position(|n| *n == name).unwrap()];
1999 assert_eq!((value(&trained, "spf_fail"), value(&trained, "dkim_pass"), value(&trained, "dmarc_pass")), (1.0, 1.0, 1.0));
2000 let unchecked = read_message(&Record { checked: None, ..record.clone() }, &trusted, &lexicon).1;
2002 assert!(value(&unchecked, "spf_fail").is_nan() && value(&unchecked, "dkim_pass").is_nan());
2003 record.checked.as_mut().unwrap().at = record.date + 365 * 86_400;
2005 let (_, old) = read_message(&record, &trusted, &lexicon);
2006 assert_eq!((value(&old, "spf_fail"), value(&old, "spf_pass"), value(&old, "dkim_pass"), value(&old, "dmarc_pass")), (0.0, 0.0, 1.0, 1.0));
2007 let unsigned = Record { checked: record.checked.clone(), ..synthetic::record_of(&synthetic::mailbox(41, 3, 3)[0], 1, 1) };
2009 assert_eq!(value(&read_message(&unsigned, &trusted, &lexicon).1, "spf_fail"), 1.0);
2010 record.checked = Some(corpus::Checked { at: record.date + 3600, gone: true, ..corpus::Checked::default() });
2012 assert!(value(&read_message(&record, &trusted, &lexicon).1, "dkim_pass").is_nan());
2013 }
2014
2015 #[test]
2021 fn the_filters_word_is_learned_until_you_say() {
2022 use sioul_core::spam::Class;
2023 use sioul_core::spam::labels::{self as logs, Entry, Moved, Source};
2024 let root = scratch("filters-word");
2025 let dirs = Dirs::under(&root);
2026 let records = records(&synthetic::mailbox(13, 200, 120));
2027 corpus::store(&dirs, &records).unwrap();
2028 let id = |r: &Record| labels::message_id(&r.header_bytes()).expect("invented mail has its Message-ID");
2029 let judged = |r: &Record, class: Class| Moved { at: r.date + 60, account: r.account.clone(), folder: r.folder.clone(), uidvalidity: r.uidvalidity, uid: r.uid, message_id: Some(id(r)), class };
2030 let inbox: Vec<&Record> = records.iter().filter(|r| r.role == sioul_core::folders::Role::Inbox && !r.flags.iter().any(|f| f == "$Junk")).take(2).collect();
2032 let junk: Vec<&Record> = records.iter().filter(|r| r.role == sioul_core::folders::Role::Junk && r.flags.is_empty()).take(2).collect();
2033 let state = &dirs.state;
2034 logs::append_to(&logs::own_flagged_log(state, "desk"), &judged(inbox[0], Class::Spam)).unwrap();
2035 logs::append_to(&logs::own_flagged_log(state, "phone"), &judged(inbox[1], Class::Unsure)).unwrap();
2036 for r in &junk {
2037 logs::append_to(&logs::own_moved_log(state, "phone"), &Moved { folder: "INBOX".into(), uid: r.uid + 10_000, ..judged(r, Class::Spam) }).unwrap();
2038 }
2039 let said = Entry { at: junk[1].date + 3600, account: junk[1].account.clone(), folder: junk[1].folder.clone(), uidvalidity: 1, uid: junk[1].uid, message_id: Some(id(junk[1])), label: Label::Ham, source: Source::NotSpam };
2040 logs::append_to(&logs::own_log(state, "desk"), &said).unwrap();
2041 let summary = train(&dirs, &trusted(), &small(), &mut |_| {}, &Cancel::new()).unwrap();
2042 let l = &summary.labels;
2043 assert_eq!((l.ham, l.spam), (199, 120), "{l:?}");
2045 assert_eq!((l.by_filter, l.unsure, l.filter_moved, l.by_log), (1, 1, 1, 1), "{l:?}");
2046 assert_eq!(summary.split.train_ham + summary.split.train_spam + summary.split.test_ham + summary.split.test_spam, 319, "the doubt left out");
2047 let _ = std::fs::remove_dir_all(root);
2048 }
2049
2050 #[test]
2051 fn the_table_is_replaced_only_when_no_worse() {
2052 assert_eq!(replacement(0, 0, 0), (true, "nothing-new"));
2053 assert_eq!(replacement(3, 3, 40), (true, "no-worse"));
2054 assert_eq!(replacement(2, 3, 40), (true, "no-worse"));
2055 assert_eq!(replacement(4, 3, 40), (false, "worse"));
2056 }
2057
2058 #[test]
2059 fn standardization_of_columns() {
2060 let rows = [vec![1.0, 5.0], vec![3.0, 5.0]];
2061 let (mean, deviation) = standardization(|k| rows[k].clone(), |_| true, rows.len(), 2, 2);
2062 assert_eq!(mean, vec![2.0, 5.0]);
2063 assert_eq!(deviation, vec![1.0, 0.0], "a column that never changes has no deviation");
2064 let rows = [vec![1.0, 4.0], vec![3.0, 6.0], vec![5.0, 0.0]];
2066 let (mean, deviation) = standardization(|k| rows[k].clone(), |k| k < 2, rows.len(), 2, 1);
2067 assert_eq!((mean, deviation), (vec![3.0, 5.0], vec![(8.0f64 / 3.0).sqrt(), 1.0]));
2068 let rows = [vec![1.0, 4.0], vec![3.0, f64::NAN], vec![5.0, 6.0]];
2070 let (mean, deviation) = standardization(|k| rows[k].clone(), |_| true, rows.len(), 2, 1);
2071 assert_eq!((mean[1], deviation[1]), (5.0, 1.0));
2072 }
2073}