From d5695b3170a4cbfc03a19df57541593b7ebeb1bd Mon Sep 17 00:00:00 2001 From: Jyong <76649700+JohnJyong@users.noreply.github.com> Date: Mon, 11 Dec 2023 16:12:11 +0800 Subject: [PATCH] check rerank document is not empty (#1740) Co-authored-by: jyong --- .../model_providers/models/reranking/cohere_reranking.py | 7 +++++-- .../models/reranking/xinference_reranking.py | 2 ++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/api/core/model_providers/models/reranking/cohere_reranking.py b/api/core/model_providers/models/reranking/cohere_reranking.py index e5b2fe61d1..fa2748734d 100644 --- a/api/core/model_providers/models/reranking/cohere_reranking.py +++ b/api/core/model_providers/models/reranking/cohere_reranking.py @@ -24,7 +24,10 @@ class CohereReranking(BaseReranking): super().__init__(model_provider, client, name) - def rerank(self, query: str, documents: List[Document], score_threshold: Optional[float], top_k: Optional[int]) -> Optional[List[Document]]: + def rerank(self, query: str, documents: List[Document], score_threshold: Optional[float], top_k: Optional[int]) -> \ + Optional[List[Document]]: + if not documents: + return [] docs = [] doc_id = [] unique_documents = [] @@ -34,7 +37,7 @@ class CohereReranking(BaseReranking): docs.append(document.page_content) unique_documents.append(document) documents = unique_documents - + results = self.client.rerank(query=query, documents=docs, model=self.name, top_n=top_k) rerank_documents = [] diff --git a/api/core/model_providers/models/reranking/xinference_reranking.py b/api/core/model_providers/models/reranking/xinference_reranking.py index 0ae9eaec6e..fda2772c70 100644 --- a/api/core/model_providers/models/reranking/xinference_reranking.py +++ b/api/core/model_providers/models/reranking/xinference_reranking.py @@ -21,6 +21,8 @@ class XinferenceReranking(BaseReranking): super().__init__(model_provider, client, name) def rerank(self, query: str, documents: List[Document], score_threshold: Optional[float], top_k: Optional[int]) -> Optional[List[Document]]: + if not documents: + return [] docs = [] doc_id = [] unique_documents = []