From cb0c55daa7bf81081d95d9adb7ada07ca6cf437e Mon Sep 17 00:00:00 2001 From: AkisAya Date: Thu, 21 Nov 2024 17:53:20 +0800 Subject: [PATCH] fix weight rerank of knowledge retrieval (#10931) --- api/core/rag/rerank/rerank_model.py | 6 +++--- api/core/rag/rerank/weight_rerank.py | 7 +++---- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/api/core/rag/rerank/rerank_model.py b/api/core/rag/rerank/rerank_model.py index fc82b2080b..6ae432a526 100644 --- a/api/core/rag/rerank/rerank_model.py +++ b/api/core/rag/rerank/rerank_model.py @@ -27,11 +27,11 @@ class RerankModelRunner(BaseRerankRunner): :return: """ docs = [] - doc_id = set() + doc_ids = set() unique_documents = [] for document in documents: - if document.provider == "dify" and document.metadata["doc_id"] not in doc_id: - doc_id.add(document.metadata["doc_id"]) + if document.provider == "dify" and document.metadata["doc_id"] not in doc_ids: + doc_ids.add(document.metadata["doc_id"]) docs.append(document.page_content) unique_documents.append(document) elif document.provider == "external": diff --git a/api/core/rag/rerank/weight_rerank.py b/api/core/rag/rerank/weight_rerank.py index b706f29bb1..4719be012f 100644 --- a/api/core/rag/rerank/weight_rerank.py +++ b/api/core/rag/rerank/weight_rerank.py @@ -37,11 +37,10 @@ class WeightRerankRunner(BaseRerankRunner): :return: """ unique_documents = [] - doc_id = set() + doc_ids = set() for document in documents: - doc_id = document.metadata.get("doc_id") - if doc_id not in doc_id: - doc_id.add(doc_id) + if document.metadata["doc_id"] not in doc_ids: + doc_ids.add(document.metadata["doc_id"]) unique_documents.append(document) documents = unique_documents