| @@ -54,7 +54,6 @@ class NvidiaRerankModel(RerankModel): | |||
| "query": {"text": query}, | |||
| "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 +70,10 @@ class NvidiaRerankModel(RerankModel): | |||
| ) | |||
| 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)) | |||