Skip to content

Commit 5a71b0b

Browse files
Adding validation for cli
1 parent a2f57f7 commit 5a71b0b

File tree

1 file changed

+7
-1
lines changed

1 file changed

+7
-1
lines changed

ads/aqua/model/model.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
read_file,
3434
upload_folder,
3535
)
36+
from ads.aqua.config.config import get_valid_tasks
3637
from ads.aqua.constants import (
3738
AQUA_MODEL_ARTIFACT_CONFIG,
3839
AQUA_MODEL_ARTIFACT_CONFIG_MODEL_NAME,
@@ -407,7 +408,12 @@ def edit_registered_model(self, id, inference_container, enable_finetuning, task
407408

408409
custom_metadata_list.remove("modelDescription")
409410
if task:
410-
freeform_tags.update({Tags.TASK: task})
411+
if task in get_valid_tasks():
412+
freeform_tags.update({Tags.TASK: task})
413+
else:
414+
raise AquaValueError(
415+
f"Failed to edit model with the given task parameter. Acceptable values are: {get_valid_tasks()}"
416+
)
411417

412418
updated_custom_metadata_list = [
413419
Metadata(**metadata)

0 commit comments

Comments
 (0)