fix: proposed fix for the word ngrams bug

Author: PyDataBlog

src/extractors/word_ngrams.rs

@@ -1,5 +1,6 @@
 use crate::FeatureExtractor;
 use lasso::{Rodeo, Spur};
+use std::fmt::Write;
 
 #[derive(Clone)]
 pub struct WordNgrams {
@@ -30,25 +31,51 @@ impl FeatureExtractor for WordNgrams {
             return vec![];
         }
 
-        let tokens = text.split(&self.splitter).filter(|s| !s.is_empty());
+        let tokens: Vec<&str> = text.split(&self.splitter).collect();
+        let mut padded_tokens = Vec::with_capacity(tokens.len() + 2);
+        padded_tokens.push(self.padder.as_str());
+        padded_tokens.extend(tokens.iter().copied());
+        padded_tokens.push(self.padder.as_str());
 
-        // an iterator that includes padding
-        let padded_tokens_iter = std::iter::once(self.padder.as_str())
-            .chain(tokens)
-            .chain(std::iter::once(self.padder.as_str()));
+        if padded_tokens.len() < self.n {
+            return vec![];
+        }
 
-        // Use a buffer to collect tokens for each n-gram
-        let mut buffer: Vec<&str> = Vec::with_capacity(self.n);
-        let mut ngrams = Vec::new();
+        let mut ngrams = Vec::with_capacity(padded_tokens.len().saturating_sub(self.n) + 1);
 
-        for token in padded_tokens_iter {
-            buffer.push(token);
-            if buffer.len() == self.n {
-                ngrams.push(buffer.join(" "));
-                buffer.remove(0);
-            }
+        for window in padded_tokens.windows(self.n) {
+            ngrams.push(encode_word_window(window));
         }
 
         super::append_feature_counts(interner, ngrams)
     }
 }
+
+fn encode_word_window(window: &[&str]) -> String {
+    let estimated_capacity: usize = window
+        .iter()
+        .map(|token| token.len() + digits(token.len()) + 1)
+        .sum::<usize>()
+        .saturating_sub(1);
+
+    let mut encoded = String::with_capacity(estimated_capacity);
+
+    for (index, token) in window.iter().enumerate() {
+        if index > 0 {
+            encoded.push('|');
+        }
+        let _ = write!(&mut encoded, "{}:", token.len());
+        encoded.push_str(token);
+    }
+
+    encoded
+}
+
+fn digits(mut value: usize) -> usize {
+    let mut digits = 1;
+    while value >= 10 {
+        value /= 10;
+        digits += 1;
+    }
+    digits
+}

tests/test_features.rs

@@ -1,4 +1,4 @@
-use lasso::Rodeo;
+use lasso::{Rodeo, Spur};
 use simstring_rust::{CharacterNgrams, FeatureExtractor, WordNgrams};
 
 #[test]
@@ -92,6 +92,21 @@ fn test_character_trigrams_prepress() {
     assert_eq!(resolved_features, expected);
 }
 
