Class: Informers::RerankingPipeline

Inherits:
Pipeline
  • Object
show all
Defined in:
lib/informers/pipelines.rb

Instance Method Summary collapse

Methods inherited from Pipeline

#initialize

Constructor Details

This class inherits a constructor from Informers::Pipeline

Instance Method Details

#call(query, documents, return_documents: false, top_k: nil) ⇒ Object 



321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
 
# File 'lib/informers/pipelines.rb', line 321

def call(
  query,
  documents,
  return_documents: false,
  top_k: nil
)
  model_inputs = @tokenizer.([query] * documents.size,
    text_pair: documents,
    padding: true,
    truncation: true
  )

  outputs = @model.(model_inputs)

  result =
    Utils.sigmoid(outputs[0].map(&:first))
      .map.with_index { |s, i| {doc_id: i, score: s} }
      .sort_by { |v| -v[:score] }

  if return_documents
    result.each do |v|
      v[:text] = documents[v[:doc_id]]
    end
  end

  top_k ? result.first(top_k) : result
end