kotoshu 0.3.0

data/lib/kotoshu/embeddings/similarity_search.rb

@@ -0,0 +1,331 @@
+# frozen_string_literal: true
+
+require_relative 'vocabulary'
+require_relative 'onnx_runtime_model'
+
+module Kotoshu
+  module Embeddings
+    # Similarity search for embedding-based nearest neighbor lookup.
+    #
+    # Efficiently finds semantically similar words using cosine similarity.
+    # Supports both on-the-fly computation and pre-computed embedding matrices.
+    #
+    # @example Basic usage
+    #   search = SimilaritySearch.new(
+    #     vocabulary: vocab,
+    #     model: model
+    #   )
+    #   neighbors = search.find_nearest('hello', k: 10)
+    #
+    # @example With pre-loaded embedding matrix (faster)
+    #   search = SimilaritySearch.new(
+    #     vocabulary: vocab,
+    #     model: model,
+    #     preload_embeddings: true
+    #   )
+    #   neighbors = search.find_nearest('hello', k: 10)
+    class SimilaritySearch
+      # @return [Vocabulary] The vocabulary
+      attr_reader :vocabulary
+
+      # @return [OnnxRuntimeModel] The ONNX model
+      attr_reader :model
+
+      # @return [Boolean] Whether embeddings are pre-loaded
+      attr_reader :embeddings_loaded
+
+      # Create a new similarity search instance.
+      #
+      # @param vocabulary [Vocabulary] Word vocabulary
+      # @param model [OnnxRuntimeModel] ONNX model for embeddings
+      # @param preload_embeddings [Boolean] Whether to preload all embeddings
+      # @param max_cache_size [Integer] Maximum embeddings to cache (if not preloading)
+      def initialize(vocabulary:, model:, preload_embeddings: false, max_cache_size: 1000)
+        @vocabulary = vocabulary
+        @model = model
+        @preload_embeddings = preload_embeddings
+        @max_cache_size = max_cache_size
+
+        # Embedding cache (word -> vector)
+        @embedding_cache = {}
+
+        # Pre-loaded embedding matrix (for faster search)
+        @embedding_matrix = nil
+
+        # Track whether embeddings are preloaded
+        @embeddings_loaded = false
+
+        # Load embeddings if requested
+        preload_embeddings! if preload_embeddings
+      end
+
+      # Find k nearest neighbors for a word.
+      #
+      # @param query_word [String] The query word
+      # @param k [Integer] Number of neighbors to return
+      # @param exclude_self [Boolean] Whether to exclude the query word itself
+      # @param min_similarity [Float] Minimum similarity threshold (0.0 to 1.0)
+      # @return [Array<Hash>] Array of {word, similarity} hashes
+      def find_nearest(query_word, k: 10, exclude_self: true, min_similarity: 0.0)
+        # Get query embedding
+        query_vec = get_embedding(query_word)
+        return [] unless query_vec
+
+        # Find neighbors
+        if @embedding_matrix
+          nearest_from_matrix(query_vec, k, exclude_self, min_similarity)
+        else
+          nearest_from_cache(query_vec, k, exclude_self, min_similarity)
+        end
+      end
+
+      # Find k nearest neighbors for multiple words.
+      #
+      # @param query_words [Array<String>] Query words
+      # @param k [Integer] Number of neighbors per word
+      # @return [Hash<String, Array<Hash>>] Word to neighbors mapping
+      def find_nearest_batch(query_words, k: 10)
+        query_words.each_with_object({}) do |word, result|
+          result[word] = find_nearest(word, k: k)
+        end
+      end
+
+      # Compute similarity between two words.
+      #
+      # @param word1 [String] First word
+      # @param word2 [String] Second word
+      # @return [Float] Cosine similarity (-1.0 to 1.0, or nil if either word not found)
+      def similarity(word1, word2)
+        vec1 = get_embedding(word1)
+        vec2 = get_embedding(word2)
+
+        return nil unless vec1 && vec2
+
+        cosine_similarity(vec1, vec2)
+      end
+
+      # Compute similarity between two embedding vectors.
+      #
+      # @param vec1 [Array<Float>] First vector
+      # @param vec2 [Array<Float>] Second vector
+      # @return [Float] Cosine similarity (-1.0 to 1.0)
+      def cosine_similarity(vec1, vec2)
+        return 0.0 if vec1.nil? || vec2.nil?
+
+        # Compute dot product
+        dot = vec1.zip(vec2).sum { |a, b| a * b }
+
+        # Compute magnitudes
+        norm1 = Math.sqrt(vec1.sum { |x| x * x })
+        norm2 = Math.sqrt(vec2.sum { |x| x * x })
+
+        return 0.0 if norm1.zero? || norm2.zero?
+
+        dot / (norm1 * norm2)
+      end
+
+      # Preload all embeddings into memory for faster search.
+      #
+      # @return [Boolean] True if loaded successfully
+      def preload_embeddings!
+        return false if @embedding_matrix
+
+        # Get all indices
+        all_indices = (0...@vocabulary.size).to_a
+
+        # Batch load embeddings
+        vectors = @model.get_embeddings(all_indices)
+        return false if vectors.nil? || vectors.empty?
+
+        # Store as hash for now (could use Numo::SFloat for efficiency)
+        @embedding_matrix = {}
+        all_indices.zip(vectors).each do |idx, vec|
+          @embedding_matrix[idx] = vec
+        end
+
+        @embeddings_loaded = true
+        true
+      rescue StandardError => e
+        warn "Failed to preload embeddings: #{e.message}"
+        false
+      end
+
+      # Clear the embedding cache.
+      #
+      # @return [self] Self for chaining
+      def clear_cache
+        @embedding_cache.clear
+        @embedding_matrix = nil
+        @embeddings_loaded = false
+        self
+      end
+
+      # Get cache statistics.
+      #
+      # @return [Hash] Cache statistics
+      def cache_stats
+        stats = {
+          size: @embedding_cache.size,
+          max_size: @max_cache_size
+        }
+        stats[:hit_rate] = @cache_hits.to_f / (@cache_hits + @cache_misses) if defined?(@cache_hits)
+        stats
+      end
+
+      # String representation.
+      #
+      # @return [String] String representation
+      def to_s
+        "SimilaritySearch(vocab_size: #{@vocabulary.size}, loaded: #{@embeddings_loaded})"
+      end
+      alias inspect to_s
+
+      private
+
+      # Get embedding for a word (with caching).
+      #
+      # @param word [String] The word
+      # @return [Array<Float>, nil] Embedding vector or nil if not found
+      def get_embedding(word)
+        # Check cache first
+        if @embedding_cache.key?(word)
+          @cache_hits += 1 if defined?(@cache_hits)
+          return @embedding_cache[word]
+        end
+
+        @cache_misses ||= 0
+        @cache_hits ||= 0
+        @cache_misses += 1
+
+        # Get from model
+        index = @vocabulary.lookup(word)
+        return nil unless index
+
+        vec = if @embedding_matrix
+                @embedding_matrix[index]
+              else
+                @model.get_embedding(index)
+              end
+
+        return nil unless vec
+
+        # Cache if not preloading (preload has all in memory already)
+        unless @preload_embeddings
+          # Evict oldest if cache is full
+          if @embedding_cache.size >= @max_cache_size
+            @embedding_cache.shift
+          end
+          @embedding_cache[word] = vec
+        end
+
+        vec
+      end
+
+      # Find nearest neighbors using pre-loaded matrix.
+      #
+      # @param query_vec [Array<Float>] Query embedding
+      # @param k [Integer] Number of neighbors
+      # @param exclude_self [Boolean] Whether to exclude exact matches
+      # @param min_similarity [Float] Minimum similarity
+      # @return [Array<Hash>] Nearest neighbors
+      def nearest_from_matrix(query_vec, k, exclude_self, min_similarity)
+        similarities = []
+
+        @vocabulary.words.each do |word|
+          index = @vocabulary.lookup(word)
+          vec = @embedding_matrix[index]
+
+          next unless vec
+
+          sim = cosine_similarity(query_vec, vec)
+
+          # Skip exact match if requested
+          next if exclude_self && sim >= 0.9999
+
+          # Skip below threshold
+          next if sim < min_similarity
+
+          similarities << { word: word, similarity: sim }
+        end
+
+        # Sort by similarity (descending) and take top k
+        similarities.sort_by { |s| -s[:similarity] }.first(k)
+      end
+
+      # Find nearest neighbors using cache (no pre-loading).
+      #
+      # @param query_vec [Array<Float>] Query embedding
+      # @param k [Integer] Number of neighbors
+      # @param exclude_self [Boolean] Whether to exclude exact matches
+      # @param min_similarity [Float] Minimum similarity
+      # @return [Array<Hash>] Nearest neighbors
+      def nearest_from_cache(query_vec, k, exclude_self, min_similarity)
+        similarities = []
+
+        # Sample from vocabulary for efficiency (or use common words)
+        sample_words = sample_vocabulary(k * 10)
+
+        sample_words.each do |word|
+          vec = get_embedding(word)
+          next unless vec
+
+          sim = cosine_similarity(query_vec, vec)
+
+          # Skip exact match if requested
+          next if exclude_self && sim >= 0.9999
+
+          # Skip below threshold
+          next if sim < min_similarity
+
+          similarities << { word: word, similarity: sim }
+        end
+
+        # Sort by similarity (descending) and take top k
+        similarities.sort_by { |s| -s[:similarity] }.first(k)
+      end
+
+      # Sample words from vocabulary for search.
+      #
+      # Prioritizes common words (first N in vocabulary).
+      #
+      # @param n [Integer] Number of words to sample
+      # @return [Array<String>] Sampled words
+      def sample_vocabulary(n)
+        # Use first N words (FastText orders by frequency)
+        # plus a random sample of the rest
+        common_size = [n / 2, 100].min
+        random_size = n - common_size
+
+        common = @vocabulary.common_words(n: common_size)
+
+        if @vocabulary.size > common_size
+          # Get a random sample from the rest
+          rest = @vocabulary.words.drop(common_size)
+          random_sample = rest.sample(random_size)
+          common + random_sample
+        else
+          common
+        end
+      end
+
+      # Create from cache.
+      #
+      # @param language_code [String] ISO 639-1 language code
+      # @param cache [Cache::ModelCache, nil] Optional cache instance
+      # @param preload [Boolean] Whether to preload embeddings
+      # @return [SimilaritySearch, nil] New search instance or nil if not available
+      def self.from_cache(language_code, cache: nil, preload: false)
+        vocab = Vocabulary.from_cache(language_code, cache: cache)
+        model = OnnxRuntimeModel.from_cache(language_code, cache: cache)
+
+        return nil unless vocab && model
+
+        new(
+          vocabulary: vocab,
+          model: model,
+          preload_embeddings: preload
+        )
+      end
+    end
+  end
+end

