WIP: Use Finite State Transducers (FST) as the backing store for language models

FST have the nice properties of both compressing the ngram data by exploiting
common suffixes and prefixes as well as being embeddable into the binary in a
form that is directly suitable for look-up thereby avoiding the separate
decompression step and indireclty using memory mappings as supplied by the
operating system for all binaries.

This is still WIP since I do not know how to regenerate the language models and
it also seems like the unique ngram models are built elsewhere.

TODO:
* Integrate unique ngram models
* Regenerate all language models
* Drop non-unified models

Author: adamreichold

Cargo.lock

@@ -72,6 +72,7 @@ brotli = "7.0.0"
 compact_str = "0.9.0"
 dashmap = "6.1.0"
 fraction = "0.15.3"
+fst = "0.4.7"
 include_dir = "0.7.4"
 itertools = "0.14.0"
 maplit = "1.0.2"

Cargo.toml

@@ -18,4 +18,6 @@ use include_dir::{include_dir, Dir};
 
 pub const GERMAN_MODELS_DIRECTORY: Dir = include_dir!("$CARGO_MANIFEST_DIR/models");
 
+pub const GERMAN_UNIFIED_MODEL: &[u8] = include_bytes!("../models/ngrams.fst");
+
 pub const GERMAN_TESTDATA_DIRECTORY: Dir = include_dir!("$CARGO_MANIFEST_DIR/testdata");

language-models/de/models/ngrams.fst

