lda-ruby 0.4.0 → 0.5.0

data/lib/lda-ruby/backends/base.rb

@@ -109,6 +109,10 @@ module Lda
 
       private
 
+      def next_random_seed
+        @random.rand(0..9_223_372_036_854_775_807)
+      end
+
       def normalize!(weights)
         total = weights.sum.to_f
 

data/lib/lda-ruby/backends/pure_ruby.rb

@@ -46,61 +46,48 @@ module Lda
       end
 
       def em(start)
-        return nil if @corpus.nil? || @corpus.num_docs.zero?
-
-        topics = Integer(num_topics)
-        raise ArgumentError, "num_topics must be greater than zero" if topics <= 0
-
-        terms = max_term_index + 1
-        raise ArgumentError, "corpus must contain terms" if terms <= 0
+        em_input = build_em_input(start)
+        return nil if em_input.nil?
 
-        document_words = @corpus.documents.map { |document| document.words.map(&:to_i) }
-        document_counts = @corpus.documents.map { |document| document.counts.map(&:to_f) }
+        run_em_iterations(em_input)
+        nil
+      end
 
-        @beta_probabilities =
-          if start.to_s.strip.casecmp("seeded").zero? || start.to_s.strip.casecmp("deterministic").zero?
-            seeded_topic_term_probabilities(topics, terms, document_words, document_counts)
-          else
-            initial_topic_term_probabilities(topics, terms)
-          end
+      # Returns an EM input snapshot that can be reused by Rust orchestration
+      # and Ruby fallback paths without re-sampling random initialization.
+      def rust_em_input(start)
+        build_em_input(start)
+      end
 
-        document_totals = document_counts.map { |counts| counts.sum.to_f }
-        document_lengths = document_words.map(&:length)
+      # Returns only the initial beta matrix for Rust compatibility paths that
+      # already hold a cached corpus snapshot.
+      def rust_initial_beta_probabilities(start, document_words, document_counts, topics, terms)
+        start_mode = start.to_s
 
-        previous_gamma = nil
+        if start_mode.strip.casecmp("seeded").zero? || start_mode.strip.casecmp("deterministic").zero?
+          seeded_topic_term_probabilities(
+            Integer(topics),
+            Integer(terms),
+            document_words,
+            document_counts
+          )
+        else
+          initial_topic_term_probabilities(Integer(topics), Integer(terms))
+        end
+      end
 
-        Integer(em_max_iter).times do
-          if @trusted_kernel_outputs && @corpus_iteration_kernel
-            current_gamma, current_phi, topic_term_counts = infer_corpus_iteration(
-              nil,
-              document_words,
-              document_counts,
-              document_totals,
-              document_lengths,
-              topics,
-              terms
-            )
-          else
-            topic_term_counts = Array.new(topics) { Array.new(terms, MIN_PROBABILITY) }
-            current_gamma, current_phi, topic_term_counts = infer_corpus_iteration(
-              topic_term_counts,
-              document_words,
-              document_counts,
-              document_totals,
-              document_lengths,
-              topics,
-              terms
-            )
-          end
+      def em_from_input(em_input)
+        return nil if em_input.nil?
 
-          @beta_probabilities, @beta_log = finalize_topic_term_counts(topic_term_counts)
-          @gamma = current_gamma
-          @phi = current_phi
+        run_em_iterations(em_input)
+        nil
+      end
 
-          break if previous_gamma && average_gamma_shift(previous_gamma, current_gamma) <= Float(em_convergence)
-
-          previous_gamma = current_gamma
-        end
+      def apply_em_state(beta_probabilities:, beta_log:, gamma:, phi:)
+        @beta_probabilities = beta_probabilities
+        @beta_log = beta_log
+        @gamma = gamma
+        @phi = phi
 
         nil
       end
@@ -147,6 +134,81 @@ module Lda
 
       private
 
