kotoshu 0.3.0

data/lib/kotoshu/embeddings/embedding_pipeline.rb

@@ -0,0 +1,244 @@
+# frozen_string_literal: true
+
+require_relative 'vocabulary'
+require_relative 'onnx_runtime_model'
+require_relative 'similarity_engine'
+require_relative 'search'
+require_relative 'lru_cache'
+
+# EmbeddingPipeline - Unified API for embedding-based similarity search
+#
+# Provides a simple, unified interface for loading vocabulary and models,
+# and performing similarity search. This is the recommended entry point.
+#
+# @example Simple usage (one line)
+#   pipeline = EmbeddingPipeline.from_cache(language: 'en')
+#
+# @example Full configuration
+#   pipeline = EmbeddingPipeline.new(
+#     vocabulary: vocab,
+#     model: model,
+#     preload: true
+#   )
+#
+# @example Finding similar words
+#   neighbors = pipeline.find_nearest('semantic', k: 5)
+#   neighbors.each { |r| puts "#{r[:word]}: #{r[:similarity].round(4)}" }
+#
+class EmbeddingPipeline
+  # @return [Vocabulary]
+  attr_reader :vocabulary
+
+  # @return [EmbeddingModel]
+  attr_reader :model
+
+  # @return [SimilarityEngine]
+  attr_reader :similarity_engine
+
+  # @return [Search]
+  attr_reader :search
+
+  # Create pipeline from cache (one-line initialization)
+  #
+  # @param language [String] ISO 639-1 language code
+  # @param cache [Cache::ModelCache] Cache instance
+  # @param preload [Boolean] Preload embeddings into memory
+  # @param index [:exact, :auto] Search index type
+  # @return [EmbeddingPipeline]
+  #
+  # @raise [ArgumentError] If no cached model found for language
+  #
+  def self.from_cache(language:, cache: nil, preload: false, index: :exact)
+    require_relative 'cache/model_cache'
+
+    cache ||= Cache::ModelCache.new
+
+    vocab_path = cache.find_vocab(language)
+    model_path = cache.find_model(language, :onnx)
+
+    unless vocab_path && model_path
+      raise ArgumentError, "No cached model for language: #{language}. " \
+                           "Run: ruby scripts/extract_vocabularies.rb --languages=#{language}"
+    end
+
+    from_files(
+      vocab_path: vocab_path,
+      model_path: model_path,
+      language: language,
+      preload: preload,
+      index: index
+    )
+  end
+
+  # Create pipeline from files
+  #
+  # @param vocab_path [String] Path to vocabulary JSON file
+  # @param model_path [String] Path to ONNX model file
+  # @param language [String] Language code
+  # @param preload [Boolean] Preload embeddings
+  # @param index [:exact, :auto] Search index type
+  # @return [EmbeddingPipeline]
+  #
+  def self.from_files(vocab_path:, model_path:, language:, preload: false, index: :exact)
+    vocab = Vocabulary.from_file(vocab_path, language_code: language)
+    model = OnnxRuntimeModel.from_file(model_path, language_code: language)
+
+    new(
+      vocabulary: vocab,
+      model: model,
+      preload: preload,
+      index: index
+    )
+  end
+
+  # Create pipeline with full configuration
+  #
+  # @param vocabulary [Vocabulary] Vocabulary instance
+  # @param model [EmbeddingModel] Model instance
+  # @param preload [Boolean] Preload embeddings
+  # @param index [:exact, :ann] Search index type (:exact = brute force, :ann = FAISS/HNSW)
+  # @param pre_normalize [Boolean] Pre-normalize vectors
+  # @param cache_size [Integer] Embedding cache size
+  #
+  def initialize(vocabulary:, model:, preload: false, index: :exact, pre_normalize: false, cache_size: 1000)
+    @vocabulary = vocabulary
+    @model = model
+    @similarity_engine = SimilarityEngine.new(pre_normalize: pre_normalize)
+    @cache_size = cache_size
+
+    # Create search engine
+    @search = Search.new(
+      vocabulary: vocabulary,
+      model: model,
+      similarity_engine: @similarity_engine,
+      pre_normalize: pre_normalize
+    )
+
+    preload_embeddings! if preload
+  end
+
+  # Find k nearest neighbors for a word
+  #
+  # @param word [String] Query word
+  # @param k [Integer] Number of neighbors
+  # @param exclude_self [Boolean] Exclude query word
+  # @param min_similarity [Float] Minimum similarity threshold
+  # @return [Array<Hash>] Array of {word, similarity, index}
+  #
+  def find_nearest(word, k: 10, exclude_self: true, min_similarity: 0.0)
+    @search.find_nearest(word, k: k, exclude_self: exclude_self, min_similarity: min_similarity)
+  end
+
+  # Find nearest neighbors for multiple words
+  #
+  # @param words [Array<String>] Query words
+  # @param k [Integer] Neighbors per word
+  # @return [Hash<String, Array<Hash>>]
+  #
+  def find_nearest_batch(words, k: 10)
+    @search.find_nearest_batch(words, k: k)
+  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)
+    @search.similarity(word1, word2)
+  end
+
+  # Get embedding for a word
+  #
+  # @param word [String] Word
+  # @return [Array<Float>, nil]
+  #
+  def get_embedding(word)
+    @model.get_embedding_for_word(word, @vocabulary)
+  end
+
+  # Get embedding by index
+  #
+  # @param index [Integer] Word index
+  # @return [Array<Float>, nil]
+  #
+  def get_embedding_by_index(index)
+    @model.get_embedding(index)
+  end
+
+  # Check if word exists in vocabulary
+  #
+  # @param word [String] Word
+  # @return [Boolean]
+  #
+  def include?(word)
+    @vocabulary.include?(word)
+  end
+
+  # Preload all embeddings into memory
+  #
+  # @return [self]
+  #
+  def preload_embeddings!
+    @model.load!
+    @search.preload_embeddings!
+    self
+  end
+
+  # Unload model from memory
+  #
+  # @return [self]
+  #
+  def unload!
+    @model.unload!
+    @search.clear_cache
+    self
+  end
+
+  # Get pipeline statistics
+  #
+  # @return [Hash]
+  #
+  def stats
+    {
+      language: @vocabulary.language_code,
+      vocabulary_size: @vocabulary.size,
+      embedding_dimension: @model.dimension,
+      model_loaded: @model.loaded?,
+      embeddings_preloaded: @search.embeddings_loaded,
+      cache_stats: @search.instance_variable_get(:@embedding_cache)&.stats
+    }
+  end
+
+  # Get model information
+  #
+  # @return [Hash]
+  #
+  def model_info
+    @model.model_info
+  end
+
+  # String representation
+  #
+  # @return [String]
+  #
+  def to_s
+    "EmbeddingPipeline(language: #{@vocabulary.language_code}, " \
+      "vocab_size: #{@vocabulary.size}, " \
+      "dimension: #{@model.dimension}, " \
+      "loaded: #{@model.loaded?})"
+  end
+  alias inspect to_s
+
+  # Convenience class methods
+  class << self
+    # Create pipeline for a specific language (shortcut)
+    #
+    # @param language [String] ISO 639-1 language code
+    # @param kwargs [Hash] Additional options
+    # @return [EmbeddingPipeline]
+    #
+    alias :[] :from_cache
+  end
+end

