kotoshu 0.3.0

data/lib/kotoshu/embeddings/search.rb

@@ -0,0 +1,192 @@
+# frozen_string_literal: true
+
+# Search - Brute force nearest neighbor search
+#
+# Performs exhaustive search over all vocabulary entries.
+# Uses min-heap for efficient top-k selection (O(n log k) instead of O(n log n)).
+#
+# @example
+#   search = ExactSearch.new(
+#     vocabulary: vocab,
+#     model: model,
+#     similarity_engine: engine
+#   )
+#   neighbors = search.find_nearest('hello', k: 5)
+#
+class Search
+  # Min-heap for top-k selection
+  class MinHeap
+    def initialize(max_size)
+      @max_size = max_size
+      @heap = []
+    end
+
+    def push(item)
+      @heap << item
+      @heap.sort_by! { |i| i[:similarity] }
+      @heap.shift if @heap.size > @max_size
+    end
+
+    def empty?
+      @heap.empty?
+    end
+
+    def size
+      @heap.size
+    end
+
+    def each(&block)
+      @heap.each(&block)
+    end
+
+    def to_a
+      @heap.dup
+    end
+  end
+
+  # @return [Vocabulary]
+  attr_reader :vocabulary
+
+  # @return [EmbeddingModel]
+  attr_reader :model
+
+  # @return [SimilarityEngine]
+  attr_reader :similarity_engine
+
+  # @return [Boolean] Whether embeddings are preloaded
+  attr_reader :embeddings_loaded
+
+  # Create a new exact search
+  #
+  # @param vocabulary [Vocabulary] Word vocabulary
+  # @param model [EmbeddingModel] Embedding provider
+  # @param similarity_engine [SimilarityEngine] Similarity calculator
+  # @param pre_normalize [Boolean] Pre-normalize vectors for faster similarity
+  #
+  def initialize(vocabulary:, model:, similarity_engine:, pre_normalize: false)
+    @vocabulary = vocabulary
+    @model = model
+    @similarity_engine = similarity_engine
+    @pre_normalize = pre_normalize
+
+    @embedding_cache = {}
+    @embeddings_loaded = false
+  end
+
+  # Find k nearest neighbors for a word
+  #
+  # @param query_word [String] Query word
+  # @param k [Integer] Number of neighbors to return
+  # @param exclude_self [Boolean] Exclude query word from results
+  # @param min_similarity [Float] Minimum similarity threshold
+  # @return [Array<Hash>] Array of {word, similarity, index}
+  #
+  def find_nearest(query_word, k: 10, exclude_self: true, min_similarity: 0.0)
+    query_vec = get_embedding_for_word(query_word)
+    return [] unless query_vec
+
+    heap = MinHeap.new(k)
+
+    @vocabulary.words.each do |word|
+      next if exclude_self && word == query_word
+
+      vec = get_embedding_for_word(word)
+      next unless vec
+
+      similarity = @similarity_engine.cosine(query_vec, vec)
+      next if similarity < min_similarity
+
+      index = @vocabulary.lookup(word)
+      heap.push(word: word, similarity: similarity, index: index)
+    end
+
+    # Return sorted by similarity descending
+    heap.to_a.sort_by { |r| -r[:similarity] }
+  end
+
+  # Find 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 results mapping
+  #
+  def find_nearest_batch(query_words, k: 10)
+    query_words.each_with_object({}) do |word, results|
+      results[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, nil] Similarity or nil if either word not found
+  #
+  def similarity(word1, word2)
+    vec1 = get_embedding_for_word(word1)
+    vec2 = get_embedding_for_word(word2)
+    return nil unless vec1 && vec2
+
+    @similarity_engine.cosine(vec1, vec2)
+  end
+
+  # Preload all embeddings into memory
+  #
+  # @return [self]
+  #
+  def preload_embeddings!
+    all_indices = (0...@vocabulary.size).to_a
+    embeddings = @model.get_embeddings(all_indices)
+
+    @vocabulary.words.each_with_index do |word, i|
+      @embedding_cache[word] = embeddings[i]
+    end
+
+    @embeddings_loaded = true
+    self
+  end
+
+  # Clear embedding cache
+  #
+  # @return [self]
+  #
+  def clear_cache
+    @embedding_cache.clear
+    @embeddings_loaded = false
+    self
+  end
+
+  # String representation
+  #
+  # @return [String]
+  #
+  def to_s
+    "ExactSearch(vocab: #{@vocabulary.size}, loaded: #{@embeddings_loaded})"
+  end
+  alias inspect to_s
+
+  private
+
+  # Get embedding for a word (with caching)
+  #
+  # @param word [String] Word
+  # @return [Array<Float>, nil]
+  #
+  def get_embedding_for_word(word)
+    # Check cache first
+    if @embedding_cache.key?(word)
+      return @embedding_cache[word]
+    end
+
+    index = @vocabulary.lookup(word)
+    return nil unless index
+
+    vec = @model.get_embedding(index)
+    return nil unless vec
+
+    # Cache if not preloaded (to avoid repeated lookups)
+    @embedding_cache[word] = vec unless @embeddings_loaded
+
+    vec
+  end
+end

