Module: OnnxRuby::Model

Defined in:
lib/onnx_ruby/model.rb

Overview

ActiveModel-style mixin for embedding generation.

Usage:

class Document
include OnnxRuby::Model

onnx_model "embeddings.onnx"
onnx_input ->(doc) { { "input_ids" => doc.token_ids, "attention_mask" => doc.mask } }
onnx_output "embeddings"
end

doc = Document.new
doc.onnx_predict  # => [0.123, -0.456, ...]

Defined Under Namespace

Modules: ClassMethods

Class Method Summary collapse

Instance Method Summary collapse

Class Method Details

.included(base) ⇒ Object 



18
19
20
 
# File 'lib/onnx_ruby/model.rb', line 18

def self.included(base)
  base.extend(ClassMethods)
end

Instance Method Details

#onnx_predict(**run_opts) ⇒ Object

Raises:



60
61
62
63
64
65
66
67
68
69
 
# File 'lib/onnx_ruby/model.rb', line 60

def onnx_predict(**run_opts)
  input_fn = self.class.onnx_input
  raise Error, "onnx_input not defined on #{self.class}" unless input_fn

  inputs = input_fn.call(self)
  result = self.class.onnx_session.run(inputs, **run_opts)

  output_name = self.class.onnx_output
  output_name ? result[output_name] : result.values.first
end