fix: fix rerank model runner usage

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

@ -52,14 +52,16 @@ class RerankModelRunner(BaseRerankRunner):
rerank_documents = [] rerank_documents = []
for result in rerank_result.docs: for result in rerank_result.docs:
# format document if score_threshold is None or result.score >= score_threshold:
rerank_document = Document( # format document
page_content=result.text, rerank_document = Document(
metadata=documents[result.index].metadata, page_content=result.text,
provider=documents[result.index].provider, metadata=documents[result.index].metadata,
) provider=documents[result.index].provider,
if rerank_document.metadata is not None: )
rerank_document.metadata["score"] = result.score if rerank_document.metadata is not None:
rerank_documents.append(rerank_document) rerank_document.metadata["score"] = result.score
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