data/lib/kotoshu/embeddings/vocabulary.rb

@@ -0,0 +1,257 @@
+# frozen_string_literal: true
+
+require 'json'
+require_relative 'protocol'
+
+# Vocabulary - Word to index mapping
+#
+# Provides efficient lookup from words to integer indices for embedding retrieval.
+# Supports JSON file loading and saving.
+#
+# @example Creating a vocabulary
+#   vocab = Kotoshu::Embeddings::Vocabulary.new(
+#     language_code: 'en',
+#     word_to_index: { 'hello' => 0, 'world' => 1 }
+#   )
+#
+# @example Loading from file
+#   vocab = Kotoshu::Embeddings::Vocabulary.from_file('/path/to/vocab.json', language_code: 'en')
+#
+class Vocabulary
+  include VocabularyProtocol
+
+  # @return [String] ISO 639-1 language code
+  attr_reader :language_code
+
+  # @return [Hash{String => Integer}] Word to index mapping
+  attr_reader :word_to_index
+
+  # @return [Array<String>] Index to word mapping (sparse array)
+  attr_reader :index_to_word
+
+  # Create a new vocabulary
+  #
+  # @param language_code [String] ISO 639-1 language code
+  # @param word_to_index [Hash{String => Integer}] Word to index mapping
+  #
+  # @raise [ArgumentError] If word_to_index is empty
+  #
+  def initialize(language_code:, word_to_index:)
+    raise ArgumentError, 'word_to_index cannot be empty' if word_to_index.nil? || word_to_index.empty?
+
+    @language_code = language_code
+    @word_to_index = word_to_index.dup.freeze
+
+    # Build reverse index (index -> word)
+    @index_to_word = Array.new(@word_to_index.size)
+    @word_to_index.each do |word, index|
+      @index_to_word[index] = word if index < @index_to_word.size
+    end
+    @index_to_word.freeze
+  end
+
+  # Look up word index
+  #
+  # @param word [String] The word to look up
+  # @return [Integer, nil] Index of the word, or nil if not found
+  #
+  def lookup(word)
+    @word_to_index[word]
+  end
+
+  # Get word by index
+  #
+  # @param index [Integer] The index to look up
+  # @return [String, nil] Word at the index, or nil if not found
+  #
+  def get_word(index)
+    @index_to_word[index]
+  end
+
+  # Check if word exists in vocabulary
+  #
+  # @param word [String] Word to check
+  # @return [Boolean] True if word exists
+  #
+  def include?(word)
+    @word_to_index.key?(word)
+  end
+
+  # Get vocabulary size
+  #
+  # @return [Integer] Number of words in vocabulary
+  #
+  def size
+    @word_to_index.size
+  end
+
+  # Check if index is valid
+  #
+  # @param index [Integer] Index to check
+  # @return [Boolean] True if index is valid
+  #
+  def valid_index?(index)
+    index.is_a?(Integer) && index >= 0 && index < @word_to_index.size
+  end
+
+  # Get common/most frequent words
+  #
+  # @param n [Integer] Number of words to return
+  # @return [Array<String>] Array of common words
+  #
+  def common_words(n: 10)
+    return [] if @word_to_index.empty?
+
+    @word_to_index.keys.first(n)
+  end
+
+  # Convert to Hash
+  #
+  # @return [Hash{String => Integer}] Copy of word_to_index mapping
+  #
+  def to_h
+    @word_to_index.dup
+  end
+
+  # Get all words as enumerator
+  #
+  # @return [Enumerator<String>] Enumerator of all words
+  #
+  def words
+    @word_to_index.each_key
+  end
+
+  # Load vocabulary from JSON file
+  #
+  # @param path [String] Path to JSON file
+  # @param language_code [String] Language code (auto-detected from filename if nil)
+  # @return [Vocabulary] New vocabulary instance
+  #
+  # @raise [ArgumentError] If file doesn't exist
+  # @raise [Json::ParserError] If file is not valid JSON
+  #
+  def self.from_file(path, language_code: nil)
+    raise ArgumentError, "File not found: #{path}" unless File.exist?(path)
+
+    language_code ||= detect_language_from_path(path)
+
+    data = JSON.parse(File.read(path))
+
+    case data
+    when Hash
+      word_to_index = data.transform_keys(&:freeze).freeze
+    when Array
+      word_to_index = {}
+      data.each_with_index do |word, index|
+        word_to_index[word.freeze] = index
+      end
+      word_to_index.freeze
+    else
+      raise ArgumentError, "Invalid vocabulary format: expected Hash or Array"
+    end
+
+    new(language_code: language_code, word_to_index: word_to_index)
+  end
+
+  # Create vocabulary from Array of words
+  #
+  # @param words [Array<String>] Array of words
+  # @param language_code [String] Language code
+  # @return [Vocabulary] New vocabulary instance
+  #
+  def self.from_words(words, language_code: 'en')
+    word_to_index = {}
+    words.each_with_index do |word, index|
+      word_to_index[word.freeze] = index
+    end
+    word_to_index.freeze
+
+    new(language_code: language_code, word_to_index: word_to_index)
+  end
+
+  # Save vocabulary to JSON file
+  #
+  # @param path [String] Path to save file
+  # @param format [Symbol] Format: :hash or :array
+  #
+  def save_to_file(path, format: :hash)
+    case format
+    when :hash
+      data = @word_to_index.dup
+    when :array
+      max_index = @index_to_word.compact.length
+      data = @index_to_word.compact.first(max_index)
+    else
+      raise ArgumentError, "Unknown format: #{format}"
+    end
+
+    File.write(path, JSON.pretty_generate(data))
+  end
+
+  # Check if vocabulary is empty
+  #
+  # @return [Boolean] True if empty
+  #
+  def empty?
+    @word_to_index.empty?
+  end
+
+  # Get a sample of words
+  #
+  # @param n [Integer] Number of words to sample
+  # @return [Array<String>] Sample of words
+  #
+  def sample(n: 10)
+    @word_to_index.keys.sample(n)
+  end
+
+  # Create a sub-vocabulary containing only specified words
+  #
+  # @param words [Array<String>] Words to include
+  # @return [Vocabulary] New vocabulary with subset of words
+  #
+  def sub_vocabulary(words)
+    filtered = @word_to_index.select { |w, _| words.include?(w) }
+    self.class.new(language_code: @language_code, word_to_index: filtered)
+  end
+
+  # Find words starting with a prefix
+  #
+  # @param prefix [String] Prefix to match
+  # @return [Array<String>] Matching words
+  #
+  def words_starting_with(prefix)
+    pattern = /^#{Regexp.escape(prefix)}/
+    @word_to_index.keys.grep(pattern)
+  end
+
+  # String representation
+  #
+  # @return [String]
+  #
+  def to_s
+    "Vocabulary(language: #{@language_code}, size: #{@word_to_index.size})"
+  end
+  alias inspect to_s
+
+  private_class_method
+
+  # Detect language code from file path
+  #
+  # @param path [String] File path
+  # @return [String] Detected language code
+  #
+  def self.detect_language_from_path(path)
+    basename = File.basename(path)
+
+    if basename =~ /(\w+)\.vocab\.json\z/
+      return $1
+    end
+
+    if basename =~ /\.(\w+)\.vocab\.json\z/
+      return $1
+    end
+
+    'unknown'
+  end
+end

