From fc60b554a1c3d1767608bb7d54a61b769ac49fa6 Mon Sep 17 00:00:00 2001 From: Ziyu Huang Date: Thu, 10 Oct 2024 15:18:07 +0800 Subject: [PATCH] Fixes #9159: Modify to make it works to llama.cpp rerank API (#9160) --- .../model_providers/jina/rerank/rerank.py | 12 ++++++++++-- .../model_providers/localai/rerank/rerank.py | 12 ++++++++++-- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/api/core/model_runtime/model_providers/jina/rerank/rerank.py b/api/core/model_runtime/model_providers/jina/rerank/rerank.py index 79ca68914f..0350207651 100644 --- a/api/core/model_runtime/model_providers/jina/rerank/rerank.py +++ b/api/core/model_runtime/model_providers/jina/rerank/rerank.py @@ -61,11 +61,19 @@ class JinaRerankModel(RerankModel): rerank_documents = [] for result in results["results"]: + index = result["index"] + if "document" in result: + text = result["document"]["text"] + else: + # llama.cpp rerank maynot return original documents + text = docs[index] + rerank_document = RerankDocument( - index=result["index"], - text=result["document"]["text"], + index=index, + text=text, score=result["relevance_score"], ) + if score_threshold is None or result["relevance_score"] >= score_threshold: rerank_documents.append(rerank_document) diff --git a/api/core/model_runtime/model_providers/localai/rerank/rerank.py b/api/core/model_runtime/model_providers/localai/rerank/rerank.py index 2b0f53bc19..075b44658d 100644 --- a/api/core/model_runtime/model_providers/localai/rerank/rerank.py +++ b/api/core/model_runtime/model_providers/localai/rerank/rerank.py @@ -70,11 +70,19 @@ class LocalaiRerankModel(RerankModel): rerank_documents = [] for result in results["results"]: + index = result["index"] + if "document" in result: + text = result["document"]["text"] + else: + # llama.cpp rerank maynot return original documents + text = docs[index] + rerank_document = RerankDocument( - index=result["index"], - text=result["document"]["text"], + index=index, + text=text, score=result["relevance_score"], ) + if score_threshold is None or result["relevance_score"] >= score_threshold: rerank_documents.append(rerank_document)