+fn sorted_features(interner: &Rodeo, features: &[Spur]) -> Vec<String> {
+    let mut resolved: Vec<String> = features
+        .iter()
+        .map(|spur| interner.resolve(spur).to_string())
+        .collect();
+    resolved.sort();
+    resolved
+}
+
+fn sorted_strings(items: Vec<&str>) -> Vec<String> {
+    let mut owned: Vec<String> = items.into_iter().map(String::from).collect();
+    owned.sort();
+    owned
+}
+
 #[cfg(test)]
 mod word_ngrams_tests {
     use super::*;
@@ -101,12 +116,8 @@ mod word_ngrams_tests {
         let mut interner = Rodeo::default();
         let extractor = WordNgrams::default(); // n=2, splitter=" ", padder=" "
         let features = extractor.features("a b", &mut interner);
-        let resolved: Vec<String> = features
-            .iter()
-            .map(|s| interner.resolve(s).to_string())
-            .collect();
-        // Corrected expectation: The padder " " and the joiner " " create two spaces.
-        let expected = vec!["  a1", "a b1", "b  1"];
+        let resolved = sorted_features(&interner, &features);
+        let expected = sorted_strings(vec!["1: |1:a1", "1:a|1:b1", "1:b|1: 1"]);
         assert_eq!(resolved, expected);
     }
 
@@ -115,11 +126,8 @@ mod word_ngrams_tests {
         let mut interner = Rodeo::default();
         let extractor = WordNgrams::new(2, " ", "-");
         let features = extractor.features("word", &mut interner);
-        let resolved: Vec<String> = features
-            .iter()
-            .map(|s| interner.resolve(s).to_string())
-            .collect();
-        let expected = vec!["- word1", "word -1"];
+        let resolved = sorted_features(&interner, &features);
+        let expected = sorted_strings(vec!["1:-|4:word1", "4:word|1:-1"]);
         assert_eq!(resolved, expected);
     }
 
@@ -128,17 +136,14 @@ mod word_ngrams_tests {
         let mut interner = Rodeo::default();
         let extractor = WordNgrams::new(3, " ", "<PAD>");
         let features = extractor.features("this is a simple test", &mut interner);
-        let resolved: Vec<String> = features
-            .iter()
-            .map(|s| interner.resolve(s).to_string())
-            .collect();
-        let expected = vec![
-            "<PAD> this is1",
-            "this is a1",
-            "is a simple1",
-            "a simple test1",
-            "simple test <PAD>1",
-        ];
+        let resolved = sorted_features(&interner, &features);
+        let expected = sorted_strings(vec![
+            "5:<PAD>|4:this|2:is1",
+            "4:this|2:is|1:a1",
+            "2:is|1:a|6:simple1",
+            "1:a|6:simple|4:test1",
+            "6:simple|4:test|5:<PAD>1",
+        ]);
         assert_eq!(resolved, expected);
     }
 
@@ -148,42 +153,44 @@ mod word_ngrams_tests {
         let extractor = WordNgrams::new(2, " ", "$");
         let s = "You are a really really really cool dude 😄🍕";
         let features = extractor.features(s, &mut interner);
-        let resolved: Vec<String> = features
-            .iter()
-            .map(|s| interner.resolve(s).to_string())
-            .collect();
-        let expected = vec![
-            "$ You1",
-            "You are1",
-            "are a1",
-            "a really1",
-            "really really1",
-            "really really2",
-            "really cool1",
-            "cool dude1",
-            "dude 😄🍕1",
-            "😄🍕 $1",
-        ];
+        let resolved = sorted_features(&interner, &features);
+        let expected = sorted_strings(vec![
+            "1:$|3:You1",
+            "3:You|3:are1",
+            "3:are|1:a1",
+            "1:a|6:really1",
+            "6:really|6:really1",
+            "6:really|6:really2",
+            "6:really|4:cool1",
+            "4:cool|4:dude1",
+            "4:dude|8:😄🍕1",
+            "8:😄🍕|1:$1",
+        ]);
         assert_eq!(resolved, expected);
     }
 
     #[test]
     fn test_word_ngram_edge_cases() {
         let mut interner = Rodeo::default();
         let extractor = WordNgrams::new(2, " ", "$");
+
         let features_empty = extractor.features("", &mut interner);
-        let resolved_empty: Vec<String> = features_empty
-            .iter()
-            .map(|s| interner.resolve(s).to_string())
-            .collect();
-        assert_eq!(resolved_empty, vec!["$ $1"]);
+        let resolved_empty = sorted_features(&interner, &features_empty);
+        assert_eq!(resolved_empty, sorted_strings(vec!["1:$|0:1", "0:|1:$1"]));
 
         let features_spaces = extractor.features("   ", &mut interner);
-        let resolved_spaces: Vec<String> = features_spaces
-            .iter()
-            .map(|s| interner.resolve(s).to_string())
-            .collect();
-        assert_eq!(resolved_spaces, vec!["$ $1"]);
+        let resolved_spaces = sorted_features(&interner, &features_spaces);
+        assert_eq!(
+            resolved_spaces,
+            sorted_strings(vec!["1:$|0:1", "0:|0:1", "0:|0:2", "0:|0:3", "0:|1:$1"])
+        );
+
+        let features_double_space = extractor.features("a  b", &mut interner);
+        let resolved_double_space = sorted_features(&interner, &features_double_space);
+        assert_eq!(
+            resolved_double_space,
+            sorted_strings(vec!["1:$|1:a1", "1:a|0:1", "0:|1:b1", "1:b|1:$1"])
+        );
     }
 
     #[test]
@@ -193,22 +200,15 @@ mod word_ngrams_tests {
         // Case 1: n=2, input="abcd", splitter=" ", padder=" "
         let extractor_case1 = WordNgrams::new(2, " ", " ");
         let features_case1 = extractor_case1.features("abcd", &mut interner);
-        let resolved1: Vec<String> = features_case1
-            .iter()
-            .map(|s| interner.resolve(s).to_string())
-            .collect();
-        // Corrected expectation
-        let expected_case1 = vec!["  abcd1", "abcd  1"];
+        let resolved1 = sorted_features(&interner, &features_case1);
+        let expected_case1 = sorted_strings(vec!["1: |4:abcd1", "4:abcd|1: 1"]);
         assert_eq!(resolved1, expected_case1, "Failed on: n=2, input='abcd'");
 
         // Case 2: n=2, input="hello world", splitter=" ", padder=" "
         let features_case2 = extractor_case1.features("hello world", &mut interner);
-        let resolved2: Vec<String> = features_case2
-            .iter()
-            .map(|s| interner.resolve(s).to_string())
-            .collect();
-        // Corrected expectation
-        let expected_case2 = vec!["  hello1", "hello world1", "world  1"];
+        let resolved2 = sorted_features(&interner, &features_case2);
+        let expected_case2 =
+            sorted_strings(vec!["1: |5:hello1", "5:hello|5:world1", "5:world|1: 1"]);
         assert_eq!(
             resolved2, expected_case2,
             "Failed on: n=2, input='hello world'"
@@ -217,12 +217,8 @@ mod word_ngrams_tests {
         // Case 3: n=3, input="hello world", splitter=" ", padder=" "
         let extractor_case3 = WordNgrams::new(3, " ", " ");
         let features_case3 = extractor_case3.features("hello world", &mut interner);
-        let resolved3: Vec<String> = features_case3
-            .iter()
-            .map(|s| interner.resolve(s).to_string())
-            .collect();
-        // Corrected expectation
-        let expected_case3 = vec!["  hello world1", "hello world  1"];
+        let resolved3 = sorted_features(&interner, &features_case3);
+        let expected_case3 = sorted_strings(vec!["1: |5:hello|5:world1", "5:hello|5:world|1: 1"]);
         assert_eq!(
             resolved3, expected_case3,
             "Failed on: n=3, input='hello world'"