Spam filter performance and accuracy improvements (part 3)

crates/nlp/src/classifier/feature.rs

@@ -15,31 +15,72 @@ pub struct Sample {
 
 pub struct Features(pub(super) HashMap<u32, f32, BuildNoHashHasher<u32>>);
 
-pub struct SampleBuilder {
+pub struct FeatureBuilder {
     pub(super) features_mask: u32,
 }
 
-impl SampleBuilder {
-    pub fn build<I>(&self, features: I, class: f32) -> Sample
-    where
-        I: IntoIterator,
-        I::Item: AsRef<[u8]>,
-    {
-        let mut features_map = HashMap::with_capacity_and_hasher(128, BuildNoHashHasher::default());
+pub trait Feature {
+    fn prefix(&self) -> u16;
+    fn value(&self) -> &[u8];
+    fn is_global_feature(&self) -> bool;
+    fn is_local_feature(&self) -> bool;
+}
 
-        for feature in features {
-            let feature = feature.as_ref();
-            let hash = xxh3_64_with_seed(feature, 0) as u32;
-            let hash_sign = xxh3_64_with_seed(feature, 1);
+impl FeatureBuilder {
+    pub fn scale<I: Feature>(&self, features: &mut HashMap<I, f32>) {
+        // Log frequency scaling
+        for x in features.values_mut() {
+            *x = x.ln_1p();
+        }
+    }
 
-            *features_map.entry(hash & self.features_mask).or_default() +=
-                if hash_sign & 1 == 0 { 1.0 } else { -1.0 };
+    pub fn build<I: Feature>(
+        &self,
+        features: &HashMap<I, f32>,
+        account_id: Option<u32>,
+    ) -> Features {
+        // Do the "hash trick"
+        let mut features_map =
+            HashMap::with_capacity_and_hasher(features.len(), BuildNoHashHasher::default());
+        let mut buf = Vec::with_capacity(2 + 4 + 63);
+        for (feature, count) in features {
+            buf.extend_from_slice(&feature.prefix().to_be_bytes());
+            buf.extend_from_slice(feature.value());
+
+            if feature.is_global_feature() {
+                let big_hash = xxh3_64_with_seed(&buf, 0);
+                let hash = big_hash as u32 & self.features_mask;
+                let sign = if big_hash & (1 << 63) == 0 { 1.0 } else { -1.0 };
+
+                *features_map.entry(hash).or_default() += sign * count;
+            }
+
+            if feature.is_local_feature()
+                && let Some(account_id) = account_id
+            {
+                buf.extend_from_slice(&account_id.to_be_bytes());
+                let big_hash = xxh3_64_with_seed(&buf, 0);
+                let hash = big_hash as u32 & self.features_mask;
+                let sign = if big_hash & (1 << 63) == 0 { 1.0 } else { -1.0 };
+
+                *features_map.entry(hash).or_default() += sign * count;
+            }
+            buf.clear();
         }
 
-        Sample {
-            features: Features(features_map),
-            class,
+        // L2 normalization
+        let sum_of_squares = features_map
+            .values()
+            .map(|&x| x as f64 * x as f64)
+            .sum::<f64>();
+        if sum_of_squares > 0.0 {
+            let norm = sum_of_squares.sqrt() as f32;
+            for x in features_map.values_mut() {
+                *x /= norm;
+            }
         }
+
+        Features(features_map)
     }
 }
 

crates/nlp/src/classifier/sgd.rs

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
  */
 
-use crate::classifier::feature::{Features, Sample, SampleBuilder};
+use crate::classifier::feature::{FeatureBuilder, Features, Sample};
 use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};
 
 #[derive(Default)]
@@ -114,8 +114,8 @@ impl SGDClassifier {
             .collect()
     }
 
-    pub fn sample_builder(&self) -> SampleBuilder {
-        SampleBuilder {
+    pub fn feature_builder(&self) -> FeatureBuilder {
+        FeatureBuilder {
             features_mask: (self.weights.len() - 1) as u32,
         }
     }
@@ -163,9 +163,13 @@ fn log1pexp(x: f32) -> f32 {
 
 #[cfg(test)]
 pub mod tests {
-    use crate::classifier::{feature::Sample, sgd::SGDClassifier};
+    use crate::classifier::{
+        feature::{Feature, Sample},
+        sgd::SGDClassifier,
+    };
     use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};
     use std::{
+        collections::HashMap,
         fs::File,
         io::{BufRead, BufReader},
         time::Instant,
@@ -262,6 +266,24 @@ pub mod tests {
         (train, test)
     }
 
+    impl Feature for String {
+        fn prefix(&self) -> u16 {
+            0
+        }
+
+        fn value(&self) -> &[u8] {
+            self.as_bytes()
+        }
+
+        fn is_global_feature(&self) -> bool {
+            true
+        }
+
+        fn is_local_feature(&self) -> bool {
+            false
+        }
+    }
+
     #[test]
     fn sgd_classifier() {
         let reader = BufReader::new(
@@ -271,7 +293,7 @@ pub mod tests {
         let mut samples = Vec::with_capacity(1024);
 
         let mut model = SGDClassifier::new(1 << 20, 1000, 0.0001, 42);
-        let builder = model.sample_builder();
+        let builder = model.feature_builder();
 
         let time = Instant::now();
 
@@ -279,14 +301,18 @@ pub mod tests {
             let line = line.unwrap();
             let (text, class) = line.trim().rsplit_once(',').unwrap();
             let text = text.trim_start_matches('"').trim_end_matches('"');
-            samples.push(
-                builder.build(
-                    text.split_whitespace(),
-                    class
-                        .parse()
-                        .unwrap_or_else(|_| panic!("Invalid class value: {line}")),
-                ),
-            );
+
+            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, None),
+                class: class
+                    .parse()
+                    .unwrap_or_else(|_| panic!("Invalid class value: {line}")),
+            });
         }
 
         println!("Loaded {} samples in {:?}", samples.len(), time.elapsed());

crates/nlp/src/tokenizers/stream.rs

@@ -15,21 +15,9 @@ use crate::{
 };
 use std::borrow::Cow;
 
-pub struct StreamTokenizer<T: Iterator<Item = StreamInputToken<I>>, I: StreamInputTokenTrait> {
-    stream: T,
+pub struct WordStemTokenizer {
     stemmer: Stemmer,
     stop_words: Option<StopwordFnc>,
-    tokens: Vec<I>,
-}
-
-pub enum StreamInputToken<T: StreamInputTokenTrait> {
-    Word(String),
-    Other(T),
-}
-
-pub trait StreamInputTokenTrait {
-    fn from_owned(word: String) -> Self;
-    fn from_borrowed(word: &str) -> Self;
 }
 
 enum Stemmer {
@@ -39,8 +27,8 @@ enum Stemmer {
     None,
 }
 
-impl<T: Iterator<Item = StreamInputToken<I>>, I: StreamInputTokenTrait> StreamTokenizer<T, I> {
-    pub fn new(text: &str, stream: T) -> Self {
+impl WordStemTokenizer {
+    pub fn new(text: &str) -> Self {
         // Detect language
         let (mut language, score) =
             LanguageDetector::detect_single(text).unwrap_or((Language::English, 1.0));
@@ -49,7 +37,6 @@ impl<T: Iterator<Item = StreamInputToken<I>>, I: StreamInputTokenTrait> StreamTo
         }
 
         Self {
-            stream,
             stemmer: match language {
                 Language::Mandarin => Stemmer::Mandarin,
                 Language::Japanese => Stemmer::Japanese,
@@ -58,66 +45,31 @@ impl<T: Iterator<Item = StreamInputToken<I>>, I: StreamInputTokenTrait> StreamTo
                     .unwrap_or(Stemmer::None),
             },
             stop_words: STOP_WORDS[language as usize],
-            tokens: vec![],
         }
     }
-}
 
-impl<T: Iterator<Item = StreamInputToken<I>>, I: StreamInputTokenTrait> Iterator
-    for StreamTokenizer<T, I>
-{
-    type Item = I;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        if let Some(prev_token) = self.tokens.pop() {
-            return Some(prev_token);
+    pub fn tokenize<'x>(&self, word: &'x str, mut cb: impl FnMut(Cow<'x, str>)) {
+        if self.stop_words.is_some_and(|sw| sw(word)) {
+            return;
         }
-
-        for token in self.stream.by_ref() {
-            return match token {
-                StreamInputToken::Word(word) => {
-                    if self.stop_words.is_some_and(|sw| sw(word.as_str())) {
-                        continue;
-                    }
-                    match &self.stemmer {
-                        Stemmer::IndoEuropean(stemmer) => match stemmer.stem(&word) {
-                            Cow::Borrowed(word) => I::from_borrowed(word),
-                            Cow::Owned(stemmed_word) => I::from_owned(stemmed_word),
-                        },
-                        Stemmer::Mandarin => {
-                            let mut result = JIEBA.cut(&word, false).into_iter();
-                            if let Some(stemmed_word) = result.next() {
-                                let stemmed_word = I::from_borrowed(stemmed_word);
-                                self.tokens = result
-                                    .rev()
-                                    .map(|word| I::from_borrowed(word))
-                                    .collect::<Vec<_>>();
-                                stemmed_word
-                            } else {
-                                // This shouldn't happen, but just in case
-                                continue;
-                            }
-                        }
-                        Stemmer::Japanese => {
-                            let mut result = japanese::tokenize(&word).into_iter();
-                            if let Some(stemmed_word) = result.next() {
-                                self.tokens =
-                                    result.rev().map(|b| I::from_owned(b)).collect::<Vec<_>>();
-                                I::from_owned(stemmed_word)
-                            } else {
-                                // This shouldn't happen, but just in case
-                                continue;
-                            }
-                        }
-                        Stemmer::None => I::from_owned(word),
-                    }
-                }
-                StreamInputToken::Other(raw) => raw,
+        match &self.stemmer {
+            Stemmer::IndoEuropean(stemmer) => {
+                cb(stemmer.stem(word));
+            }
+            Stemmer::Mandarin => {
+                for word in JIEBA.cut(word, false) {
+                    cb(Cow::from(word));
+                }
+            }
+            Stemmer::Japanese => {
+                for word in japanese::tokenize(word) {
+                    cb(Cow::from(word));
+                }
+            }
+            Stemmer::None => {
+                cb(Cow::from(word));
             }
-            .into();
         }
-
-        None
     }
 }
 
@@ -7872,89 +7824,12 @@ pub fn symbols(input: &str) -> bool {
     )
 }
 
-impl StreamInputTokenTrait for Vec<u8> {
-    fn from_owned(word: String) -> Self {
-        word.into_bytes()
-    }
-
-    fn from_borrowed(word: &str) -> Self {
-        word.as_bytes().to_vec()
-    }
-}
-
 #[cfg(test)]
 pub mod tests {
-    use super::{StreamInputToken, symbols};
     use crate::tokenizers::{
-        stream::StreamTokenizer,
+        stream::WordStemTokenizer,
         types::{TokenType, TypesTokenizer},
     };
-    use std::{borrow::Cow, net::IpAddr};
-
-    pub trait ToStreamToken {
-        fn to_stream_token(&self) -> Option<StreamInputToken<Vec<u8>>>;
-    }
-
-    impl<T: AsRef<str>, E: AsRef<str>, U: AsRef<str>, I: AsRef<str>> ToStreamToken
-        for TokenType<T, E, U, I>
-    {
-        fn to_stream_token(&self) -> Option<StreamInputToken<Vec<u8>>> {
-            match self {
-                TokenType::Alphabetic(word) => {
-                    Some(StreamInputToken::Word(word.as_ref().to_lowercase()))
-                }
-                TokenType::Url(word) => {
-                    let word = word.as_ref();
-                    word.split_once("://")
-                        .map(|(_, host)| StreamInputToken::Other(url_host_as_bytes(host)))
-                }
-                TokenType::IpAddr(word) => word.as_ref().parse::<IpAddr>().ok().map(|ip| {
-                    StreamInputToken::Other(match ip {
-                        IpAddr::V4(ip) => ip.octets().to_vec(),
-                        IpAddr::V6(ip) => ip.octets().to_vec(),
-                    })
-                }),
-                TokenType::UrlNoScheme(word) => {
-                    StreamInputToken::Other(url_host_as_bytes(word.as_ref())).into()
-                }
-                TokenType::Alphanumeric(word) | TokenType::UrlNoHost(word) => {
-                    StreamInputToken::Other(word.as_ref().to_lowercase().into_bytes()).into()
-                }
-                TokenType::Email(word) => {
-                    StreamInputToken::Other(word.as_ref().to_lowercase().into_bytes()).into()
-                }
-                TokenType::Other(ch) => {
-                    let ch = ch.to_string();
-                    if symbols(&ch) {
-                        Some(StreamInputToken::Other(ch.into_bytes()))
-                    } else {
-                        None
-                    }
-                }
-                TokenType::Integer(word) => number_to_tag(false, word.as_ref()).into(),
-                TokenType::Float(word) => number_to_tag(true, word.as_ref()).into(),
-                TokenType::Punctuation(_) | TokenType::Space => None,
-            }
-        }
-    }
-
-    fn url_host_as_bytes(host: &str) -> Vec<u8> {
-        host.split_once('/')
-            .map_or(host, |(h, _)| h.rsplit_once(':').map_or(h, |(h, _)| h))
-            .to_lowercase()
-            .into_bytes()
-    }
-
-    fn number_to_tag(is_float: bool, num: &str) -> StreamInputToken<Vec<u8>> {
-        let t = match (is_float, num.starts_with('-')) {
-            (true, true) => b'F',
-            (true, false) => b'f',
-            (false, true) => b'I',
-            (false, false) => b'i',
-        };
-
-        StreamInputToken::Other([t, num.len() as u8].to_vec())
-    }
 
     #[test]
     fn stream_tokenizer() {
@@ -8040,15 +7915,17 @@ pub mod tests {
         ];
 
         for (input, expect) in inputs.iter() {
-            let input = StreamTokenizer::new(
-                input,
-                TypesTokenizer::new(input).filter_map(|t| t.word.to_stream_token()),
-            )
-            .map(|word| String::from_utf8(word).unwrap())
-            .collect::<Vec<_>>();
-            let expect = expect.iter().copied().map(Cow::from).collect::<Vec<_>>();
+            let tokenizer = WordStemTokenizer::new(input);
+            let mut result = Vec::new();
+            for token in TypesTokenizer::new(input) {
+                if let TokenType::Alphabetic(word) = token.word {
+                    tokenizer.tokenize(word, |t| {
+                        result.push(t.into_owned());
+                    });
+                }
+            }
 
-            assert_eq!(input, expect,);
+            assert_eq!(&result, expect,);
         }
     }
 }