Adding task for unverified model import

Author: kumar-shivam-ranjan

ads/aqua/config/config.py

@@ -29,3 +29,11 @@ def get_evaluation_service_config(
         .get(ContainerSpec.CONTAINER_SPEC, {})
         .get(container, {})
     )
+
+def get_valid_tasks():
+    return [
+        "text_generation",
+        "code_synthesis",
+        "image_text_to_text",
+        "feature_extraction",
+    ]

ads/aqua/extension/model_handler.py

@@ -13,6 +13,7 @@
     get_hf_model_info,
     list_hf_models,
 )
+from ads.aqua.config.config import get_valid_tasks
 from ads.aqua.extension.base_handler import AquaAPIhandler
 from ads.aqua.extension.errors import Errors
 from ads.aqua.model import AquaModelApp
@@ -119,6 +120,9 @@ def post(self, *args, **kwargs):
         os_path = input_data.get("os_path")
         if not os_path:
             raise HTTPError(400, Errors.MISSING_REQUIRED_PARAMETER.format("os_path"))
+        task = input_data.get("task")
+        if task not in get_valid_tasks():
+            raise HTTPError(400, Errors.INVALID_VALUE_OF_PARAMETER.format("task"))
 
         inference_container = input_data.get("inference_container")
         finetuning_container = input_data.get("finetuning_container")
@@ -128,7 +132,6 @@ def post(self, *args, **kwargs):
         download_from_hf = (
             str(input_data.get("download_from_hf", "false")).lower() == "true"
         )
-
         return self.finish(
             AquaModelApp().register(
                 model=model,
@@ -139,6 +142,7 @@ def post(self, *args, **kwargs):
                 compartment_id=compartment_id,
                 project_id=project_id,
                 model_file=model_file,
+                task=task,
             )
         )
 
@@ -164,6 +168,8 @@ def put(self, id):
 
         enable_finetuning = input_data.get("enable_finetuning")
         task = input_data.get("task")
+        if task not in get_valid_tasks():
+            raise HTTPError(400, Errors.INVALID_VALUE_OF_PARAMETER.format("task"))
         return self.finish(
             AquaModelApp().edit_registered_model(
                 id, inference_container, enable_finetuning, task

ads/aqua/model/entities.py

@@ -284,6 +284,7 @@ class ImportModelDetails(CLIBuilderMixin):
     local_dir: Optional[str] = None
     inference_container: Optional[str] = None
     finetuning_container: Optional[str] = None
+    task: Optional[str] = None
     compartment_id: Optional[str] = None
     project_id: Optional[str] = None
     model_file: Optional[str] = None

ads/aqua/model/model.py

@@ -735,6 +735,7 @@ def _create_model_catalog_entry(
         model_name: str,
         inference_container: str,
         finetuning_container: str,
+        task: Optional[str],
         verified_model: DataScienceModel,
         validation_result: ModelValidationResult,
         compartment_id: Optional[str],
@@ -790,6 +791,7 @@ def _create_model_catalog_entry(
                     json_dict=verified_model.model_file_description
                 )
         else:
+            tags.update({Tags.TASK: task})
             metadata = ModelCustomMetadata()
             if not inference_container:
                 raise AquaRuntimeError(
@@ -1336,6 +1338,7 @@ def register(
             model_name=model_name,
             inference_container=import_model_details.inference_container,
             finetuning_container=import_model_details.finetuning_container,
+            task=import_model_details.task,
             verified_model=verified_model,
             validation_result=validation_result,
             compartment_id=import_model_details.compartment_id,