lda-ruby 0.4.0 → 0.5.0

data/ext/lda-ruby-rust/src/lib.rs

@@ -1,4 +1,7 @@
 use magnus::{define_module, function, Error, Module, Object};
+use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, Mutex, OnceLock};
 
 fn available() -> bool {
     true
@@ -40,6 +43,99 @@ fn normalize_in_place(weights: &mut [f64]) {
     }
 }
 
+#[derive(Clone, PartialEq)]
+struct SessionConfig {
+    topics: usize,
+    max_iter: i64,
+    convergence: f64,
+    em_max_iter: i64,
+    em_convergence: f64,
+    init_alpha: f64,
+    min_probability: f64,
+}
+
+struct CorpusSessionData {
+    document_words: Vec<Vec<usize>>,
+    document_counts: Vec<Vec<f64>>,
+    terms: usize,
+}
+
+struct CorpusSession {
+    data: Arc<CorpusSessionData>,
+    config: Option<SessionConfig>,
+}
+
+static CORPUS_SESSIONS: OnceLock<Mutex<HashMap<u64, CorpusSession>>> = OnceLock::new();
+static NEXT_CORPUS_SESSION_ID: AtomicU64 = AtomicU64::new(1);
+
+fn corpus_sessions() -> &'static Mutex<HashMap<u64, CorpusSession>> {
+    CORPUS_SESSIONS.get_or_init(|| Mutex::new(HashMap::new()))
+}
+
+fn corpus_session_count() -> i64 {
+    match corpus_sessions().lock() {
+        Ok(sessions) => sessions.len() as i64,
+        Err(_) => 0,
+    }
+}
+
+fn corpus_session_exists(session_id: i64) -> bool {
+    if session_id <= 0 {
+        return false;
+    }
+
+    let session_key = session_id as u64;
+    match corpus_sessions().lock() {
+        Ok(sessions) => sessions.contains_key(&session_key),
+        Err(_) => false,
+    }
+}
+
+fn empty_em_output() -> (Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<Vec<f64>>>) {
+    (Vec::new(), Vec::new(), Vec::new(), Vec::new())
+}
+
+fn empty_managed_session_em_output(
+) -> (
+    i64,
+    Vec<Vec<f64>>,
+    Vec<Vec<f64>>,
+    Vec<Vec<f64>>,
+    Vec<Vec<Vec<f64>>>,
+) {
+    (0, Vec::new(), Vec::new(), Vec::new(), Vec::new())
+}
+
+struct XorShift64 {
+    state: u64,
+}
+
+impl XorShift64 {
+    fn new(seed: i64) -> Self {
+        let mut state = seed as u64;
+        if state == 0 {
+            state = 0x9E37_79B9_7F4A_7C15;
+        }
+
+        Self { state }
+    }
+
+    fn next_u64(&mut self) -> u64 {
+        let mut x = self.state;
+        x ^= x >> 12;
+        x ^= x << 25;
+        x ^= x >> 27;
+        self.state = x;
+        x.wrapping_mul(0x2545_F491_4F6C_DD1D)
+    }
+
+    fn next_f64_unit(&mut self) -> f64 {
+        // Keep 53 random bits to map uniformly into [0, 1).
+        let value = self.next_u64() >> 11;
+        value as f64 / ((1_u64 << 53) as f64)
+    }
+}
+
 fn compute_topic_weights(
     beta_probabilities: &[Vec<f64>],
     gamma: &[f64],
@@ -158,7 +254,7 @@ fn normalize_topic_term_counts(
     (beta_probabilities, beta_log)
 }
 
-fn average_gamma_shift(previous_gamma: Vec<Vec<f64>>, current_gamma: Vec<Vec<f64>>) -> f64 {
+fn average_gamma_shift_internal(previous_gamma: &[Vec<f64>], current_gamma: &[Vec<f64>]) -> f64 {
     let mut sum = 0.0_f64;
     let mut count = 0_usize;
 
@@ -183,6 +279,10 @@ fn average_gamma_shift(previous_gamma: Vec<Vec<f64>>, current_gamma: Vec<Vec<f64
     }
 }
 
+fn average_gamma_shift(previous_gamma: Vec<Vec<f64>>, current_gamma: Vec<Vec<f64>>) -> f64 {
+    average_gamma_shift_internal(previous_gamma.as_slice(), current_gamma.as_slice())
+}
+
 fn topic_document_probability(
     phi_tensor: Vec<Vec<Vec<f64>>>,
     document_counts: Vec<Vec<f64>>,
@@ -222,9 +322,9 @@ fn topic_document_probability(
     output
 }
 
-fn seeded_topic_term_probabilities(
-    document_words: Vec<Vec<usize>>,
-    document_counts: Vec<Vec<f64>>,
+fn seeded_topic_term_probabilities_internal(
+    document_words: &[Vec<usize>],
+    document_counts: &[Vec<f64>],
     topics: usize,
     terms: usize,
     min_probability: f64,
@@ -264,6 +364,427 @@ fn seeded_topic_term_probabilities(
     topic_term_counts
 }
 
+fn seeded_topic_term_probabilities(
+    document_words: Vec<Vec<usize>>,
+    document_counts: Vec<Vec<f64>>,
+    topics: usize,
+    terms: usize,
+    min_probability: f64,
+) -> Vec<Vec<f64>> {
+    seeded_topic_term_probabilities_internal(
+        document_words.as_slice(),
+        document_counts.as_slice(),
+        topics,
+        terms,
+        min_probability,
+    )
+}
+
+fn random_topic_term_probabilities(
+    topics: usize,
+    terms: usize,
+    min_probability: f64,
+    random_seed: i64,
+) -> Vec<Vec<f64>> {
+    if topics == 0 || terms == 0 {
+        return Vec::new();
+    }
+
+    let floor = floor_value(min_probability);
+    let mut rng = XorShift64::new(random_seed);
+    let mut matrix = Vec::with_capacity(topics);
+
+    for _ in 0..topics {
+        let mut weights = Vec::with_capacity(terms);
+        for _ in 0..terms {
+            weights.push(rng.next_f64_unit() + floor);
+        }
+        normalize_in_place(&mut weights);
+        matrix.push(weights);
+    }
+
+    matrix
+}
+
+fn corpus_session_data(
+    document_words: &[Vec<usize>],
+    document_counts: &[Vec<f64>],
+    terms: usize,
+) -> Arc<CorpusSessionData> {
+    Arc::new(CorpusSessionData {
+        document_words: document_words.to_vec(),
+        document_counts: document_counts.to_vec(),
+        terms,
+    })
+}
+
+fn create_corpus_session_internal(
+    document_words: &[Vec<usize>],
+    document_counts: &[Vec<f64>],
+    terms: usize,
+) -> i64 {
+    let session_id = NEXT_CORPUS_SESSION_ID.fetch_add(1, Ordering::Relaxed);
+    let session = CorpusSession {
+        data: corpus_session_data(document_words, document_counts, terms),
+        config: None,
+    };
+
+    match corpus_sessions().lock() {
+        Ok(mut sessions) => {
+            sessions.insert(session_id, session);
+            session_id as i64
+        }
+        Err(_) => 0,
+    }
+}
+
+fn create_corpus_session(
+    document_words: Vec<Vec<usize>>,
+    document_counts: Vec<Vec<f64>>,
+    terms: usize,
+) -> i64 {
+    create_corpus_session_internal(document_words.as_slice(), document_counts.as_slice(), terms)
+}
+
+fn replace_corpus_session_internal(
+    session_id: i64,
+    document_words: &[Vec<usize>],
+    document_counts: &[Vec<f64>],
+    terms: usize,
+) -> i64 {
+    if terms == 0 {
+        return 0;
+    }
+
+    let replacement_data = corpus_session_data(document_words, document_counts, terms);
+    match corpus_sessions().lock() {
+        Ok(mut sessions) => {
+            if session_id > 0 {
+                let session_key = session_id as u64;
+                if let Some(session) = sessions.get_mut(&session_key) {
+                    session.data = replacement_data;
+                    session.config = None;
+                    return session_id;
+                }
+            }
+
+            let new_session_id = NEXT_CORPUS_SESSION_ID.fetch_add(1, Ordering::Relaxed);
+            sessions.insert(
+                new_session_id,
+                CorpusSession {
+                    data: replacement_data,
+                    config: None,
+                },
+            );
+            new_session_id as i64
+        }
+        Err(_) => 0,
+    }
+}
+
+fn replace_corpus_session(
+    session_id: i64,
+    document_words: Vec<Vec<usize>>,
+    document_counts: Vec<Vec<f64>>,
+    terms: usize,
+) -> i64 {
+    replace_corpus_session_internal(
+        session_id,
+        document_words.as_slice(),
+        document_counts.as_slice(),
+        terms,
+    )
+}
+
+fn ensure_corpus_session(
+    session_id: i64,
+    document_words: &[Vec<usize>],
+    document_counts: &[Vec<f64>],
+    terms: usize,
+) -> i64 {
+    if terms == 0 {
+        return 0;
+    }
+
+    if session_id > 0 && corpus_session_exists(session_id) {
+        return session_id;
+    }
+
+    create_corpus_session_internal(document_words, document_counts, terms)
+}
+
+fn drop_corpus_session(session_id: i64) -> bool {
+    if session_id <= 0 {
+        return false;
+    }
+
+    let session_key = session_id as u64;
+    match corpus_sessions().lock() {
+        Ok(mut sessions) => sessions.remove(&session_key).is_some(),
+        Err(_) => false,
+    }
+}
+
+fn configure_corpus_session(
+    session_id: i64,
+    topics: usize,
+    max_iter: i64,
+    convergence: f64,
+    em_max_iter: i64,
+    em_convergence: f64,
+    init_alpha: f64,
+    min_probability: f64,
+) -> bool {
+    if session_id <= 0 || topics == 0 {
+        return false;
+    }
+
+    let session_key = session_id as u64;
+    match corpus_sessions().lock() {
+        Ok(mut sessions) => {
+            let Some(session) = sessions.get_mut(&session_key) else {
+                return false;
+            };
+
+            session.config = Some(SessionConfig {
+                topics,
+                max_iter,
+                convergence,
+                em_max_iter,
+                em_convergence,
+                init_alpha,
+                min_probability,
+            });
+
+            true
+        }
+        Err(_) => false,
+    }
+}
+
+fn run_em_on_session_with_start_seed(
+    session_id: i64,
+    start: String,
+    topics: usize,
+    max_iter: i64,
+    convergence: f64,
+    em_max_iter: i64,
+    em_convergence: f64,
+    init_alpha: f64,
+    min_probability: f64,
+    random_seed: i64,
+) -> (Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<Vec<f64>>>) {
+    if session_id <= 0 {
+        return empty_em_output();
+    }
+
+    let session_key = session_id as u64;
+    let session_data = match corpus_sessions().lock() {
+        Ok(sessions) => sessions
+            .get(&session_key)
+            .map(|session| Arc::clone(&session.data)),
+        Err(_) => None,
+    };
+
+    let Some(session_data) = session_data else {
+        return empty_em_output();
+    };
+
+    run_em_with_start_seed_internal(
+        start.as_str(),
+        session_data.document_words.as_slice(),
+        session_data.document_counts.as_slice(),
+        topics,
+        session_data.terms,
+        max_iter,
+        convergence,
+        em_max_iter,
+        em_convergence,
+        init_alpha,
+        min_probability,
+        random_seed,
+    )
+}
+
+fn run_em_on_session(
+    session_id: i64,
+    start: String,
+    topics: usize,
+    max_iter: i64,
+    convergence: f64,
+    em_max_iter: i64,
+    em_convergence: f64,
+    init_alpha: f64,
+    min_probability: f64,
+    random_seed: i64,
+) -> (Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<Vec<f64>>>) {
+    if session_id <= 0 || topics == 0 {
+        return empty_em_output();
+    }
+
+    let desired_config = SessionConfig {
+        topics,
+        max_iter,
+        convergence,
+        em_max_iter,
+        em_convergence,
+        init_alpha,
+        min_probability,
+    };
+
+    let session_key = session_id as u64;
+    let session_data = match corpus_sessions().lock() {
+        Ok(mut sessions) => {
+            let Some(session) = sessions.get_mut(&session_key) else {
+                return empty_em_output();
+            };
+
+            if session.config.as_ref() != Some(&desired_config) {
+                session.config = Some(desired_config.clone());
+            }
+
+            Arc::clone(&session.data)
+        }
+        Err(_) => return empty_em_output(),
+    };
+
+    run_em_with_start_seed_internal(
+        start.as_str(),
+        session_data.document_words.as_slice(),
+        session_data.document_counts.as_slice(),
+        desired_config.topics,
+        session_data.terms,
+        desired_config.max_iter,
+        desired_config.convergence,
+        desired_config.em_max_iter,
+        desired_config.em_convergence,
+        desired_config.init_alpha,
+        desired_config.min_probability,
+        random_seed,
+    )
+}
+
+fn run_em_on_session_with_corpus(
+    session_id: i64,
+    document_words: Vec<Vec<usize>>,
+    document_counts: Vec<Vec<f64>>,
+    terms: usize,
+    start: String,
+    topics: usize,
+    max_iter: i64,
+    convergence: f64,
+    em_max_iter: i64,
+    em_convergence: f64,
+    init_alpha: f64,
+    min_probability: f64,
+    random_seed: i64,
+) -> (
+    i64,
+    Vec<Vec<f64>>,
+    Vec<Vec<f64>>,
+    Vec<Vec<f64>>,
+    Vec<Vec<Vec<f64>>>,
+) {
+    if topics == 0 || terms == 0 {
+        return empty_managed_session_em_output();
+    }
+
+    let active_session_id = ensure_corpus_session(
+        session_id,
+        document_words.as_slice(),
+        document_counts.as_slice(),
+        terms,
+    );
+
+    if active_session_id > 0 {
+        let (beta_probabilities, beta_log, gamma, phi) = run_em_on_session(
+            active_session_id,
+            start.clone(),
+            topics,
+            max_iter,
+            convergence,
+            em_max_iter,
+            em_convergence,
+            init_alpha,
+            min_probability,
+            random_seed,
+        );
+
+        if !(beta_probabilities.is_empty()
+            && beta_log.is_empty()
+            && gamma.is_empty()
+            && phi.is_empty())
+        {
+            return (active_session_id, beta_probabilities, beta_log, gamma, phi);
+        }
+    }
+
+    let (beta_probabilities, beta_log, gamma, phi) = run_em_with_start_seed_internal(
+        start.as_str(),
+        document_words.as_slice(),
+        document_counts.as_slice(),
+        topics,
+        terms,
+        max_iter,
+        convergence,
+        em_max_iter,
+        em_convergence,
+        init_alpha,
+        min_probability,
+        random_seed,
+    );
+
+    if beta_probabilities.is_empty() && beta_log.is_empty() && gamma.is_empty() && phi.is_empty() {
+        return empty_managed_session_em_output();
+    }
+
+    (active_session_id, beta_probabilities, beta_log, gamma, phi)
+}
+
+fn run_em_on_session_start(
+    session_id: i64,
+    start: String,
+    random_seed: i64,
+) -> (Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<Vec<f64>>>) {
+    if session_id <= 0 {
+        return empty_em_output();
+    }
+
+    let session_key = session_id as u64;
+    let session_data = match corpus_sessions().lock() {
+        Ok(sessions) => sessions.get(&session_key).map(|session| {
+            (
+                Arc::clone(&session.data),
+                session.config.clone(),
+            )
+        }),
+        Err(_) => None,
+    };
+
+    let Some((session_data, config)) = session_data else {
+        return empty_em_output();
+    };
+
+    let Some(config) = config else {
+        return empty_em_output();
+    };
+
+    run_em_with_start_seed_internal(
+        start.as_str(),
+        session_data.document_words.as_slice(),
+        session_data.document_counts.as_slice(),
+        config.topics,
+        session_data.terms,
+        config.max_iter,
+        config.convergence,
+        config.em_max_iter,
+        config.em_convergence,
+        config.init_alpha,
+        config.min_probability,
+        random_seed,
+    )
+}
+
 fn infer_document_internal(
     beta_probabilities: &[Vec<f64>],
     gamma_initial: &[f64],
@@ -360,10 +881,10 @@ fn infer_document(
     output
 }
 
-fn infer_corpus_iteration(
-    beta_probabilities: Vec<Vec<f64>>,
-    document_words: Vec<Vec<usize>>,
-    document_counts: Vec<Vec<f64>>,
+fn infer_corpus_iteration_internal(
+    beta_probabilities: &[Vec<f64>],
+    document_words: &[Vec<usize>],
+    document_counts: &[Vec<f64>],
     max_iter: i64,
     convergence: f64,
     min_probability: f64,
@@ -392,7 +913,7 @@ fn infer_corpus_iteration(
         let gamma_initial = vec![init_alpha_value + (total / topics as f64); topics];
 
         let (gamma_d, phi_d) = infer_document_internal(
-            beta_probabilities.as_slice(),
+            beta_probabilities,
             gamma_initial.as_slice(),
             words.as_slice(),
             counts.as_slice(),
@@ -416,6 +937,264 @@ fn infer_corpus_iteration(
     (gamma_matrix, phi_tensor, topic_term_counts)
 }
 
+fn infer_corpus_iteration(
+    beta_probabilities: Vec<Vec<f64>>,
+    document_words: Vec<Vec<usize>>,
+    document_counts: Vec<Vec<f64>>,
+    max_iter: i64,
+    convergence: f64,
+    min_probability: f64,
+    init_alpha: f64,
+) -> (Vec<Vec<f64>>, Vec<Vec<Vec<f64>>>, Vec<Vec<f64>>) {
+    infer_corpus_iteration_internal(
+        beta_probabilities.as_slice(),
+        document_words.as_slice(),
+        document_counts.as_slice(),
+        max_iter,
+        convergence,
+        min_probability,
+        init_alpha,
+    )
+}
+
+fn start_uses_seeded_initialization(start: &str) -> bool {
+    let normalized = start.trim().to_ascii_lowercase();
+    normalized == "seeded" || normalized == "deterministic"
+}
+
+fn start_uses_random_initialization(start: &str) -> bool {
+    start.trim().eq_ignore_ascii_case("random")
+}
+
+fn run_em_internal(
+    mut beta_probabilities: Vec<Vec<f64>>,
+    document_words: &[Vec<usize>],
+    document_counts: &[Vec<f64>],
+    max_iter: i64,
+    convergence: f64,
+    em_max_iter: i64,
+    em_convergence: f64,
+    init_alpha: f64,
+    min_probability: f64,
+) -> (Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<Vec<f64>>>) {
+    let em_max_iter_value = if em_max_iter <= 0 { 0 } else { em_max_iter as usize };
+    let em_convergence_value = if em_convergence.is_finite() && em_convergence >= 0.0 {
+        em_convergence
+    } else {
+        1.0e-4
+    };
+
+    let mut previous_gamma: Option<Vec<Vec<f64>>> = None;
+    let mut beta_log: Vec<Vec<f64>> = Vec::new();
+    let mut gamma: Vec<Vec<f64>> = Vec::new();
+    let mut phi: Vec<Vec<Vec<f64>>> = Vec::new();
+
+    for _ in 0..em_max_iter_value {
+        let (current_gamma, current_phi, topic_term_counts) = infer_corpus_iteration_internal(
+            beta_probabilities.as_slice(),
+            document_words,
+            document_counts,
+            max_iter,
+            convergence,
+            min_probability,
+            init_alpha,
+        );
+
+        let (next_beta_probabilities, next_beta_log) =
+            normalize_topic_term_counts(topic_term_counts, min_probability);
+        let should_stop = previous_gamma
+            .as_ref()
+            .map(|prev| {
+                average_gamma_shift_internal(prev.as_slice(), current_gamma.as_slice())
+                    <= em_convergence_value
+            })
+            .unwrap_or(false);
+
+        beta_probabilities = next_beta_probabilities;
+        beta_log = next_beta_log;
+        gamma = current_gamma;
+        phi = current_phi;
+
+        if should_stop {
+            break;
+        }
+
+        previous_gamma = Some(gamma.clone());
+    }
+
+    (beta_probabilities, beta_log, gamma, phi)
+}
+
+fn run_em(
+    beta_probabilities: Vec<Vec<f64>>,
+    document_words: Vec<Vec<usize>>,
+    document_counts: Vec<Vec<f64>>,
+    max_iter: i64,
+    convergence: f64,
+    em_max_iter: i64,
+    em_convergence: f64,
+    init_alpha: f64,
+    min_probability: f64,
+) -> (Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<Vec<f64>>>) {
+    run_em_internal(
+        beta_probabilities,
+        document_words.as_slice(),
+        document_counts.as_slice(),
+        max_iter,
+        convergence,
+        em_max_iter,
+        em_convergence,
+        init_alpha,
+        min_probability,
+    )
+}
+
+fn run_em_with_start_internal(
+    start: &str,
+    document_words: &[Vec<usize>],
+    document_counts: &[Vec<f64>],
+    topics: usize,
+    terms: usize,
+    max_iter: i64,
+    convergence: f64,
+    em_max_iter: i64,
+    em_convergence: f64,
+    init_alpha: f64,
+    min_probability: f64,
+) -> (Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<Vec<f64>>>) {
+    let initial_beta =
+        if start_uses_seeded_initialization(start) || start_uses_random_initialization(start) {
+            seeded_topic_term_probabilities_internal(
+                document_words,
+                document_counts,
+                topics,
+                terms,
+                min_probability,
+            )
+        } else {
+            // Unknown start modes default to seeded initialization for a stable fallback.
+            seeded_topic_term_probabilities_internal(
+                document_words,
+                document_counts,
+                topics,
+                terms,
+                min_probability,
+            )
+        };
+
+    run_em_internal(
+        initial_beta,
+        document_words,
+        document_counts,
+        max_iter,
+        convergence,
+        em_max_iter,
+        em_convergence,
+        init_alpha,
+        min_probability,
+    )
+}
+
+fn run_em_with_start(
+    start: String,
+    document_words: Vec<Vec<usize>>,
+    document_counts: Vec<Vec<f64>>,
+    topics: usize,
+    terms: usize,
+    max_iter: i64,
+    convergence: f64,
+    em_max_iter: i64,
+    em_convergence: f64,
+    init_alpha: f64,
+    min_probability: f64,
+) -> (Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<Vec<f64>>>) {
+    run_em_with_start_internal(
+        start.as_str(),
+        document_words.as_slice(),
+        document_counts.as_slice(),
+        topics,
+        terms,
+        max_iter,
+        convergence,
+        em_max_iter,
+        em_convergence,
+        init_alpha,
+        min_probability,
+    )
+}
+
+fn run_em_with_start_seed_internal(
+    start: &str,
+    document_words: &[Vec<usize>],
+    document_counts: &[Vec<f64>],
+    topics: usize,
+    terms: usize,
+    max_iter: i64,
+    convergence: f64,
+    em_max_iter: i64,
+    em_convergence: f64,
+    init_alpha: f64,
+    min_probability: f64,
+    random_seed: i64,
+) -> (Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<Vec<f64>>>) {
+    let initial_beta = if start_uses_seeded_initialization(start) {
+        seeded_topic_term_probabilities_internal(
+            document_words,
+            document_counts,
+            topics,
+            terms,
+            min_probability,
+        )
+    } else if start_uses_random_initialization(start) {
+        random_topic_term_probabilities(topics, terms, min_probability, random_seed)
+    } else {
+        // Unknown start modes follow Ruby's non-seeded fallback behavior.
+        random_topic_term_probabilities(topics, terms, min_probability, random_seed)
+    };
+
+    run_em_internal(
+        initial_beta,
+        document_words,
+        document_counts,
+        max_iter,
+        convergence,
+        em_max_iter,
+        em_convergence,
+        init_alpha,
+        min_probability,
+    )
+}
+
+fn run_em_with_start_seed(
+    start: String,
+    document_words: Vec<Vec<usize>>,
+    document_counts: Vec<Vec<f64>>,
+    topics: usize,
+    terms: usize,
+    max_iter: i64,
+    convergence: f64,
+    em_max_iter: i64,
+    em_convergence: f64,
+    init_alpha: f64,
+    min_probability: f64,
+    random_seed: i64,
+) -> (Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<Vec<Vec<f64>>>) {
+    run_em_with_start_seed_internal(
+        start.as_str(),
+        document_words.as_slice(),
+        document_counts.as_slice(),
+        topics,
+        terms,
+        max_iter,
+        convergence,
+        em_max_iter,
+        em_convergence,
+        init_alpha,
+        min_probability,
+        random_seed,
+    )
+}
+
 #[magnus::init]
 fn init() -> Result<(), Error> {
     let lda_module = define_module("Lda")?;
@@ -423,6 +1202,8 @@ fn init() -> Result<(), Error> {
 
     rust_backend_module.define_singleton_method("available?", function!(available, 0))?;
     rust_backend_module.define_singleton_method("abi_version", function!(abi_version, 0))?;
+    rust_backend_module.define_singleton_method("corpus_session_count", function!(corpus_session_count, 0))?;
+    rust_backend_module.define_singleton_method("corpus_session_exists", function!(corpus_session_exists, 1))?;
     rust_backend_module.define_singleton_method("before_em", function!(before_em, 3))?;
     rust_backend_module.define_singleton_method(
         "topic_weights_for_word",
@@ -451,6 +1232,32 @@ fn init() -> Result<(), Error> {
         "seeded_topic_term_probabilities",
         function!(seeded_topic_term_probabilities, 5),
     )?;
+    rust_backend_module.define_singleton_method(
+        "random_topic_term_probabilities",
+        function!(random_topic_term_probabilities, 4),
+    )?;
+    rust_backend_module
+        .define_singleton_method("create_corpus_session", function!(create_corpus_session, 3))?;
+    rust_backend_module
+        .define_singleton_method("replace_corpus_session", function!(replace_corpus_session, 4))?;
+    rust_backend_module
+        .define_singleton_method("drop_corpus_session", function!(drop_corpus_session, 1))?;
+    rust_backend_module
+        .define_singleton_method("configure_corpus_session", function!(configure_corpus_session, 8))?;
+    rust_backend_module.define_singleton_method("run_em", function!(run_em, 9))?;
+    rust_backend_module
+        .define_singleton_method("run_em_with_start", function!(run_em_with_start, 11))?;
+    rust_backend_module
+        .define_singleton_method("run_em_with_start_seed", function!(run_em_with_start_seed, 12))?;
+    rust_backend_module.define_singleton_method(
+        "run_em_on_session_with_start_seed",
+        function!(run_em_on_session_with_start_seed, 10),
+    )?;
+    rust_backend_module.define_singleton_method("run_em_on_session", function!(run_em_on_session, 10))?;
+    rust_backend_module
+        .define_singleton_method("run_em_on_session_with_corpus", function!(run_em_on_session_with_corpus, 13))?;
+    rust_backend_module
+        .define_singleton_method("run_em_on_session_start", function!(run_em_on_session_start, 3))?;
 
     Ok(())
 }