Class: Kotoshu::Embeddings::EmbeddingPipeline

Inherits:
Object
  • Object
show all
Defined in:
lib/kotoshu/embeddings/embedding_pipeline.rb

Instance Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(vocabulary:, model:, preload: false, index: :exact, pre_normalize: false, cache_size: 1000) ⇒ EmbeddingPipeline

Create pipeline with full configuration

Parameters:

  • vocabulary (Vocabulary) —

    Vocabulary instance

  • model (EmbeddingModel) —

    Model instance

  • preload (Boolean) (defaults to: false) —

    Preload embeddings

  • index (:exact, :ann) (defaults to: :exact) —

    Search index type (:exact = brute force, :ann = FAISS/HNSW)

  • pre_normalize (Boolean) (defaults to: false) —

    Pre-normalize vectors

  • cache_size (Integer) (defaults to: 1000) —

    Embedding cache size



97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 97

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

Instance Attribute Details

#model ⇒ EmbeddingModel (readonly)

Returns:

  • (EmbeddingModel) 


29
30
31
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 29

def model
  @model
end

#search ⇒ Search (readonly)

Returns:



35
36
37
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 35

def search
  @search
end

#similarity_engine ⇒ SimilarityEngine (readonly)

Returns:



32
33
34
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 32

def similarity_engine
  @similarity_engine
end

#vocabulary ⇒ Vocabulary (readonly)

Returns:



26
27
28
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 26

def vocabulary
  @vocabulary
end

Class Method Details

.from_cache(language:, cache: nil, preload: false, index: :exact) ⇒ EmbeddingPipeline Also known as: []

Create pipeline from cache (one-line initialization)

Parameters:

  • language (String) —

    ISO 639-1 language code

  • cache (Cache::ModelCache) (defaults to: nil) —

    Cache instance

  • preload (Boolean) (defaults to: false) —

    Preload embeddings into memory

  • index (:exact, :auto) (defaults to: :exact) —

    Search index type

Returns:

Raises:

  • (ArgumentError) —

    If no cached model found for language



47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 47

def self.from_cache(language:, cache: nil, preload: false, index: :exact)
  cache ||= Kotoshu::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

.from_files(vocab_path:, model_path:, language:, preload: false, index: :exact) ⇒ EmbeddingPipeline

Create pipeline from files

Parameters:

  • vocab_path (String) —

    Path to vocabulary JSON file

  • model_path (String) —

    Path to ONNX model file

  • language (String) —

    Language code

  • preload (Boolean) (defaults to: false) —

    Preload embeddings

  • index (:exact, :auto) (defaults to: :exact) —

    Search index type

Returns:



76
77
78
79
80
81
82
83
84
85
86
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 76

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

Instance Method Details

#find_nearest(word, k: 10, exclude_self: true, min_similarity: 0.0) ⇒ Array<Hash>

Find k nearest neighbors for a word

Parameters:

  • word (String) —

    Query word

  • k (Integer) (defaults to: 10) —

    Number of neighbors

  • exclude_self (Boolean) (defaults to: true) —

    Exclude query word

  • min_similarity (Float) (defaults to: 0.0) —

    Minimum similarity threshold

Returns:

  • (Array<Hash>) —

    Array of similarity, index



122
123
124
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 122

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_batch(words, k: 10) ⇒ Hash<String, Array<Hash>>

Find nearest neighbors for multiple words

Parameters:

  • words (Array<String>) —

    Query words

  • k (Integer) (defaults to: 10) —

    Neighbors per word

Returns:

  • (Hash<String, Array<Hash>>) 


132
133
134
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 132

def find_nearest_batch(words, k: 10)
  @search.find_nearest_batch(words, k: k)
end

#get_embedding(word) ⇒ Array<Float>?

Get embedding for a word

Parameters:

  • word (String) —

    Word

Returns:

  • (Array<Float>, nil) 


151
152
153
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 151

def get_embedding(word)
  @model.get_embedding_for_word(word, @vocabulary)
end

#get_embedding_by_index(index) ⇒ Array<Float>?

Get embedding by index

Parameters:

  • index (Integer) —

    Word index

Returns:

  • (Array<Float>, nil) 


160
161
162
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 160

def get_embedding_by_index(index)
  @model.get_embedding(index)
end

#include?(word) ⇒ Boolean

Check if word exists in vocabulary

Parameters:

  • word (String) —

    Word

Returns:

  • (Boolean) 


169
170
171
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 169

def include?(word)
  @vocabulary.include?(word)
end

#model_info ⇒ Hash

Get model information

Returns:

  • (Hash) 


212
213
214
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 212

def model_info
  @model.model_info
end

#preload_embeddings! ⇒ self

Preload all embeddings into memory

Returns:

  • (self) 


177
178
179
180
181
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 177

def preload_embeddings!
  @model.load!
  @search.preload_embeddings!
  self
end

#similarity(word1, word2) ⇒ Float?

Compute similarity between two words

Parameters:

  • word1 (String) —

    First word

  • word2 (String) —

    Second word

Returns:

  • (Float, nil) —

    Similarity or nil if either word not found



142
143
144
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 142

def similarity(word1, word2)
  @search.similarity(word1, word2)
end

#stats ⇒ Hash

Get pipeline statistics

Returns:

  • (Hash) 


197
198
199
200
201
202
203
204
205
206
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 197

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.cache_size
  }
end

#to_s ⇒ String Also known as: inspect

String representation

Returns:

  • (String) 


220
221
222
223
224
225
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 220

def to_s
  "EmbeddingPipeline(language: #{@vocabulary.language_code}, " \
    "vocab_size: #{@vocabulary.size}, " \
    "dimension: #{@model.dimension}, " \
    "loaded: #{@model.loaded?})"
end

#unload! ⇒ self

Unload model from memory

Returns:

  • (self) 


187
188
189
190
191
 
# File 'lib/kotoshu/embeddings/embedding_pipeline.rb', line 187

def unload!
  @model.unload!
  @search.clear_cache
  self
end