Class: OnnxRuby::Embedder

Inherits:
Object
  • Object
show all
Includes:
TokenizerSupport
Defined in:
lib/onnx_ruby/embedder.rb

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(model_path, tokenizer: nil, normalize: true, **session_opts) ⇒ Embedder

Returns a new instance of Embedder.



9
10
11
12
13
 
# File 'lib/onnx_ruby/embedder.rb', line 9

def initialize(model_path, tokenizer: nil, normalize: true, **session_opts)
  @session = Session.new(model_path, **session_opts)
  @normalize = normalize
  @tokenizer = resolve_tokenizer(tokenizer)
end

Instance Attribute Details

#session ⇒ Object (readonly)

Returns the value of attribute session.



7
8
9
 
# File 'lib/onnx_ruby/embedder.rb', line 7

def session
  @session
end

Instance Method Details

#embed(input) ⇒ Array<Float>

Embed a single text or pre-tokenized input

Parameters:

  • input (String, Hash) —

    text string (requires tokenizer) or hash of input tensors

Returns:

  • (Array<Float>) —

    embedding vector



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

def embed(input)
  embed_batch([input]).first
end

#embed_batch(inputs) ⇒ Array<Array<Float>>

Embed a batch of texts or pre-tokenized inputs

Parameters:

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

    batch of texts or tensor hashes

Returns:

  • (Array<Array<Float>>) —

    array of embedding vectors



25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
 
# File 'lib/onnx_ruby/embedder.rb', line 25

def embed_batch(inputs)
  @_masks = nil
  feed = prepare_inputs(inputs)
  result = @session.run(feed)

  raw = find_output(result, %w[embeddings sentence_embedding output last_hidden_state])
  return [] if raw.nil? || raw.empty?

  # If output is 3D (batch, seq_len, dim) — do mean pooling
  embeddings = if raw.first.is_a?(Array) && raw.first.first.is_a?(Array)
                 mean_pool(raw, @_masks)
               else
                 raw
               end

  embeddings.map { |vec| @normalize ? l2_normalize(vec) : vec }
end