data/lib/kotoshu/embeddings/similarity_engine.rb

@@ -0,0 +1,248 @@
+# frozen_string_literal: true
+
+require_relative 'protocol'
+
+# SimilarityEngine - Compute similarity between embedding vectors
+#
+# Provides various similarity/distance metrics with optimizations like
+# norm caching and pre-normalized vector support.
+#
+# @example Basic usage
+#   engine = SimilarityEngine.new
+#   engine.cosine([1.0, 0.0], [1.0, 0.0])  # => 1.0
+#
+# @example Pre-normalized vectors (faster)
+#   engine = SimilarityEngine.new(pre_normalize: true)
+#   engine.pre_normalize([1.0, 0.0])  # => [1.0, 0.0]
+#
+class SimilarityEngine
+  include SimilarityEngineProtocol
+
+  # Default embedding dimension for norm cache initialization
+  DEFAULT_CACHE_SIZE = 10_000
+
+  # @return [Boolean] Whether vectors are pre-normalized
+  attr_reader :pre_normalize
+
+  # @return [Integer] Number of cache hits
+  attr_reader :cache_hits
+
+  # @return [Integer] Number of cache misses
+  attr_reader :cache_misses
+
+  # Create a new similarity engine
+  #
+  # @param pre_normalize [Boolean] Whether to pre-normalize vectors
+  # @param cache_norms [Boolean] Whether to cache vector norms
+  #
+  def initialize(pre_normalize: false, cache_norms: true)
+    @pre_normalize = pre_normalize
+    @cache_norms = cache_norms
+    @norm_cache = cache_norms ? {} : nil
+    @cache_hits = 0
+    @cache_misses = 0
+  end
+
+  # Compute cosine similarity between two vectors
+  #
+  # Cosine similarity = dot(v1, v2) / (||v1|| * ||v2||)
+  # Range: -1.0 (opposite) to 1.0 (identical)
+  #
+  # @param vec1 [Array<Float>] First vector
+  # @param vec2 [Array<Float>] Second vector
+  # @return [Float] Cosine similarity, or 0.0 if either vector is nil/empty
+  #
+  def cosine(vec1, vec2)
+    return 0.0 if vec1.nil? || vec2.nil? || vec1.empty? || vec2.empty?
+
+    norm1 = get_norm(vec1)
+    norm2 = get_norm(vec2)
+
+    return 0.0 if norm1.zero? || norm2.zero?
+
+    dot = dot_product(vec1, vec2)
+    dot / (norm1 * norm2)
+  end
+
+  # Compute dot product between two vectors
+  #
+  # @param vec1 [Array<Float>] First vector
+  # @param vec2 [Array<Float>] Second vector
+  # @return [Float] Dot product
+  #
+  def dot_product(vec1, vec2)
+    return 0.0 if vec1.nil? || vec2.nil? || vec1.empty? || vec2.empty?
+
+    vec1.zip(vec2).sum { |a, b| a * b }
+  end
+
+  # Compute Euclidean distance between two vectors
+  #
+  # @param vec1 [Array<Float>] First vector
+  # @param vec2 [Array<Float>] Second vector
+  # @return [Float] Euclidean distance
+  #
+  def euclidean(vec1, vec2)
+    return 0.0 if vec1.nil? || vec2.nil? || vec1.empty? || vec2.empty?
+    return 0.0 if vec1.equal?(vec2)
+
+    sum = 0.0
+    vec1.zip(vec2) do |a, b|
+      diff = a - b
+      sum += diff * diff
+    end
+    Math.sqrt(sum)
+  end
+
+  # Compute Manhattan (L1) distance between two vectors
+  #
+  # @param vec1 [Array<Float>] First vector
+  # @param vec2 [Array<Float>] Second vector
+  # @return [Float] Manhattan distance
+  #
+  def manhattan(vec1, vec2)
+    return 0.0 if vec1.nil? || vec2.nil? || vec1.empty? || vec2.empty?
+
+    vec1.zip(vec2).sum { |a, b| (a - b).abs }
+  end
+
+  # Pre-normalize a vector to unit length
+  #
+  # @param vec [Array<Float>] Vector to normalize
+  # @return [Array<Float>] Normalized vector
+  #
+  def pre_normalize(vec)
+    return vec.dup if vec.nil? || vec.empty?
+
+    norm = get_norm(vec)
+    return vec.dup if norm.zero?
+
+    vec.map { |x| x / norm }
+  end
+
+  # Normalize and compute similarity in one pass
+  #
+  # For pre-normalized vectors, this is just dot product (much faster).
+  #
+  # @param vec1 [Array<Float>] First vector
+  # @param vec2 [Array<Float>] Second vector
+  # @return [Float] Cosine similarity
+  #
+  def normalize_and_compute(vec1, vec2)
+    return 0.0 if vec1.nil? || vec2.nil? || vec1.empty? || vec2.empty?
+
+    if @pre_normalize
+      # For normalized vectors, cosine similarity = dot product
+      dot_product(vec1, vec2)
+    else
+      cosine(vec1, vec2)
+    end
+  end
+
+  # Check if vectors are normalized (unit length)
+  #
+  # @param vec [Array<Float>] Vector to check
+  # @return [Boolean] True if vector is normalized
+  #
+  def is_normalized?(vec)
+    return true if vec.nil? || vec.empty?
+
+    norm = get_norm(vec)
+    (norm - 1.0).abs < Float::EPSILON * 10
+  end
+
+  # Check if normalization is required for accurate similarity
+  #
+  # @return [Boolean] True if normalization should be applied
+  #
+  def normalization_required?
+    !@pre_normalize
+  end
+
+  # Clear the norm cache
+  #
+  # @return [self]
+  #
+  def clear_cache
+    @norm_cache&.clear
+    @cache_hits = 0
+    @cache_misses = 0
+    self
+  end
+
+  # Get cache statistics
+  #
+  # @return [Hash] Cache statistics
+  #
+  def cache_stats
+    total = @cache_hits + @cache_misses
+    {
+      hits: @cache_hits,
+      misses: @cache_misses,
+      hit_rate: total.zero? ? 0.0 : @cache_hits.to_f / total,
+      cache_size: @norm_cache&.size || 0
+    }
+  end
+
+  # Compute similarity for a batch of vector pairs
+  #
+  # More efficient than calling cosine() repeatedly.
+  #
+  # @param pairs [Array<Array<Array<Float>>>] Array of [vec1, vec2] pairs
+  # @return [Array<Float>] Array of similarities
+  #
+  def cosine_batch(pairs)
+    pairs.map { |v1, v2| cosine(v1, v2) }
+  end
+
+  # Compute all pairwise similarities for a set of vectors
+  #
+  # @param vectors [Array<Array<Float>>>] Array of vectors
+  # @return [Array<Array<Float>>] Similarity matrix
+  #
+  def compute_all_pairs(vectors)
+    n = vectors.length
+    matrix = Array.new(n) { Array.new(n, 0.0) }
+
+    (0...n).each do |i|
+      matrix[i][i] = 1.0
+      ((i + 1)...n).each do |j|
+        sim = cosine(vectors[i], vectors[j])
+        matrix[i][j] = sim
+        matrix[j][i] = sim
+      end
+    end
+
+    matrix
+  end
+
+  private
+
+  # Get norm with caching
+  #
+  # @param vec [Array<Float>] Vector
+  # @return [Float] Vector norm (magnitude)
+  #
+  def get_norm(vec)
+    return 0.0 if vec.nil? || vec.empty?
+
+    if @norm_cache && @norm_cache.key?(vec.object_id)
+      @cache_hits += 1
+      return @norm_cache[vec.object_id]
+    end
+
+    @cache_misses += 1 if @norm_cache
+
+    norm = Math.sqrt(vec.sum { |x| x * x })
+
+    if @norm_cache
+      # Avoid memory leaks by limiting cache size
+      if @norm_cache.size >= 100_000
+        @norm_cache.shift
+      end
+      @norm_cache[vec.object_id] = norm
+    end
+
+    norm
+  end
+end