sais python functions

Author: norabelrose

src/in_memory_index.rs

@@ -50,6 +50,34 @@ impl InMemoryIndex {
         })
     }
 
+    #[staticmethod]
+    pub fn from_token_file_sais(
+        path: String,
+        token_limit: Option<usize>,
+    ) -> PyResult<Self> {
+        let mut buffer = Vec::new();
+        let mut file = File::open(&path)?;
+
+        if let Some(max_tokens) = token_limit {
+            // Limit on the number of tokens to consider is provided
+            let max_bytes = max_tokens * std::mem::size_of::<u16>();
+            file.take(max_bytes as u64).read_to_end(&mut buffer)?;
+        } else {
+            file.read_to_end(&mut buffer)?;
+        };
+
+        Ok(InMemoryIndex {
+            table: SuffixTable::new_sais(transmute_slice(buffer.as_slice())),
+        })
+    }
+
+    #[staticmethod]
+    pub fn sais(tokens: Vec<u16>) -> Self {
+        InMemoryIndex {
+            table: SuffixTable::new_sais(tokens),
+        }
+    }
+
     pub fn contains(&self, query: Vec<u16>) -> bool {
         self.table.contains(&query)
     }

src/sais.rs

@@ -32,6 +32,7 @@ extern crate utf16_literal;
 use rayon::prelude::*;
 use std::u64;
 
+use crate::util::par_bincount;
 use self::SuffixType::{Ascending, Descending, Valley};
 
 
@@ -55,6 +56,7 @@ fn sais<T: Text + ?Sized>(sa: &mut [u64], stypes: &mut SuffixTypes, bins: &mut B
     }
     sa.fill(0);
 
+    // TODO: Parallelize this step across batches / chunks / documents in the corpus
     stypes.compute(text);
     bins.find_sizes((0..text.len()).map(|i| text.char_at(i)));
     bins.find_tail_pointers();

src/util.rs

@@ -1,3 +1,32 @@
+use funty::Unsigned;
+use rayon::prelude::*;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+/// Essentially np.bincount(data) in parallel.
+pub fn par_bincount<'a, I, T>(data: &'a I) -> Vec<usize>
+where
+    I: IntoParallelRefIterator<'a, Item = T>,
+    T: Unsigned,
+{
+    // Find the maximum value in the data
+    let max = match data.par_iter().max() {
+        Some(m) => m,
+        None => return Vec::new(),
+    };
+
+    // Create a vector of atomic counters
+    let mut counts = Vec::with_capacity(max.as_usize() + 1);
+    for _ in 0..=max.as_usize() {
+        counts.push(AtomicUsize::new(0));
+    }
+
+    // Increment the counters in parallel
+    data.par_iter().for_each(|x| {
+        counts[x.as_usize()].fetch_add(1, Ordering::Relaxed);
+    });
+    counts.into_iter().map(|c| c.into_inner()).collect()
+}
+
 /// Return a zero-copy view of the given slice with the given type.
 /// The resulting view has the same lifetime as the provided slice.
 #[inline]

tests/tests.rs

@@ -15,7 +15,6 @@ fn qc<T: Testable>(f: T) {
 
 
 // Do some testing on substring search.
-
 #[test]
 fn empty_find_empty() {
     let sa = sais("");

tokengrams/tokengrams.pyi

@@ -7,6 +7,10 @@ class InMemoryIndex:
     @staticmethod
     def from_token_file(path: str, verbose: bool, token_limit: int | None) -> "InMemoryIndex":
         """Construct a `InMemoryIndex` from a file containing raw little-endian tokens."""
+    
+    @staticmethod
+    def from_token_file_sais(path: str, verbose: bool, token_limit: int | None) -> "InMemoryIndex":
+        """Construct a `InMemoryIndex` from a file containing raw little-endian tokens."""
 
     def contains(self, query: list[int]) -> bool:
         """Check if `query` has nonzero count. Faster than `count(query) > 0`."""
@@ -46,7 +50,7 @@ class MemmapIndex:
     @staticmethod
     def build(token_file: str, index_file: str) -> "MemmapIndex":
         """Build a memory-mapped index from a token file."""
-    
+
     def contains(self, query: list[int]) -> bool:
         """Check if `query` has nonzero count. Faster than `count(query) > 0`."""