refactor: improve QwenRerank logic (#14388)

### What problem does this PR solve?

improve QwenRerank logic

### Type of change

- [x] Refactoring
This commit is contained in:
Stephen Hu
2026-04-28 20:17:34 +08:00
committed by GitHub
parent 0d18b293f5
commit 345bec812d

View File

@@ -365,7 +365,7 @@ class VoyageRerank(Base):
class QWenRerank(Base):
_FACTORY_NAME = "Tongyi-Qianwen"
def __init__(self, key, model_name="gte-rerank", base_url=None, **kwargs):
def __init__(self, key, model_name="gte-rerank", **kwargs):
import dashscope
self.api_key = key
@@ -376,18 +376,19 @@ class QWenRerank(Base):
import dashscope
# qwen3-rerank does not support return_documents parameter
if self.model_name.startswith("qwen3-rerank"):
resp = dashscope.TextReRank.call(
api_key=self.api_key, model=self.model_name,
query=query, documents=texts, top_n=len(texts)
)
else:
resp = dashscope.TextReRank.call(
api_key=self.api_key, model=self.model_name,
query=query, documents=texts,
top_n=len(texts), return_documents=False
)
# Build call parameters
call_kwargs = {
"api_key": self.api_key,
"model": self.model_name,
"query": query,
"documents": texts,
"top_n": len(texts)
}
# qwen3-rerank does not support return_documents parameter
if not self.model_name.startswith("qwen3-rerank"):
call_kwargs["return_documents"] = False
resp = dashscope.TextReRank.call(**call_kwargs)
rank = np.zeros(len(texts), dtype=float)
if resp.status_code == HTTPStatus.OK: