Class: Kotoshu::Embeddings::Search

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

Defined Under Namespace

Classes: MinHeap

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(vocabulary:, model:, similarity_engine:, pre_normalize: false) ⇒ Search

Create a new exact search

Parameters:

  • vocabulary (Vocabulary) —

    Word vocabulary

  • model (EmbeddingModel) —

    Embedding provider

  • similarity_engine (SimilarityEngine) —

    Similarity calculator

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

    Pre-normalize vectors for faster similarity



68
69
70
71
72
73
74
75
76
 
# File 'lib/kotoshu/embeddings/search.rb', line 68

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

Instance Attribute Details

#embeddings_loaded ⇒ Boolean (readonly)

Returns Whether embeddings are preloaded.

Returns:

  • (Boolean) —

    Whether embeddings are preloaded



59
60
61
 
# File 'lib/kotoshu/embeddings/search.rb', line 59

def embeddings_loaded
  @embeddings_loaded
end

#model ⇒ EmbeddingModel (readonly)

Returns:

  • (EmbeddingModel) 


53
54
55
 
# File 'lib/kotoshu/embeddings/search.rb', line 53

def model
  @model
end

#similarity_engine ⇒ SimilarityEngine (readonly)

Returns:



56
57
58
 
# File 'lib/kotoshu/embeddings/search.rb', line 56

def similarity_engine
  @similarity_engine
end

#vocabulary ⇒ Vocabulary (readonly)

Returns:



50
51
52
 
# File 'lib/kotoshu/embeddings/search.rb', line 50

def vocabulary
  @vocabulary
end

Instance Method Details

#cache_size ⇒ Integer

Number of cached embeddings

Returns:

  • (Integer) 


165
166
167
 
# File 'lib/kotoshu/embeddings/search.rb', line 165

def cache_size
  @embedding_cache.size
end

#clear_cache ⇒ self

Clear embedding cache

Returns:

  • (self) 


155
156
157
158
159
 
# File 'lib/kotoshu/embeddings/search.rb', line 155

def clear_cache
  @embedding_cache.clear
  @embeddings_loaded = false
  self
end

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

Find k nearest neighbors for a word

Parameters:

  • query_word (String) —

    Query word

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

    Number of neighbors to return

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

    Exclude query word from results

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

    Minimum similarity threshold

Returns:

  • (Array<Hash>) —

    Array of similarity, index



86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
 
# File 'lib/kotoshu/embeddings/search.rb', line 86

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

Find nearest neighbors for multiple words

Parameters:

  • query_words (Array<String>) —

    Query words

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

    Number of neighbors per word

Returns:

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

    Word to results mapping



115
116
117
118
119
 
# File 'lib/kotoshu/embeddings/search.rb', line 115

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

#preload_embeddings! ⇒ self

Preload all embeddings into memory

Returns:

  • (self) 


139
140
141
142
143
144
145
146
147
148
149
 
# File 'lib/kotoshu/embeddings/search.rb', line 139

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

#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



127
128
129
130
131
132
133
 
# File 'lib/kotoshu/embeddings/search.rb', line 127

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

#to_s ⇒ String Also known as: inspect

String representation

Returns:

  • (String) 


173
174
175
 
# File 'lib/kotoshu/embeddings/search.rb', line 173

def to_s
  "ExactSearch(vocab: #{@vocabulary.size}, loaded: #{@embeddings_loaded})"
end