data/lib/kotoshu/embeddings.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+# Load all components
+require_relative 'embeddings/protocol'
+require_relative 'embeddings/lru_cache'
+require_relative 'embeddings/vocabulary'
+require_relative 'embeddings/onnx_runtime_model'
+require_relative 'embeddings/similarity_engine'
+require_relative 'embeddings/search'
+require_relative 'embeddings/embedding_pipeline'
+require_relative 'embeddings/registry'
+
+# Embeddings module for FastText ONNX integration.
+#
+# Provides semantic spell checking using FastText word embeddings.
+# Supports 157 languages through pre-converted ONNX models.
+#
+# @example Simple usage (recommended)
+#   pipeline = Kotoshu::Embeddings.from_cache(language: 'en')
+#   neighbors = pipeline.find_nearest('semantic', k: 5)
+#
+# @example Advanced usage
+#   vocab = Kotoshu::Embeddings::Vocabulary.from_file('vocab.json')
+#   model = Kotoshu::Embeddings::OnnxRuntimeModel.from_file('model.onnx')
+#   engine = Kotoshu::Embeddings::SimilarityEngine.new(pre_normalize: true)
+#
+module Kotoshu
+  module Embeddings
+    # Constants
+    DEFAULT_DIMENSION = 300
+    MAX_VOCABULARY_SIZE = 100_000
+    VERSION = '2.0.0'
+
+    # Expose classes
+    Vocabulary = ::Vocabulary
+    OnnxRuntimeModel = ::OnnxRuntimeModel
+    SimilarityEngine = ::SimilarityEngine
+    Search = ::Search
+    EmbeddingPipeline = ::EmbeddingPipeline
+    LruCache = ::LruCache
+    Registry = ::EmbeddingRegistry
+
+    # Protocols namespace
+    module Protocols
+      EmbeddingModel = ::EmbeddingModelProtocol
+      SimilarityEngine = ::SimilarityEngineProtocol
+      Vocabulary = ::VocabularyProtocol
+    end
+
+    # Create an EmbeddingPipeline from cache
+    #
+    # @param language [String] ISO 639-1 language code
+    # @param preload [Boolean] Preload embeddings into memory
+    # @return [EmbeddingPipeline]
+    #
+    def self.from_cache(language:, preload: false, index: :exact)
+      EmbeddingPipeline.from_cache(language: language, preload: preload, index: index)
+    end
+
+    # Check if a language is supported
+    #
+    # @param language [String] ISO 639-1 language code
+    # @return [Boolean]
+    #
+    def self.language_supported?(language)
+      require_relative '../cache/model_cache'
+      cache = Cache::ModelCache.new
+      cache.available_models_for(language.to_sym).include?(:onnx)
+    end
+
+    # List all supported languages
+    #
+    # @return [Array<String>]
+    #
+    def self.supported_languages
+      require_relative '../cache/model_cache'
+      cache = Cache::ModelCache.new
+      cache.all_available_models[:onnx].keys.map(&:to_s)
+    end
+
+    # Create a custom embedding pipeline
+    #
+    # @param vocabulary [Vocabulary] Vocabulary instance
+    # @param model [EmbeddingModel] Model instance
+    # @param preload [Boolean] Preload embeddings
+    # @return [EmbeddingPipeline]
+    #
+    def self.create_pipeline(vocabulary:, model:, preload: false, pre_normalize: false)
+      EmbeddingPipeline.new(
+        vocabulary: vocabulary,
+        model: model,
+        preload: preload,
+        pre_normalize: pre_normalize
+      )
+    end
+  end
+end