Web spell as dedicated module (#240)

mikkeldenker · web-flow · commit 040d04413e0f · 2024-11-29T15:15:18.000+01:00
* separate web-spell into a dedicated module

* web-spell readme
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -16,6 +16,7 @@ members = [
   "crates/leechy",
   "crates/common",
   "crates/simple-wal",
+  "crates/web-spell",
   "fuzz",
 ]
 resolver = "2"
@@ -187,6 +188,7 @@ xxhash-rust = { version = "0.8.10", features = ["xxh3", "const_xxh3"] }
 zipf = "7.0.0"
 zstd = { version = "0.13", features = ["experimental"] }
 urlencoding = "2.1.3"
+file_store = { path = "crates/file-store" }
 
 [profile.test.package]
 flate2.opt-level = 3
diff --git a/assets/licenses.html b/assets/licenses.html
@@ -46,7 +46,7 @@ <h2>Overview of licenses:</h2>
         <ul class="licenses-overview">
             <li><a href="#Apache-2.0">Apache License 2.0</a> (411)</li>
             <li><a href="#MIT">MIT License</a> (191)</li>
-            <li><a href="#AGPL-3.0">GNU Affero General Public License v3.0</a> (9)</li>
+            <li><a href="#AGPL-3.0">GNU Affero General Public License v3.0</a> (10)</li>
             <li><a href="#BSD-3-Clause">BSD 3-Clause &quot;New&quot; or &quot;Revised&quot; License</a> (9)</li>
             <li><a href="#MPL-2.0">Mozilla Public License 2.0</a> (8)</li>
             <li><a href="#Unicode-3.0">Unicode License v3</a> (4)</li>
@@ -76,6 +76,7 @@ <h4>Used by:</h4>
                     <li><a href=" https://crates.io/crates/robotstxt ">robotstxt 0.1.0</a></li>
                     <li><a href=" https://crates.io/crates/simple_wal ">simple_wal 0.1.0</a></li>
                     <li><a href=" https://crates.io/crates/speedy_kv ">speedy_kv 0.1.0</a></li>
+                    <li><a href=" https://crates.io/crates/web-spell ">web-spell 0.1.0</a></li>
                 </ul>
                 <pre class="license-text">GNU AFFERO GENERAL PUBLIC LICENSE
 Version 3, 19 November 2007
diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml
@@ -117,6 +117,7 @@ uuid.workspace = true
 whatlang.workspace = true
 zimba = { path = "../zimba" }
 urlencoding.workspace = true
+web-spell = { path = "../web-spell" }
 
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 tikv-jemallocator.workspace = true
diff --git a/crates/core/src/config/defaults.rs b/crates/core/src/config/defaults.rs
@@ -156,22 +156,6 @@ impl SearchQuery {
     }
 }
 
