github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-10 13:29:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(rag): Fix rerank bug of EmbeddingRetriever (#1499)

Browse Source 
Co-authored-by: dong <dongzhancai@iie2.com>
This commit is contained in:
Cooper
2024-05-09 11:30:50 +08:00
committed by GitHub
parent b38305b72b
commit 87a67cce7d
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
dbgpt/rag/retriever/embedding.py
View File

@@ -120,7 +120,7 @@ class EmbeddingRetriever(BaseRetriever):
new_candidates_with_score = cast( new_candidates_with_score = cast(
List[Chunk], reduce(lambda x, y: x + y, candidates_with_score) 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 return new_candidates_with_score
async def _aretrieve( async def _aretrieve(
@@ -207,7 +207,9 @@ class EmbeddingRetriever(BaseRetriever):
"rerank_cls": self._rerank.__class__.__name__, "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 return new_candidates_with_score
async def _similarity_search( async def _similarity_search(