Skip to content

Commit 850c159

Browse files
committed
fix: vllm reranker passes top_n twice to cohere rerank()
1 parent cdd974d commit 850c159

1 file changed

Lines changed: 4 additions & 2 deletions

File tree

  • apps/models_provider/impl/vllm_model_provider/model

apps/models_provider/impl/vllm_model_provider/model/reranker.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,12 +31,14 @@ def is_cache_model():
3131
@staticmethod
3232
def new_instance(model_type, model_name, model_credential: Dict[str, object], **model_kwargs):
3333
r_url = model_credential.get('api_url')[:-3] if model_credential.get('api_url').endswith('/v1') else model_credential.get('api_url')
34+
optional_params = MaxKBBaseModel.filter_optional_params(model_kwargs)
35+
top_n = optional_params.pop('top_n', 3)
3436
return VllmBgeReranker(
3537
model=model_name,
3638
api_key=model_credential.get('api_key'),
3739
api_url=r_url,
38-
params=model_kwargs,
39-
**model_kwargs
40+
top_n=top_n,
41+
params=optional_params,
4042
)
4143

4244
def compress_documents(self, documents: Sequence[Document], query: str, callbacks: Optional[Callbacks] = None) -> \

0 commit comments

Comments
 (0)