data/lib/kotoshu/embeddings/lru_cache.rb

@@ -0,0 +1,233 @@
+# frozen_string_literal: true
+
+# LruCache - Least Recently Used Cache
+#
+# Provides efficient O(1) LRU caching with optional TTL support.
+# Used for caching embeddings during similarity search.
+#
+# @example Basic usage
+#   cache = LruCache.new(max_size: 1000)
+#   cache[:key] = value
+#   cache[:key]  # => value
+#
+# @example With TTL
+#   cache = LruCache.new(max_size: 1000, ttl: 300)  # 5 minutes
+#
+class LruCache
+  # @return [Integer] Maximum number of entries
+  attr_reader :max_size
+
+  # @return [Integer, nil] TTL in seconds
+  attr_reader :ttl
+
+  # @return [Integer] Number of cache hits
+  attr_reader :hits
+
+  # @return [Integer] Number of cache misses
+  attr_reader :misses
+
+  # Create a new LRU cache
+  #
+  # @param max_size [Integer] Maximum number of entries (default: 1000)
+  # @param ttl [Integer, nil] Time-to-live in seconds (default: nil = no expiry)
+  #
+  def initialize(max_size: 1000, ttl: nil)
+    @max_size = max_size
+    @ttl = ttl
+    @cache = {}  # key -> {value: v, accessed_at: t, created_at: t}
+    @order = []  # Ordered list of keys (most recently used first)
+    @hits = 0
+    @misses = 0
+  end
+
+  # Get value for key
+  #
+  # @param key [Object] Cache key
+  # @return [Object, nil] Cached value or nil if not found/expired
+  #
+  def [](key)
+    entry = @cache[key]
+    return nil unless entry
+
+    # Check TTL
+    if @ttl && (Time.now - entry[:created_at]) > @ttl
+      delete(key)
+      @misses += 1
+      return nil
+    end
+
+    # Update access order (move to front = most recently used)
+    @order.delete(key)
+    @order.unshift(key)
+    entry[:accessed_at] = Time.now
+
+    @hits += 1
+    entry[:value]
+  end
+
+  # Set value for key
+  #
+  # @param key [Object] Cache key
+  # @param value [Object] Value to cache
+  # @return [Object] The value
+  #
+  def []=(key, value)
+    # Evict LRU if at capacity
+    if @cache.key?(key)
+      # Update existing entry
+      @cache[key][:value] = value
+      @cache[key][:accessed_at] = Time.now
+      # Move to front
+      @order.delete(key)
+      @order.unshift(key)
+      return value
+    end
+
+    if @cache.size >= @max_size
+      evict_lru
+    end
+
+    @cache[key] = {
+      value: value,
+      accessed_at: Time.now,
+      created_at: Time.now
+    }
+    @order.unshift(key)
+
+    value
+  end
+
+  # Check if key exists
+  #
+  # @param key [Object] Cache key
+  # @return [Boolean] True if key exists and not expired
+  #
+  def key?(key)
+    entry = @cache[key]
+    return false unless entry
+
+    if @ttl && (Time.now - entry[:created_at]) > @ttl
+      delete(key)
+      return false
+    end
+
+    true
+  end
+
+  # Delete key from cache
+  #
+  # @param key [Object] Cache key
+  # @return [Object, nil] Deleted value or nil
+  #
+  def delete(key)
+    entry = @cache.delete(key)
+    @order.delete(key)
+    entry&.[](:value)
+  end
+
+  # Clear all entries
+  #
+  # @return [self]
+  #
+  def clear
+    @cache.clear
+    @order.clear
+    self
+  end
+
+  # Get current size
+  #
+  # @return [Integer] Number of entries
+  #
+  def size
+    @cache.size
+  end
+
+  # Check if empty
+  #
+  # @return [Boolean]
+  #
+  def empty?
+    @cache.empty?
+  end
+
+  # Get least recently used key-value pair
+  #
+  # @return [Array<Object, Object>, nil]
+  #
+  def lru
+    return nil if @order.empty?
+
+    key = @order.last
+    [key, @cache[key][:value]]
+  end
+
+  # Get most recently used key-value pair
+  #
+  # @return [Array<Object, Object>, nil]
+  #
+  def mru
+    return nil if @order.empty?
+
+    key = @order.first
+    [key, @cache[key][:value]]
+  end
+
+  # Get all keys
+  #
+  # @return [Array<Object>] Array of keys
+  #
+  def keys
+    @order.dup
+  end
+
+  # Get all values
+  #
+  # @return [Array<Object>] Array of values
+  #
+  def values
+    @order.map { |key| @cache[key][:value] }
+  end
+
+  # Get cache statistics
+  #
+  # @return [Hash] Statistics
+  #
+  def stats
+    total = @hits + @misses
+    {
+      size: size,
+      max_size: @max_size,
+      hits: @hits,
+      misses: @misses,
+      hit_rate: total.zero? ? 0.0 : @hits.to_f / total,
+      ttl: @ttl
+    }
+  end
+
+  # Fetch with block (cache-aside pattern)
+  #
+  # @param key [Object] Cache key
+  # @return [Object] Cached value or block result
+  #
+  def fetch(key, &block)
+    result = self[key]
+    return result if result || key?(key)
+
+    value = block.call
+    self[key] = value
+    value
+  end
+
+  private
+
+  # Evict least recently used entry
+  #
+  def evict_lru
+    return if @order.empty?
+
+    lru_key = @order.last
+    @cache.delete(lru_key)
+    @order.pop
+  end
+end