fix: add correct default values in DeployRequest

UKP-SQuARE · Mar 27, 2024 · 7831fa0 · 7831fa0
1 parent 20bec8e
commit 7831fa0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git 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
@@ -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"],
             )