AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-08-12 22:29:03 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: better WeightRerankRunner run logic use O(1) and delete unused code (#10849)

Browse Source 
Signed-off-by: yihong0618 <zouzou0208@gmail.com>
This commit is contained in:
yihong 2024-11-19 20:12:13 +08:00 committed by GitHub
parent bc1013dacf
commit 58a9d9eb9a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
api/core/rag/rerank/weight_rerank.py
View File

@ -36,23 +36,21 @@ class WeightRerankRunner(BaseRerankRunner):
:return:
"""
docs = []
doc_id = []
unique_documents = []
doc_id = set()
for document in documents:
if document.metadata["doc_id"] not in doc_id:
doc_id.append(document.metadata["doc_id"])
docs.append(document.page_content)
doc_id = document.metadata.get("doc_id")
if doc_id not in doc_id:
doc_id.add(doc_id)
unique_documents.append(document)
documents = unique_documents
rerank_documents = []
query_scores = self._calculate_keyword_score(query, documents)
query_vector_scores = self._calculate_cosine(self.tenant_id, query, documents, self.weights.vector_setting)
rerank_documents = []
for document, query_score, query_vector_score in zip(documents, query_scores, query_vector_scores):
# format document
score = (
self.weights.vector_setting.vector_weight * query_vector_score
+ self.weights.keyword_setting.keyword_weight * query_score
@ -61,7 +59,8 @@ class WeightRerankRunner(BaseRerankRunner):
continue
document.metadata["score"] = score
rerank_documents.append(document)
rerank_documents = sorted(rerank_documents, key=lambda x: x.metadata["score"], reverse=True)
rerank_documents.sort(key=lambda x: x.metadata["score"], reverse=True)
return rerank_documents[:top_n] if top_n else rerank_documents
def _calculate_keyword_score(self, query: str, documents: list[Document]) -> list[float]: