Skip to content

Commit c5937c6

Browse files
committed
fix: vllm reranker passes top_n twice to cohere rerank()
1 parent bc2bf62 commit c5937c6

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
@@ -32,12 +32,14 @@ def is_cache_model():
3232
@staticmethod
3333
def new_instance(model_type, model_name, model_credential: Dict[str, object], **model_kwargs):
3434
r_url = model_credential.get('api_url')[:-3] if model_credential.get('api_url').endswith('/v1') else model_credential.get('api_url')
35+
optional_params = MaxKBBaseModel.filter_optional_params(model_kwargs)
36+
top_n = optional_params.pop('top_n', 3)
3537
return VllmBgeReranker(
3638
model=model_name,
3739
api_key=model_credential.get('api_key'),
3840
api_url=r_url,
39-
params=model_kwargs,
40-
**model_kwargs
41+
top_n=top_n,
42+
params=optional_params,
4143
)
4244

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

0 commit comments

Comments
 (0)