nvidia rerank missed top n

langgenius · Sep 10, 2024 · 72df152 · 72df152
1 parent c8df92d
commit 72df152
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/api/core/model_runtime/model_providers/nvidia/rerank/rerank.py b/api/core/model_runtime/model_providers/nvidia/rerank/rerank.py
@@ -52,9 +52,9 @@ def _invoke(self, model: str, credentials: dict,
             payload = {
                 "model": model,
                 "query": {"text": query},
+                "top_n": 1,
                 "passages": [{"text": doc} for doc in docs],
             }
-
             session = requests.Session()
             response = session.post(invoke_url, headers=headers, json=payload)
             response.raise_for_status()
@@ -71,7 +71,10 @@ def _invoke(self, model: str, credentials: dict,
                 )
 
                 rerank_documents.append(rerank_document)
-
+            if rerank_documents:
+                rerank_documents = sorted(rerank_documents, key=lambda x: x.score, reverse=True)
+                if top_n:
+                    rerank_documents = rerank_documents[:top_n]
             return RerankResult(model=model, docs=rerank_documents)
         except requests.HTTPError as e:
             raise InvokeServerUnavailableError(str(e))