浏览代码

Resolve 8475 support rerank model from infinity (#10939)

Co-authored-by: linyanxu <linyanxu2@qq.com>
tags/0.12.0
LastHopeOfGPNU 11 个月前
父节点
当前提交
1a6b961b5f
没有帐户链接到提交者的电子邮件
共有 1 个文件被更改,包括 8 次插入2 次删除
  1. 8
    2
      api/core/model_runtime/model_providers/openai_api_compatible/rerank/rerank.py

+ 8
- 2
api/core/model_runtime/model_providers/openai_api_compatible/rerank/rerank.py 查看文件

@@ -64,7 +64,7 @@ class OAICompatRerankModel(RerankModel):

# TODO: Do we need truncate docs to avoid llama.cpp return error?

data = {"model": model_name, "query": query, "documents": docs, "top_n": top_n}
data = {"model": model_name, "query": query, "documents": docs, "top_n": top_n, "return_documents": True}

try:
response = post(str(URL(url) / "rerank"), headers=headers, data=dumps(data), timeout=60)
@@ -83,7 +83,13 @@ class OAICompatRerankModel(RerankModel):
index = result["index"]

# Retrieve document text (fallback if llama.cpp rerank doesn't return it)
text = result.get("document", {}).get("text", docs[index])
text = docs[index]
document = result.get("document", {})
if document:
if isinstance(document, dict):
text = document.get("text", docs[index])
elif isinstance(document, str):
text = document

# Normalize the score
normalized_score = (result["relevance_score"] - min_score) / score_range

正在加载...
取消
保存