From 87a67cce7d061baa6dd0a4a1c674d19a74b85b2a Mon Sep 17 00:00:00 2001 From: Cooper <42096311+FOkvj@users.noreply.github.com> Date: Thu, 9 May 2024 11:30:50 +0800 Subject: [PATCH] fix(rag): Fix rerank bug of EmbeddingRetriever (#1499) Co-authored-by: dong --- dbgpt/rag/retriever/embedding.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/dbgpt/rag/retriever/embedding.py b/dbgpt/rag/retriever/embedding.py index e2be7b7be..0ce896761 100644 --- a/dbgpt/rag/retriever/embedding.py +++ b/dbgpt/rag/retriever/embedding.py @@ -120,7 +120,7 @@ def _retrieve_with_score( new_candidates_with_score = cast( List[Chunk], reduce(lambda x, y: x + y, candidates_with_score) ) - new_candidates_with_score = self._rerank.rank(new_candidates_with_score) + new_candidates_with_score = self._rerank.rank(new_candidates_with_score, query) return new_candidates_with_score async def _aretrieve( @@ -207,7 +207,9 @@ async def _aretrieve_with_score( "rerank_cls": self._rerank.__class__.__name__, }, ): - new_candidates_with_score = self._rerank.rank(new_candidates_with_score) + new_candidates_with_score = self._rerank.rank( + new_candidates_with_score, query + ) return new_candidates_with_score async def _similarity_search(