Class: Kotoshu::Embeddings::OnnxRuntimeModel

Inherits:
Object
  • Object
show all
Includes:
EmbeddingModelProtocol
Defined in:
lib/kotoshu/embeddings/onnx_runtime_model.rb

Constant Summary collapse

DEFAULT_DIMENSION =

Default dimension for FastText models

300
BATCH_SIZE =

Batch size for batch inference

32

Instance Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Methods included from Protocol

#assert_implemented_by!, #compliance_errors, #optional, #optional_methods, #required, #required_methods

Constructor Details

#initialize(language_code:, onnx_path:, dimension: DEFAULT_DIMENSION) ⇒ OnnxRuntimeModel

Create a new ONNX Runtime model

Parameters:

  • language_code (String) —

    ISO 639-1 language code

  • onnx_path (String) —

    Path to .onnx file

  • dimension (Integer) (defaults to: DEFAULT_DIMENSION) —

    Embedding dimension (default: 300)



53
54
55
56
57
58
59
60
61
62
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 53

def initialize(language_code:, onnx_path:, dimension: DEFAULT_DIMENSION)
  @language_code = language_code
  @onnx_path = onnx_path
  @dimension = dimension
  @session = nil
  @loaded = false
  @input_name = nil
  @output_name = nil
  @inference_count = 0
end

Instance Attribute Details

#dimension ⇒ Integer (readonly)

Returns Embedding dimension.

Returns:

  • (Integer) —

    Embedding dimension



34
35
36
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 34

def dimension
  @dimension
end

#inference_count ⇒ Integer (readonly)

Returns Number of inference calls.

Returns:

  • (Integer) —

    Number of inference calls



45
46
47
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 45

def inference_count
  @inference_count
end

#language_code ⇒ String (readonly)

Returns Language code (ISO 639-1).

Returns:

  • (String) —

    Language code (ISO 639-1)



31
32
33
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 31

def language_code
  @language_code
end

#onnx_path ⇒ String (readonly)

Returns Path to ONNX model file.

Returns:

  • (String) —

    Path to ONNX model file



37
38
39
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 37

def onnx_path
  @onnx_path
end

Class Method Details

.detect_language_from_path(path) ⇒ String

Detect language from file path

Parameters:

  • path (String)

Returns:

  • (String) 


390
391
392
393
394
395
396
397
398
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 390

def self.detect_language_from_path(path)
  basename = File.basename(path)

  if basename =~ /\.([a-z]{2})\./i
    Regexp.last_match(1).downcase
  else
    'en'
  end
end

.from_cache(language_code, cache: nil) ⇒ OnnxRuntimeModel?

Create model from cache

Parameters:

  • language_code (String) —

    ISO 639-1 language code

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

    Cache instance

Returns:



270
271
272
273
274
275
276
277
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 270

def self.from_cache(language_code, cache: nil)
  cache ||= Kotoshu::Cache::ModelCache.new

  onnx_path = cache.get_onnx_model(language_code)
  return nil unless onnx_path

  from_file(onnx_path, language_code: language_code)
end

.from_file(onnx_path, language_code: nil, dimension: nil) ⇒ OnnxRuntimeModel

Create model from file

Parameters:

  • onnx_path (String) —

    Path to .onnx file

  • language_code (String) (defaults to: nil) —

    Language code (auto-detected if nil)

  • dimension (Integer) (defaults to: nil) —

    Embedding dimension

Returns:



248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 248

def self.from_file(onnx_path, language_code: nil, dimension: nil)
  unless File.exist?(onnx_path)
    raise ArgumentError,
          "ONNX file not found: #{onnx_path}"
  end

  language_code ||= detect_language_from_path(onnx_path)
  dimension ||= DEFAULT_DIMENSION

  new(
    language_code: language_code,
    onnx_path: onnx_path,
    dimension: dimension,
  )
end

Instance Method Details

#batch_size ⇒ Integer

Get batch size for batch inference

Returns:

  • (Integer) 


214
215
216
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 214

def batch_size
  BATCH_SIZE
end

#get_embedding(index) ⇒ Array<Float>

Get embedding for a single word index

Parameters:

  • index (Integer) —

    Word index in vocabulary

Returns:

  • (Array<Float>) —

    Embedding vector

Raises:

  • (RuntimeError) —

    if model is not loaded

  • (ArgumentError) —

    if index is invalid



119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 119

def get_embedding(index)
  ensure_loaded

  unless valid_index?(index)
    raise ArgumentError,
          "Invalid word index: #{index}"
  end

  output = @session.run(
    [@output_name],
    { @input_name => [index] },
  )

  @inference_count += 1

  extract_embedding(output.first)
end

#get_embedding_for_word(word, vocabulary) ⇒ Array<Float>?

Get embedding for a word using vocabulary

Parameters:

  • word (String) —

    The word to lookup

  • vocabulary (Vocabulary) —

    Vocabulary for word-to-index mapping

Returns:

  • (Array<Float>, nil) —

    Embedding vector or nil if word not found



180
181
182
183
184
185
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 180

def get_embedding_for_word(word, vocabulary)
  index = vocabulary.lookup(word)
  return nil unless index

  get_embedding(index)
end

#get_embeddings(indices) ⇒ Array<Array<Float>>

Get embeddings for multiple indices (batched)

More efficient than individual calls for batch operations.

Parameters:

  • indices (Array<Integer>) —

    Word indices

Returns:

  • (Array<Array<Float>>) —

    Array of embedding vectors



144
145
146
147
148
149
150
151
152
153
154
155
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 144

def get_embeddings(indices)
  ensure_loaded
  return [] if indices.nil? || indices.empty?

  valid_indices = indices.select { |i| valid_index?(i) }
  return [] if valid_indices.empty?

  # Process in batches for memory efficiency
  valid_indices.each_slice(BATCH_SIZE).flat_map do |batch|
    run_batch_inference(batch)
  end
end

#get_embeddings_for_words(words, vocabulary) ⇒ Hash<String, Array<Float>>

Get embeddings for multiple words using vocabulary

Parameters:

  • words (Array<String>) —

    Words to lookup

  • vocabulary (Vocabulary) —

    Vocabulary for word-to-index mapping

Returns:

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

    Word to embedding mapping



193
194
195
196
197
198
199
200
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 193

def get_embeddings_for_words(words, vocabulary)
  result = {}
  words.each do |word|
    embedding = get_embedding_for_word(word, vocabulary)
    result[word] = embedding if embedding
  end
  result
end

#load! ⇒ self

Load the ONNX model into memory

Returns:

  • (self)

Raises:

  • (Kotoshu::Models::OnnxUnavailable) —

    if onnxruntime gem is missing

  • (ArgumentError) —

    if model file doesn't exist



71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 71

def load!
  return self if @loaded

  raise Kotoshu::Models::OnnxModel::OnnxUnavailable unless Kotoshu::Models::OnnxModel::ONNX_LOADED

  unless File.exist?(@onnx_path)
    raise ArgumentError,
          "ONNX file not found: #{@onnx_path}"
  end

  @session = OnnxRuntime::InferenceSession.new(@onnx_path)

  # Detect input/output names
  @input_name = detect_input_name
  @output_name = detect_output_name

  @loaded = true
  self
end

#loaded? ⇒ Boolean

Returns Whether the model is loaded.

Returns:

  • (Boolean) —

    Whether the model is loaded



40
41
42
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 40

def loaded?
  @loaded
end

#model_info ⇒ Hash

Get model information

Returns:

  • (Hash) 


230
231
232
233
234
235
236
237
238
239
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 230

def model_info
  {
    type: 'onnx',
    language: @language_code,
    dimension: @dimension,
    path: @onnx_path,
    loaded: @loaded,
    inference_count: @inference_count,
  }
end

#model_type ⇒ String

Get model type identifier

Returns:

  • (String) 


222
223
224
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 222

def model_type
  'onnx'
end

#preload_embeddings!(vocabulary) ⇒ Hash<Integer, Array<Float>>

Preload all embeddings into memory

For small vocabularies, this provides O(1) lookup after loading.

Parameters:

  • vocabulary (Vocabulary) —

    Vocabulary with complete word list

Returns:

  • (Hash<Integer, Array<Float>>) —

    Index to embedding mapping



164
165
166
167
168
169
170
171
172
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 164

def preload_embeddings!(vocabulary)
  ensure_loaded

  all_indices = (0...vocabulary.size).to_a
  embeddings = get_embeddings(all_indices)

  # Build index mapping
  all_indices.zip(embeddings).to_h
end

#ready? ⇒ Boolean

Check if model is ready for inference

Returns:

  • (Boolean) 


107
108
109
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 107

def ready?
  @loaded && !@session.nil?
end

#supports_batching? ⇒ Boolean

Check if batching is supported

Returns:

  • (Boolean) 


206
207
208
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 206

def supports_batching?
  true
end

#to_s ⇒ String Also known as: inspect

String representation

Returns:

  • (String) 


283
284
285
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 283

def to_s
  "OnnxRuntimeModel(language: #{@language_code}, dimension: #{@dimension}, loaded: #{@loaded})"
end

#unload! ⇒ self

Unload the model from memory

Returns:

  • (self) 


95
96
97
98
99
100
101
 
# File 'lib/kotoshu/embeddings/onnx_runtime_model.rb', line 95

def unload!
  @session = nil
  @input_name = nil
  @output_name = nil
  @loaded = false
  self
end