+      def build_em_input(start)
+        return nil if @corpus.nil? || @corpus.num_docs.zero?
+
+        topics = Integer(num_topics)
+        raise ArgumentError, "num_topics must be greater than zero" if topics <= 0
+
+        terms = max_term_index + 1
+        raise ArgumentError, "corpus must contain terms" if terms <= 0
+
+        document_words = @corpus.documents.map { |document| document.words.map(&:to_i) }
+        document_counts = @corpus.documents.map { |document| document.counts.map(&:to_f) }
+
+        {
+          topics: topics,
+          terms: terms,
+          document_words: document_words,
+          document_counts: document_counts,
+          document_totals: document_counts.map { |counts| counts.sum.to_f },
+          document_lengths: document_words.map(&:length),
+          initial_beta_probabilities: rust_initial_beta_probabilities(
+            start,
+            document_words,
+            document_counts,
+            topics,
+            terms
+          ),
+          min_probability: MIN_PROBABILITY
+        }
+      end
+
+      def run_em_iterations(em_input)
+        topics = em_input.fetch(:topics)
+        terms = em_input.fetch(:terms)
+        document_words = em_input.fetch(:document_words)
+        document_counts = em_input.fetch(:document_counts)
+        document_totals = em_input.fetch(:document_totals)
+        document_lengths = em_input.fetch(:document_lengths)
+
+        @beta_probabilities = em_input.fetch(:initial_beta_probabilities)
+        previous_gamma = nil
+
+        Integer(em_max_iter).times do
+          if @trusted_kernel_outputs && @corpus_iteration_kernel
+            current_gamma, current_phi, topic_term_counts = infer_corpus_iteration(
+              nil,
+              document_words,
+              document_counts,
+              document_totals,
+              document_lengths,
+              topics,
+              terms
+            )
+          else
+            topic_term_counts = Array.new(topics) { Array.new(terms, MIN_PROBABILITY) }
+            current_gamma, current_phi, topic_term_counts = infer_corpus_iteration(
+              topic_term_counts,
+              document_words,
+              document_counts,
+              document_totals,
+              document_lengths,
+              topics,
+              terms
+            )
+          end
+
+          @beta_probabilities, @beta_log = finalize_topic_term_counts(topic_term_counts)
+          @gamma = current_gamma
+          @phi = current_phi
+
+          break if previous_gamma && average_gamma_shift(previous_gamma, current_gamma) <= Float(em_convergence)
+
+          previous_gamma = current_gamma
+        end
+      end
+
       def max_term_index
         return -1 if @corpus.nil? || @corpus.documents.empty?
 

data/lib/lda-ruby/backends/rust.rb

@@ -4,6 +4,7 @@ module Lda
   module Backends
     class Rust < Base
       SETTINGS = %i[max_iter convergence em_max_iter em_convergence num_topics init_alpha est_alpha verbose].freeze
+      MIN_PROBABILITY = 1e-12
 
       def self.available?
         return false unless defined?(::Lda::RUST_EXTENSION_LOADED) && ::Lda::RUST_EXTENSION_LOADED
@@ -32,6 +33,12 @@ module Lda
         super(random_seed: random_seed)
         raise LoadError, "Rust backend is unavailable for this environment" unless self.class.available?
 
+        @rust_corpus_session_id = nil
+        @rust_corpus_terms = nil
+        @rust_document_lengths = nil
+        @rust_document_words = nil
+        @rust_document_counts = nil
+
         @fallback = PureRuby.new(random_seed: random_seed)
         @fallback.topic_weights_kernel = method(:rust_topic_weights_for_word)
         @fallback.topic_term_accumulator_kernel = method(:rust_accumulate_topic_term_counts)
@@ -49,20 +56,22 @@ module Lda
       end
 
       def corpus=(corpus)
+        previous_session_id = @rust_corpus_session_id
         @corpus = corpus
         @fallback.corpus = corpus
+        register_rust_corpus_session(previous_session_id)
         true
       end
 
       def fast_load_corpus_from_file(filename)
         loaded = @fallback.fast_load_corpus_from_file(filename)
-        @corpus = @fallback.corpus
+        self.corpus = @fallback.corpus
         loaded
       end
 
       def load_settings(settings_file)
         loaded = @fallback.load_settings(settings_file)
-        @corpus = @fallback.corpus
+        self.corpus = @fallback.corpus
         loaded
       end
 
@@ -71,7 +80,10 @@ module Lda
       end
 
       def em(start)
-        rust_before_em(start)
+        start_mode = start.to_s
+        rust_before_em(start_mode)
+        return nil if rust_orchestrated_em(start_mode)
+
         @fallback.em(start)
       end
 
@@ -93,6 +105,375 @@ module Lda
 
       private
 
