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 



925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
 
# File 'lib/informers/pipelines.rb', line 925

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