fix: fix rerank model runner usage

pull/19008/head
Ray Wu 1 year ago
parent b8bb45b106
commit 009d391400

@ -52,6 +52,7 @@ class RerankModelRunner(BaseRerankRunner):
rerank_documents = [] rerank_documents = []
for result in rerank_result.docs: for result in rerank_result.docs:
if score_threshold is None or result.score >= score_threshold:
# format document # format document
rerank_document = Document( rerank_document = Document(
page_content=result.text, page_content=result.text,
@ -62,4 +63,5 @@ class RerankModelRunner(BaseRerankRunner):
rerank_document.metadata["score"] = result.score rerank_document.metadata["score"] = result.score
rerank_documents.append(rerank_document) rerank_documents.append(rerank_document)
return rerank_documents rerank_documents.sort(key=lambda x: x.metadata.get('score', 0.0), reverse=True)
return rerank_documents[:top_n] if top_n else rerank_documents

Loading…
Cancel
Save