From 7831fa027203636e901b6337d6f202fe4bf64e9b Mon Sep 17 00:00:00 2001 From: muhammed-shihebi <58932402+muhammed-shihebi@users.noreply.github.com> Date: Wed, 27 Mar 2024 22:06:17 +0000 Subject: [PATCH] fix: add correct default values in DeployRequest --- model-manager/model_manager/app/models/management.py | 4 ++-- model-manager/model_manager/app/routers/management.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/model-manager/model_manager/app/models/management.py b/model-manager/model_manager/app/models/management.py index 74121b061..cfccd5dce 100644 --- a/model-manager/model_manager/app/models/management.py +++ b/model-manager/model_manager/app/models/management.py @@ -57,8 +57,8 @@ class DeployRequest(BaseModel): disable_gpu: Optional[bool] = Field( True, description="whether to use gpu for inference" ) - batch_size: int = Field("", description="input batch size") - max_input: int = Field("", description="max input length") + batch_size: int = Field(32, description="input batch size") + max_input: int = Field(512, description="max input length") transformers_cache: Optional[str] = Field( "../.cache", description="path to cache models" ) diff --git a/model-manager/model_manager/app/routers/management.py b/model-manager/model_manager/app/routers/management.py index 4551954d9..4fb9a17f0 100644 --- a/model-manager/model_manager/app/routers/management.py +++ b/model-manager/model_manager/app/routers/management.py @@ -59,8 +59,8 @@ async def get_all_models(): # token: str = Depends(client_credentials)): model_type=model["MODEL_TYPE"], model_name=model["MODEL_NAME"], disable_gpu=model["DISABLE_GPU"], - batch_size=model["BATCH_SIZE"], - max_input=model["MAX_INPUT_SIZE"], + batch_size=model.get("BATCH_SIZE", 32), + max_input=model.get("MAX_INPUT_SIZE", 512), model_class=model["MODEL_CLASS"], return_plaintext_arrays=model["RETURN_PLAINTEXT_ARRAYS"], )