Class: OnnxRuby::Classifier

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

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(model_path, tokenizer: nil, labels: nil, **session_opts) ⇒ Classifier

Returns a new instance of Classifier.



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

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

Instance Attribute Details

#labels ⇒ Object (readonly)

Returns the value of attribute labels.



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

def labels
  @labels
end

#session ⇒ Object (readonly)

Returns the value of attribute session.



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

def session
  @session
end

Instance Method Details

#predict(input) ⇒ Hash

Classify a single input

Parameters:

  • input (String, Array<Float>) —

    text (requires tokenizer) or feature vector

Returns:

  • (Hash) —

    { label:, score:, scores: }



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

def predict(input)
  predict_batch([input]).first
end

#predict_batch(inputs) ⇒ Array<Hash>

Classify a batch of inputs

Parameters:

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

    batch of texts or feature vectors

Returns:

  • (Array<Hash>) —

    array of { label:, score:, scores: }



25
26
27
28
29
30
31
32
 
# File 'lib/onnx_ruby/classifier.rb', line 25

def predict_batch(inputs)
  feed = prepare_inputs(inputs)
  result = @session.run(feed)

  logits = find_output(result, %w[logits output probabilities scores])

  logits.map { |row| format_prediction(row) }
end