@@ -127,7 +127,7 @@ fn deserialize_ngram_probabilities<'de, D: Deserializer<'de>>(
 
 pub(crate) struct TrainingDataLanguageModel {
     pub(crate) absolute_frequencies: HashMap<Ngram, u32>,
-    ngram_probability_model: NgramProbabilityModel,
+    pub(crate) ngram_probability_model: NgramProbabilityModel,
 }
 
 impl TrainingDataLanguageModel {
@@ -260,6 +260,72 @@ fn get_utf8_slice(string: &str, start: usize, end: usize) -> &str {
         .unwrap()
 }
 
+pub(crate) struct UnifiedNgramModel<'a> {
+    map: fst::Map<&'a [u8]>,
+}
+
+impl<'a> UnifiedNgramModel<'a> {
+    pub(crate) const PROBABILITY: u8 = 0;
+    pub(crate) const UNIQUE: u8 = 1;
+    pub(crate) const MOST_COMMON: u8 = 2;
+
+    pub(crate) fn load(language: Language) -> Result<Self, fst::Error> {
+        let data = match language {
+            #[cfg(feature = "german")]
+            Language::German => lingua_german_language_model::GERMAN_UNIFIED_MODEL,
+
+            _ => unimplemented!(),
+        };
+
+        fst::Map::new(data).map(|map| Self { map })
+    }
+
+    pub(crate) fn get_probability(&self, ngram: &str) -> Option<f64> {
+        self.get(ngram, Self::PROBABILITY).map(f64::from_bits)
+    }
+
+    pub(crate) fn is_unique(&self, ngram: &str) -> bool {
+        self.get(ngram, Self::UNIQUE).is_some()
+    }
+
+    pub(crate) fn is_most_common(&self, ngram: &str) -> bool {
+        self.get(ngram, Self::MOST_COMMON).is_some()
+    }
+
+    fn get(&self, ngram: &str, kind: u8) -> Option<u64> {
+        let key = UnifiedNgramKey::new(ngram, kind);
+
+        self.map.get(key)
+    }
+}
+
+#[derive(Clone, Copy)]
+pub(crate) struct UnifiedNgramKey {
+    len: usize,
+    key: [u8; Self::MAX_LEN],
+}
+
+impl UnifiedNgramKey {
+    // Maximum UTF-8-encoded length of fivegrams plus one kind byte.
+    pub(crate) const MAX_LEN: usize = 5 * 4 + 1;
+
+    pub(crate) fn new(ngram: &str, kind: u8) -> Self {
+        let len = ngram.len();
+
+        let mut key = [0; Self::MAX_LEN];
+        key[..len].copy_from_slice(ngram.as_bytes());
+        key[len] = kind;
+
+        Self { len, key }
+    }
+}
+
+impl AsRef<[u8]> for UnifiedNgramKey {
+    fn as_ref(&self) -> &[u8] {
+        &self.key[..=self.len]
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use itertools::Itertools;

language-models/de/src/lib.rs

@@ -21,11 +21,12 @@ use std::io::{BufRead, BufReader, LineWriter, Write};
 use std::path::Path;
 
 use brotli::CompressorWriter;
+use fraction::ToPrimitive;
 use itertools::Itertools;
 use regex::Regex;
 
 use crate::constant::{MULTIPLE_WHITESPACE, NUMBERS, PUNCTUATION};
-use crate::model::TrainingDataLanguageModel;
+use crate::model::{TrainingDataLanguageModel, UnifiedNgramKey, UnifiedNgramModel};
 use crate::ngram::Ngram;
 use crate::Language;
 
@@ -127,6 +128,18 @@ impl LanguageModelFilesWriter {
             "fivegrams.json",
         )?;
 
+        Self::write_unified_language_model(
+            [
+                &unigram_model,
+                &bigram_model,
+                &trigram_model,
+                &quadrigram_model,
+                &fivegram_model,
+            ],
+            output_directory_path,
+            "ngrams.fst",
+        )?;
+
         Ok(())
     }
 
@@ -167,6 +180,53 @@ impl LanguageModelFilesWriter {
         compressed_file.write_all(model.to_json().as_bytes())?;
         Ok(())
     }
+
+    fn write_unified_language_model(
+        models: [&TrainingDataLanguageModel; 5],
+        output_directory_path: &Path,
+        file_name: &str,
+    ) -> io::Result<()> {
+        let mut pairs = models
+            .iter()
+            .flat_map(|model| {
+                let probabilities =
+                    model
+                        .ngram_probability_model
+                        .ngrams
+                        .iter()
+                        .map(|(ngram, probability)| {
+                            let key = UnifiedNgramKey::new(ngram, UnifiedNgramModel::PROBABILITY);
+                            let value = probability.to_f64().unwrap().ln().to_bits();
+
+                            (key, value)
+                        });
+
+                let most_common = model
+                    .absolute_frequencies
+                    .iter()
+                    .k_largest_by_key(25, |(_, frequency)| *frequency)
+                    .map(|(ngram, frequency)| {
+                        let key =
+                            UnifiedNgramKey::new(&ngram.value, UnifiedNgramModel::MOST_COMMON);
+                        let value = *frequency as u64;
+
+                        (key, value)
+                    });
+
+                probabilities.chain(most_common)
+            })
+            .collect_vec();
+
+        pairs.sort_unstable_by(|(lhs, _), (rhs, _)| lhs.as_ref().cmp(rhs.as_ref()));
+
+        let mut builder = fst::MapBuilder::memory();
+        builder.extend_iter(pairs).unwrap();
+        let buffer = builder.into_inner().unwrap();
+
+        let file_path = output_directory_path.join(file_name);
+        let mut file = File::create(file_path)?;
+        file.write_all(&buffer)
+    }
 }
 
 impl TestDataFilesWriter {

src/detector.rs

@@ -62,12 +62,13 @@ def test_language_model_files_writer(language_model_files_text):
 
     files = read_directory_content(output_directory_path)
 
-    assert len(files) == 5
-    assert files[4] == "unigrams.json.br"
+    assert len(files) == 6
+    assert files[5] == "unigrams.json.br"
     assert files[0] == "bigrams.json.br"
-    assert files[3] == "trigrams.json.br"
-    assert files[2] == "quadrigrams.json.br"
+    assert files[4] == "trigrams.json.br"
+    assert files[3] == "quadrigrams.json.br"
     assert files[1] == "fivegrams.json.br"
+    assert files[2] == "ngrams.fst"
 
 
 def test_test_data_files_writer(test_data_files_text):