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

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

View File

@@ -120,7 +120,7 @@ class EmbeddingRetriever(BaseRetriever):
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 @@ class EmbeddingRetriever(BaseRetriever):
"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(