diff --git a/ads/aqua/model/constants.py b/ads/aqua/model/constants.py index 194245fe4..7daeaa57b 100644 --- a/ads/aqua/model/constants.py +++ b/ads/aqua/model/constants.py @@ -26,6 +26,7 @@ class ModelTask(ExtendedEnum): TEXT_GENERATION = "text-generation" IMAGE_TEXT_TO_TEXT = "image-text-to-text" IMAGE_TO_TEXT = "image-to-text" + TIME_SERIES_FORECASTING = "time-series-forecasting" class FineTuningMetricCategories(ExtendedEnum): diff --git a/ads/aqua/modeldeployment/deployment.py b/ads/aqua/modeldeployment/deployment.py index 9834eddd1..20d91ead0 100644 --- a/ads/aqua/modeldeployment/deployment.py +++ b/ads/aqua/modeldeployment/deployment.py @@ -48,7 +48,11 @@ ) from ads.aqua.data import AquaResourceIdentifier from ads.aqua.model import AquaModelApp -from ads.aqua.model.constants import AquaModelMetadataKeys, ModelCustomMetadataFields +from ads.aqua.model.constants import ( + AquaModelMetadataKeys, + ModelCustomMetadataFields, + ModelTask, +) from ads.aqua.model.utils import ( extract_base_model_from_ft, extract_fine_tune_artifacts_path, @@ -215,6 +219,14 @@ def create( freeform_tags=freeform_tags, defined_tags=defined_tags, ) + task_tag = aqua_model.freeform_tags.get(Tags.TASK, UNKNOWN) + if ( + task_tag == ModelTask.TIME_SERIES_FORECASTING + or task_tag == ModelTask.TIME_SERIES_FORECASTING.replace("-", "_") + ): + create_deployment_details.env_var.update( + {Tags.TASK.upper(): ModelTask.TIME_SERIES_FORECASTING} + ) return self._create( aqua_model=aqua_model, create_deployment_details=create_deployment_details,