Browse Source

Fixes #9159: Modify to make it works to llama.cpp rerank API (#9160)

Ziyu Huang 6 months ago
parent
commit
fc60b554a1

+ 10 - 2
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)
 

+ 10 - 2
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)