wip

Author: luciaquirke

Cargo.lock

@@ -13,6 +13,7 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 anyhow = "1.0.81"
 bincode = "1.3.3"
+derive_builder = "0.20.0"
 funty = "2.0.0"
 indicatif = "0.17.8"
 memmap2 = "0.9.4"

Cargo.toml

@@ -0,0 +1,50 @@
+``` rust
+struct SuffixTable {}
+impl SuffixTable {
+    pub fn count_next(&self);
+
+    /// To be deleted - not a natural fit for a suffix table data structure, and depends on count_next
+    /// which is "overridden" in ShardedMemmapIndex. I want to move it into a sampler struct.
+    pub fn sample(&self) {
+        counts = self.count_next()
+        do_sample(counts)
+    }
+}
+
+struct MemmapIndex {
+    table: SuffixTable
+}
+impl MemmapIndex {
+    ...lots of index/table creation code
+
+    pub fn count_next(&self) {
+        self.table.count_next()
+    }
+
+    pub fn sample(&self) {
+        self.table.sample()
+    }
+}
+
+struct ShardedInMemoryIndex {
+    shards: Vec<MemmapIndex>
+}
+impl ShardedMemmapIndex {
+    pub fn count_next(&self) {
+        counts = self.shards.iter().map(|shard| shard.count_next())
+        merge_counts(counts)
+    }
+}
+
+// WIP
+struct Sampler {
+    // Not possible due to PyO3 not supporting lifetimes
+    index: Box<Index>
+}
+impl Sampler {
+    pub fn sample(&self) {
+        counts = self.index.count_next()
+        do_sample(counts)
+    }
+}
+```

src/architecture.md

@@ -0,0 +1,10 @@
+use std::collections::HashMap;
+
+pub trait CountableIndex: Send + Sync {
+    fn count_next(&self, query: Vec<u16>, vocab: Option<u16>) -> Vec<usize>;
+    
+    /// Generate a frequency map from an occurrence frequency 
+    /// to the number of n-grams in the data structure with that 
+    /// frequency.
+    fn count_ngrams(&self, n: usize) -> HashMap<usize, usize>;
+}

src/countable_index.rs

@@ -1,14 +1,16 @@
 use bincode::{deserialize, serialize};
-use pyo3::exceptions::PyValueError;
+use std::collections::HashMap;
 use pyo3::prelude::*;
 use std::fs::File;
 use std::io::Read;
 
 use crate::table::SuffixTable;
+use crate::countable_index::CountableIndex;
 use crate::util::transmute_slice;
 
 /// An in-memory index exposes suffix table functionality over text corpora small enough to fit in memory.
 #[pyclass]
+#[derive(Builder)]
 pub struct InMemoryIndex {
     table: SuffixTable,
 }
@@ -59,14 +61,14 @@ impl InMemoryIndex {
         self.table.contains(&query)
     }
 
-    pub fn count(&self, query: Vec<u16>) -> usize {
-        self.table.positions(&query).len()
-    }
-
     pub fn positions(&self, query: Vec<u16>) -> Vec<u64> {
         self.table.positions(&query).to_vec()
     }
 
+    pub fn count(&self, query: Vec<u16>) -> usize {
+        self.table.positions(&query).len()
+    }
+
     pub fn count_next(&self, query: Vec<u16>, vocab: Option<u16>) -> Vec<usize> {
         self.table.count_next(&query, vocab)
     }
@@ -75,48 +77,20 @@ impl InMemoryIndex {
         self.table.batch_count_next(&queries, vocab)
     }
 
-    pub fn sample_unsmoothed(
-        &self,
-        query: Vec<u16>,
-        n: usize,
-        k: usize,
-        num_samples: usize,
-        vocab: Option<u16>
-    ) -> Result<Vec<Vec<u16>>, PyErr> {
-        self.table
-            .sample_unsmoothed(&query, n, k, num_samples, vocab)
-            .map_err(|error| PyValueError::new_err(error.to_string()))
-    }
-
-    pub fn sample_smoothed(
-        &mut self,
-        query: Vec<u16>,
-        n: usize,
-        k: usize,
-        num_samples: usize,
-        vocab: Option<u16>
-    ) -> Result<Vec<Vec<u16>>, PyErr> {
-        self.table
-            .sample_smoothed(&query, n, k, num_samples, vocab)
-            .map_err(|error| PyValueError::new_err(error.to_string()))
-    }
-
-    pub fn smoothed_probs(&mut self, query: Vec<u16>, vocab: Option<u16>) -> Vec<f64> {
-        self.table.get_smoothed_probs(&query, vocab)
-    }
-
-    pub fn batch_smoothed_probs(&mut self, queries: Vec<Vec<u16>>, vocab: Option<u16>) -> Vec<Vec<f64>> {
-        self.table.batch_get_smoothed_probs(&queries, vocab)
-    }
-
-    pub fn estimate_deltas(&mut self, n: usize) {
-        self.table.estimate_deltas(n);
-    }
-
     pub fn save(&self, path: String) -> PyResult<()> {
         // TODO: handle errors here
         let bytes = serialize(&self.table).unwrap();
         std::fs::write(&path, bytes)?;
         Ok(())
     }
 }
+
+impl CountableIndex for InMemoryIndex {
+    fn count_next(&self, query: Vec<u16>, vocab: Option<u16>) -> Vec<usize> {
+        self.table.count_next(&query, vocab)
+    }
+
+    fn count_ngrams(&self, n: usize) -> HashMap<usize, usize> {
+        self.table.count_ngrams(n)
+    }
+}

src/in_memory_index.rs

@@ -1,23 +1,28 @@
 pub mod mmap_slice;
 pub use in_memory_index::InMemoryIndex;
 pub use memmap_index::MemmapIndex;
-pub use sharded_index::ShardedIndex;
+pub use sharded_memmap_index::ShardedMemmapIndex;
 pub use table::SuffixTable;
 
 /// Python bindings
 use pyo3::prelude::*;
 
+mod sharded_memmap_index;
 mod in_memory_index;
 mod memmap_index;
-mod sharded_index;
-mod par_quicksort;
+mod countable_index;
+mod sampler;
 mod table;
+mod par_quicksort;
 mod util;
 
+#[macro_use]
+extern crate derive_builder;
+
 #[pymodule]
 fn tokengrams(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<InMemoryIndex>()?;
     m.add_class::<MemmapIndex>()?;
-    m.add_class::<ShardedIndex>()?;
+    m.add_class::<ShardedMemmapIndex>()?;
     Ok(())
 }