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 



1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
 
# File 'lib/informers/pipelines.rb', line 1054

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