Class: NerRuby::Pipeline

Inherits:
Object
  • Object
show all
Defined in:
lib/ner_ruby/pipeline.rb

Instance Method Summary collapse

Constructor Details

#initialize(model:, tokenizer:, decoder: nil) ⇒ Pipeline

Returns a new instance of Pipeline.



5
6
7
8
9
 
# File 'lib/ner_ruby/pipeline.rb', line 5

def initialize(model:, tokenizer:, decoder: nil)
  @model = model
  @tokenizer = tokenizer
  @decoder = decoder || Decoder.new
end

Instance Method Details

#call(text) ⇒ Object 



11
12
13
14
15
16
17
18
19
20
21
22
 
# File 'lib/ner_ruby/pipeline.rb', line 11

def call(text)
  encoding = @tokenizer.encode(text)
  tokens = encoding[:tokens] || encoding["tokens"]
  input_ids = encoding[:ids] || encoding["ids"]

  logits = @model.predict(input_ids)

  predictions = logits.map { |row| row.each_with_index.max_by { |v, _| v }.last }
  scores = logits.map { |row| softmax(row).max }

  @decoder.decode(tokens, predictions, scores: scores, original_text: text)
end