Austin/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2024-11-16 03:32:23 +08:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

nvidia rerank top n missed (#8185)

Browse Source
This commit is contained in:
Jyong 2024-09-10 13:17:48 +08:00 committed by GitHub
parent fede54be77
commit 2d690801d1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
api/core/model_runtime/model_providers/nvidia/rerank/rerank.py
View File

@ -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))