+      def rust_orchestrated_em(start)
+        managed_orchestrated = rust_orchestrated_em_with_managed_corpus(start)
+        return true if managed_orchestrated
+
+        direct_orchestrated = rust_orchestrated_em_with_start_seed(start)
+        return true if direct_orchestrated
+
+        rust_orchestrated_em_with_beta(start)
+      end
+
+      def rust_orchestrated_em_with_managed_corpus(start)
+        return false unless defined?(::Lda::RustBackend)
+        return false unless ensure_rust_corpus_snapshot
+
+        random_seed = Integer(next_random_seed)
+        if ::Lda::RustBackend.respond_to?(:run_em_on_session_with_corpus)
+          managed_output = ::Lda::RustBackend.run_em_on_session_with_corpus(
+            Integer(@rust_corpus_session_id || 0),
+            @rust_document_words,
+            @rust_document_counts,
+            Integer(@rust_corpus_terms),
+            start.to_s,
+            *current_rust_session_config_signature,
+            random_seed
+          )
+
+          return false unless managed_output.is_a?(Array) && managed_output.size == 5
+
+          session_id, beta_probabilities, beta_log, gamma, phi = managed_output
+          output = [beta_probabilities, beta_log, gamma, phi]
+          return false unless valid_rust_em_output?(output, @rust_document_lengths, Integer(num_topics), Integer(@rust_corpus_terms))
+
+          @rust_corpus_session_id =
+            if session_id.is_a?(Numeric) && session_id.positive?
+              Integer(session_id)
+            end
+          @fallback.apply_em_state(
+            beta_probabilities: beta_probabilities,
+            beta_log: beta_log,
+            gamma: gamma,
+            phi: phi
+          )
+          return true
+        end
+
+        return false unless ::Lda::RustBackend.respond_to?(:run_em_on_session)
+        return false unless @rust_corpus_session_id
+
+        output = ::Lda::RustBackend.run_em_on_session(
+          Integer(@rust_corpus_session_id),
+          start.to_s,
+          *current_rust_session_config_signature,
+          random_seed
+        )
+
+        return false unless valid_rust_em_output?(output, @rust_document_lengths, Integer(num_topics), Integer(@rust_corpus_terms))
+
+        beta_probabilities, beta_log, gamma, phi = output
+        @fallback.apply_em_state(
+          beta_probabilities: beta_probabilities,
+          beta_log: beta_log,
+          gamma: gamma,
+          phi: phi
+        )
+        true
+      rescue StandardError
+        false
+      end
+
+      def rust_orchestrated_em_with_start_seed(start)
+        return false unless defined?(::Lda::RustBackend)
+        return false unless ::Lda::RustBackend.respond_to?(:run_em_with_start_seed)
+        return false unless ensure_rust_corpus_snapshot
+
+        topics = Integer(num_topics)
+        return false unless topics.positive?
+
+        random_seed = Integer(next_random_seed)
+        output = ::Lda::RustBackend.run_em_with_start_seed(
+          start.to_s,
+          @rust_document_words,
+          @rust_document_counts,
+          topics,
+          Integer(@rust_corpus_terms),
+          Integer(max_iter),
+          Float(convergence),
+          Integer(em_max_iter),
+          Float(em_convergence),
+          Float(init_alpha),
+          MIN_PROBABILITY,
+          random_seed
+        )
+
+        return false unless valid_rust_em_output?(
+          output,
+          @rust_document_lengths,
+          topics,
+          Integer(@rust_corpus_terms)
+        )
+
+        beta_probabilities, beta_log, gamma, phi = output
+        @fallback.apply_em_state(
+          beta_probabilities: beta_probabilities,
+          beta_log: beta_log,
+          gamma: gamma,
+          phi: phi
+        )
+        true
+      rescue StandardError
+        false
+      end
+
+      def rust_orchestrated_em_with_beta(start)
+        return false unless defined?(::Lda::RustBackend)
+        return false unless ::Lda::RustBackend.respond_to?(:run_em)
+
+        em_input =
+          if ensure_rust_corpus_snapshot && @fallback.respond_to?(:rust_initial_beta_probabilities)
+            topics = Integer(num_topics)
+            terms = Integer(@rust_corpus_terms)
+            initial_beta_probabilities = @fallback.rust_initial_beta_probabilities(
+              start,
+              @rust_document_words,
+              @rust_document_counts,
+              topics,
+              terms
+            )
+
+            {
+              topics: topics,
+              terms: terms,
+              document_words: @rust_document_words,
+              document_counts: @rust_document_counts,
+              document_totals: @rust_document_counts.map { |counts| counts.sum.to_f },
+              document_lengths: @rust_document_lengths,
+              initial_beta_probabilities: initial_beta_probabilities,
+              min_probability: MIN_PROBABILITY
+            }
+          else
+            @fallback.rust_em_input(start)
+          end
+
+        return true if em_input.nil?
+
+        output = ::Lda::RustBackend.run_em(
+          em_input.fetch(:initial_beta_probabilities),
+          em_input.fetch(:document_words),
+          em_input.fetch(:document_counts),
+          Integer(max_iter),
+          Float(convergence),
+          Integer(em_max_iter),
+          Float(em_convergence),
+          Float(init_alpha),
+          Float(em_input.fetch(:min_probability))
+        )
+
+        unless valid_rust_em_output?(
+          output,
+          em_input.fetch(:document_lengths),
+          em_input.fetch(:topics),
+          em_input.fetch(:terms)
+        )
+          @fallback.em_from_input(em_input)
+          return true
+        end
+
+        beta_probabilities, beta_log, gamma, phi = output
+        @fallback.apply_em_state(
+          beta_probabilities: beta_probabilities,
+          beta_log: beta_log,
+          gamma: gamma,
+          phi: phi
+        )
+        true
+      rescue StandardError
+        if defined?(em_input) && em_input
+          @fallback.em_from_input(em_input)
+          return true
+        end
+
+        false
+      end
+
+      def rust_em_corpus_input
+        return nil if @corpus.nil? || @corpus.num_docs.zero?
+
+        topics = Integer(num_topics)
+        raise ArgumentError, "num_topics must be greater than zero" if topics <= 0
+
+        terms = max_term_index + 1
+        raise ArgumentError, "corpus must contain terms" if terms <= 0
+
+        document_words = @corpus.documents.map { |document| document.words.map(&:to_i) }
+        document_counts = @corpus.documents.map { |document| document.counts.map(&:to_f) }
+
+        {
+          topics: topics,
+          terms: terms,
+          document_words: document_words,
+          document_counts: document_counts,
+          document_lengths: document_words.map(&:length),
+          min_probability: MIN_PROBABILITY
+        }
+      end
+
+      def max_term_index
+        return -1 if @corpus.nil? || @corpus.documents.empty?
+
+        @corpus.documents
+          .flat_map(&:words)
+          .max || -1
+      end
+
+      def register_rust_corpus_session(previous_session_id = nil)
+        @rust_corpus_session_id = nil
+        @rust_corpus_terms = nil
+        @rust_document_lengths = nil
+        @rust_document_words = nil
+        @rust_document_counts = nil
+
+        if @corpus.nil?
+          drop_rust_corpus_session_by_id(previous_session_id)
+          return
+        end
+
+        return unless defined?(::Lda::RustBackend)
+
+        em_input = rust_em_corpus_input
+        if em_input.nil?
+          drop_rust_corpus_session_by_id(previous_session_id)
+          return
+        end
+
+        @rust_corpus_terms = Integer(em_input.fetch(:terms))
+        @rust_document_lengths = em_input.fetch(:document_lengths)
+        @rust_document_words = em_input.fetch(:document_words)
+        @rust_document_counts = em_input.fetch(:document_counts)
+
+        session_id =
+          if ::Lda::RustBackend.respond_to?(:replace_corpus_session)
+            ::Lda::RustBackend.replace_corpus_session(
+              Integer(previous_session_id || 0),
+              @rust_document_words,
+              @rust_document_counts,
+              Integer(@rust_corpus_terms)
+            )
+          elsif ::Lda::RustBackend.respond_to?(:create_corpus_session)
+            drop_rust_corpus_session_by_id(previous_session_id)
+            ::Lda::RustBackend.create_corpus_session(
+              @rust_document_words,
+              @rust_document_counts,
+              Integer(@rust_corpus_terms)
+            )
+          end
+
+        unless session_id.is_a?(Numeric) && session_id.positive?
+          drop_rust_corpus_session_by_id(previous_session_id)
+          return
+        end
+
+        @rust_corpus_session_id = Integer(session_id)
+      rescue StandardError
+        @rust_corpus_session_id = nil
+        @rust_corpus_terms = nil
+        @rust_document_lengths = nil
+        @rust_document_words = nil
+        @rust_document_counts = nil
+        drop_rust_corpus_session_by_id(previous_session_id)
+      end
+
+      def ensure_rust_corpus_snapshot
+        has_session_data = @rust_corpus_terms && @rust_document_lengths && @rust_document_words && @rust_document_counts
+        return true if has_session_data
+
+        register_rust_corpus_session(@rust_corpus_session_id)
+        @rust_corpus_terms && @rust_document_lengths && @rust_document_words && @rust_document_counts
+      rescue StandardError
+        false
+      end
+
+      def release_rust_corpus_session
+        session_id = @rust_corpus_session_id
+
+        @rust_corpus_session_id = nil
+        @rust_corpus_terms = nil
+        @rust_document_lengths = nil
+        @rust_document_words = nil
+        @rust_document_counts = nil
+
+        drop_rust_corpus_session_by_id(session_id)
+      rescue StandardError
+        nil
+      end
+
+      def drop_rust_corpus_session_by_id(session_id)
+        return unless session_id
+        return unless defined?(::Lda::RustBackend)
+        return unless ::Lda::RustBackend.respond_to?(:drop_corpus_session)
+
+        ::Lda::RustBackend.drop_corpus_session(Integer(session_id))
+      rescue StandardError
+        nil
+      end
+
+      def current_rust_session_config_signature
+        [
+          Integer(num_topics),
+          Integer(max_iter),
+          Float(convergence),
+          Integer(em_max_iter),
+          Float(em_convergence),
+          Float(init_alpha),
+          MIN_PROBABILITY
+        ]
+      end
+
+      def valid_rust_em_output?(output, document_lengths, topics, terms)
+        return false unless output.is_a?(Array)
+        return false unless output.size == 4
+
+        beta_probabilities, beta_log, gamma, phi = output
+
+        valid_topic_term_matrix?(beta_probabilities, topics, terms) &&
+          valid_topic_term_matrix?(beta_log, topics, terms) &&
+          valid_gamma_matrix?(gamma, document_lengths.size, topics) &&
+          valid_phi_tensor?(phi, document_lengths, topics)
+      end
+
+      def valid_topic_term_matrix?(matrix, topics, terms)
+        return false unless matrix.is_a?(Array)
+        return false unless matrix.size == topics
+
+        matrix.all? do |row|
+          row.is_a?(Array) &&
+            row.size == terms &&
+            row.all? { |value| finite_numeric?(value) }
+        end
+      end
+
+      def valid_gamma_matrix?(gamma, expected_docs, topics)
+        return false unless gamma.is_a?(Array)
+        return false unless gamma.size == expected_docs
+
+        gamma.all? do |row|
+          row.is_a?(Array) &&
+            row.size == topics &&
+            row.all? { |value| finite_numeric?(value) && value.positive? }
+        end
+      end
+
+      def valid_phi_tensor?(phi, document_lengths, topics)
+        return false unless phi.is_a?(Array)
+        return false unless phi.size == document_lengths.size
+
+        phi.each_with_index.all? do |doc_phi, doc_index|
+          doc_phi.is_a?(Array) &&
+            doc_phi.size == document_lengths[doc_index] &&
+            doc_phi.all? do |row|
+              row.is_a?(Array) &&
+                row.size == topics &&
+                row.all? { |value| finite_numeric?(value) }
+            end
+        end
+      end
+
+      def finite_numeric?(value)
+        value.is_a?(Numeric) && value.finite?
+      end
+
       def rust_before_em(start)
         return unless defined?(::Lda::RustBackend)
         return unless ::Lda::RustBackend.respond_to?(:before_em)

data/lib/lda-ruby/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Lda
-  VERSION = "0.4.0"
+  VERSION = "0.5.0"
 end

data/test/benchmark_scripts_test.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require_relative "test_helper"
+require "open3"
+
+class BenchmarkScriptsTest < Test::Unit::TestCase
+  def setup
+    @repo_root = File.expand_path("..", __dir__)
+    @check_rust_benchmark = File.join(@repo_root, "bin", "check-rust-benchmark")
+  end
+
+  def test_check_rust_benchmark_help
+    stdout, stderr, status = Open3.capture3(@check_rust_benchmark, "--help", chdir: @repo_root)
+    assert(status.success?, stderr)
+    assert_match(/Usage: \.\/bin\/check-rust-benchmark/, stdout)
+  end
+
+  def test_check_rust_benchmark_rejects_unknown_argument
+    _stdout, stderr, status = Open3.capture3(@check_rust_benchmark, "--unknown", chdir: @repo_root)
+    assert(!status.success?, "expected check-rust-benchmark to fail for unknown args")
+    assert_match(/Unknown argument/, stderr)
+  end
+end