diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py index 3a07e60067..a1ec3a17ea 100644 --- a/rag/llm/rerank_model.py +++ b/rag/llm/rerank_model.py @@ -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: