Skip to content

Commit 70207b1

Browse files
committed
refactor: enhance VllmBgeReranker initialization with optional parameters and default top_n value
1 parent 0817d10 commit 70207b1

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
@@ -35,12 +35,14 @@ def new_instance(model_type, model_name, model_credential: Dict[str, object], **
3535
if model_credential.get("api_url").endswith("/v1")
3636
else model_credential.get("api_url")
3737
)
38+
optional_params = MaxKBBaseModel.filter_optional_params(model_kwargs)
39+
top_n = optional_params.pop("top_n", 3)
3840
return VllmBgeReranker(
3941
model=model_name,
4042
api_key=model_credential.get("api_key"),
4143
api_url=r_url,
42-
params=model_kwargs,
43-
**model_kwargs,
44+
top_n=top_n,
45+
params=optional_params,
4446
)
4547

4648
def compress_documents(

0 commit comments

Comments
 (0)