refac: rerank

This commit is contained in:
Timothy Jaeryang Baek 2025-11-09 21:33:50 -05:00
parent e76f77bcb7
commit bc739de024

View file

@ -516,11 +516,13 @@ def get_reranking_function(reranking_engine, reranking_model, reranking_function
if reranking_function is None:
return None
if reranking_engine == "external":
return lambda sentences, user=None: reranking_function.predict(
sentences, user=user
return lambda query, documents, user=None: reranking_function.predict(
[(query, doc.page_content) for doc in documents], user=user
)
else:
return lambda sentences, user=None: reranking_function.predict(sentences)
return lambda query, documents, user=None: reranking_function.predict(
[(query, doc.page_content) for doc in documents]
)
def get_sources_from_items(
@ -1064,9 +1066,7 @@ class RerankCompressor(BaseDocumentCompressor):
scores = None
if reranking:
scores = self.reranking_function(
[(query, doc.page_content) for doc in documents]
)
scores = self.reranking_function(query, documents)
else:
from sentence_transformers import util