-pub struct Correction;
-
-impl Correction {
-    pub fn misspelled_prob() -> f64 {
-        0.1
-    }
-
-    pub fn correction_threshold() -> f64 {
-        50.0 // logprob difference
-    }
-
-    pub fn lm_prob_weight() -> f64 {
-        5.77
-    }
-}
-
 pub struct Widgets;
 
 impl Widgets {
diff --git a/crates/core/src/config/mod.rs b/crates/core/src/config/mod.rs
@@ -16,6 +16,8 @@
 
 pub mod defaults;
 
+pub use web_spell::CorrectionConfig;
+
 use super::Result;
 use crate::ampc::dht;
 use crate::distributed::member::ShardId;
@@ -656,33 +658,6 @@ pub struct SiteStatsConfig {
     pub skip_warc_files: Option<usize>,
 }
 
-#[derive(Clone, Copy, Debug, serde::Deserialize, serde::Serialize)]
-pub struct CorrectionConfig {
-    /// The probability that a word is misspelled
-    #[serde(default = "defaults::Correction::misspelled_prob")]
-    pub misspelled_prob: f64,
-
-    /// Lambda in eq. 2 (http://static.googleusercontent.com/media/research.google.com/en/us/pubs/archive/36180.pdf)
-    #[serde(default = "defaults::Correction::lm_prob_weight")]
-    pub lm_prob_weight: f64,
-
-    /// The threshold that the difference between the log probability of the best
-    /// correction and the observed word must be above for the word to be
-    /// corrected
-    #[serde(default = "defaults::Correction::correction_threshold")]
-    pub correction_threshold: f64,
-}
-
-impl Default for CorrectionConfig {
-    fn default() -> Self {
-        Self {
-            misspelled_prob: defaults::Correction::misspelled_prob(),
-            lm_prob_weight: defaults::Correction::lm_prob_weight(),
-            correction_threshold: defaults::Correction::correction_threshold(),
-        }
-    }
-}
-
 #[derive(Clone, Debug, serde::Deserialize, serde::Serialize)]
 pub struct GossipConfig {
     pub seed_nodes: Option<Vec<SocketAddr>>,
diff --git a/crates/core/src/entrypoint/web_spell.rs b/crates/core/src/entrypoint/web_spell.rs
@@ -26,9 +26,9 @@ use whatlang::Lang;
 use crate::{
     config::{self, WebSpellConfig},
     entrypoint::download_all_warc_files,
-    web_spell::{FirstTrainer, FirstTrainerResult, SecondTrainer},
     webpage::Html,
 };
+use web_spell::{FirstTrainer, FirstTrainerResult, SecondTrainer};
 
 pub struct SpellWorker {
     prev_trained: Vec<FnvHashMap<Lang, FirstTrainerResult>>,
diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs
@@ -92,7 +92,6 @@ pub mod tokenizer;
 #[allow(unused)]
 mod ttl_cache;
 pub mod warc;
-pub mod web_spell;
 pub mod webgraph;
 pub mod webpage;
 mod widgets;
diff --git a/crates/core/src/search_prettifier/mod.rs b/crates/core/src/search_prettifier/mod.rs
@@ -34,9 +34,9 @@ use crate::{
     ranking::{SignalEnumDiscriminants, SignalScore},
     searcher::SearchQuery,
     snippet::TextSnippet,
-    web_spell::{self, CorrectionTerm},
     webpage::url_ext::UrlExt,
 };
+use web_spell::{self, CorrectionTerm};
 
 pub use self::stack_overflow::{create_stackoverflow_sidebar, CodeOrText};
 pub use entity::DisplayedEntity;
diff --git a/crates/core/src/searcher/api/mod.rs b/crates/core/src/searcher/api/mod.rs
@@ -40,7 +40,6 @@ use crate::ranking::{
     SignalScore,
 };
 use crate::search_prettifier::{DisplayedSidebar, DisplayedWebpage, HighlightedSpellCorrection};
-use crate::web_spell::SpellChecker;
 use crate::webgraph::remote::RemoteWebgraph;
 use crate::webgraph::EdgeLimit;
 use crate::webpage::html::links::RelFlags;
@@ -51,6 +50,7 @@ use crate::{
     ranking::{models::lambdamart::LambdaMART, pipeline::RankingPipeline},
 };
 use crate::{query, webgraph, Result};
+use web_spell::SpellChecker;
 
 use self::sidebar::SidebarManager;
 use self::widget::WidgetManager;
@@ -361,33 +361,29 @@ where
             .terms
             .into_iter()
             .filter_map(|t| match t {
-                crate::web_spell::CorrectionTerm::Corrected { orig, correction } => {
+                web_spell::CorrectionTerm::Corrected { orig, correction } => {
                     Some((orig, correction))
                 }
-                crate::web_spell::CorrectionTerm::NotCorrected(_) => None,
+                web_spell::CorrectionTerm::NotCorrected(_) => None,
             })
             .collect();
 
-        let mut correction = crate::web_spell::Correction::empty(query);
+        let mut correction = web_spell::Correction::empty(query);
 
         for term in terms {
             match term {
                 query::parser::Term::SimpleOrPhrase(query::parser::SimpleOrPhrase::Simple(t)) => {
                     if let Some(term_correction) = correction_map.get(t.as_str()) {
-                        correction.push(crate::web_spell::CorrectionTerm::Corrected {
+                        correction.push(web_spell::CorrectionTerm::Corrected {
                             orig: String::from(t),
                             correction: term_correction.to_string(),
                         });
                     } else {
-                        correction.push(crate::web_spell::CorrectionTerm::NotCorrected(
-                            String::from(t),
-                        ));
+                        correction.push(web_spell::CorrectionTerm::NotCorrected(String::from(t)));
                     }
                 }
                 _ => {
-                    correction.push(crate::web_spell::CorrectionTerm::NotCorrected(
-                        term.to_string(),
-                    ));
+                    correction.push(web_spell::CorrectionTerm::NotCorrected(term.to_string()));
                 }
             }
         }
diff --git a/crates/core/src/snippet.rs b/crates/core/src/snippet.rs
@@ -22,11 +22,11 @@ use crate::query::Query;
 use crate::tokenizer::fields::{
     BigramTokenizer, DefaultTokenizer, FieldTokenizer, Stemmed, TrigramTokenizer,
 };
-use crate::web_spell::sentence_ranges;
 use crate::webpage::region::Region;
 use hashbrown::{HashMap, HashSet};
 use lending_iter::LendingIterator;
 use utoipa::ToSchema;
+use web_spell::sentence_ranges;
 
 use itertools::Itertools;
 use whatlang::Lang;
diff --git a/crates/web-spell/Cargo.toml b/crates/web-spell/Cargo.toml
@@ -0,0 +1,27 @@
+[package]
+edition = "2021"
+license = "AGPL-3.0"
+name = "web-spell"
+version = "0.1.0"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+anyhow.workspace = true
+thiserror.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+bincode.workspace = true
+fst.workspace = true
+itertools.workspace = true
+indicatif.workspace = true
+rayon.workspace = true
+uuid.workspace = true
+tracing.workspace = true
+memmap2.workspace = true
+whatlang.workspace = true
+fnv.workspace = true
+
+[dev-dependencies]
+file_store.workspace = true
+proptest.workspace = true
diff --git a/crates/web-spell/README.md b/crates/web-spell/README.md
@@ -0,0 +1,13 @@
+# Web Spell
+
+Automatic spelling correction from web data. It is based on the paper
+[Using the Web for Language Independent Spellchecking and
+Autocorrection](http://static.googleusercontent.com/media/research.google.com/en/us/pubs/archive/36180.pdf)
+from google.
+
+## Usage
+```rust
+let checker = SpellChecker::open("<path-to-model>", CorrectionConfig::default()).unwrap();
+let correction = checker.correct("hwllo", Lang::Eng);
+assert_eq!(correction.unwrap().terms, vec![CorrectionTerm::Corrected { orig: "hwllo".to_string(), correction: "hello".to_string() }]);
+```
diff --git a/crates/web-spell/src/config.rs b/crates/web-spell/src/config.rs
@@ -0,0 +1,58 @@
+// Stract is an open source web search engine.
+// Copyright (C) 2024 Stract ApS
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+fn misspelled_prob() -> f64 {
+    0.1
+}
+
+fn correction_threshold() -> f64 {
+    50.0 // logprob difference
+}
+
+fn lm_prob_weight() -> f64 {
+    5.77
+}
+
+#[derive(Clone, Copy, Debug, serde::Deserialize, serde::Serialize)]
+pub struct CorrectionConfig {
+    /// The probability that a word is misspelled
+    #[serde(default = "misspelled_prob")]
+    pub misspelled_prob: f64,
+
+    /// Lambda in eq. 2 (http://static.googleusercontent.com/media/research.google.com/en/us/pubs/archive/36180.pdf)
+    #[serde(default = "lm_prob_weight")]
+    pub lm_prob_weight: f64,
+
+    /// The threshold that the difference between the log probability of the best
+    /// correction and the observed word must be above for the word to be
+    /// corrected
+    #[serde(default = "correction_threshold")]
+    pub correction_threshold: f64,
+}
+
+impl Default for CorrectionConfig {
+    fn default() -> Self {
+        Self {
+            misspelled_prob: misspelled_prob(),
+            lm_prob_weight: lm_prob_weight(),
+            correction_threshold: correction_threshold(),
+        }
+    }
+}
+
+pub fn bincode_config() -> bincode::config::Configuration {
+    bincode::config::standard()
+}
diff --git a/crates/web-spell/src/error_model.rs b/crates/web-spell/src/error_model.rs
@@ -1,5 +1,5 @@
 // Stract is an open source web search engine.
-// Copyright (C) 2023 Stract ApS
+// Copyright (C) 2024 Stract ApS
 //
 // This program is free software: you can redistribute it and/or modify
 // it under the terms of the GNU Affero General Public License as
diff --git a/crates/web-spell/src/lib.rs b/crates/web-spell/src/lib.rs
diff --git a/crates/web-spell/src/spell_checker.rs b/crates/web-spell/src/spell_checker.rs
diff --git a/crates/web-spell/src/stupid_backoff.rs b/crates/web-spell/src/stupid_backoff.rs
diff --git a/crates/web-spell/src/term_freqs.rs b/crates/web-spell/src/term_freqs.rs
diff --git a/crates/web-spell/src/trainer.rs b/crates/web-spell/src/trainer.rs
