Fixes the validation of acceptable model types for registration in AQUA.

Author: mrDzurb

ads/aqua/extension/model_handler.py

@@ -177,10 +177,8 @@ def _find_matching_aqua_model(model_id: str) -> Optional[AquaModelSummary]:
 
         return None
 
-
-
     @handle_exceptions
-    def get(self,*args, **kwargs):
+    def get(self, *args, **kwargs):
         """
         Finds a list of matching models from hugging face based on query string provided from users.
 
@@ -194,13 +192,11 @@ def get(self,*args, **kwargs):
             Returns the matching model ids string
         """
 
-        query=self.get_argument("query",default=None)
+        query = self.get_argument("query", default=None)
         if not query:
-            raise HTTPError(400,Errors.MISSING_REQUIRED_PARAMETER.format("query"))
-        models=list_hf_models(query)
-        return self.finish({"models":models})
-
-
+            raise HTTPError(400, Errors.MISSING_REQUIRED_PARAMETER.format("query"))
+        models = list_hf_models(query)
+        return self.finish({"models": models})
 
     @handle_exceptions
     def post(self, *args, **kwargs):
@@ -235,9 +231,9 @@ def post(self, *args, **kwargs):
             )
 
         # Check pipeline_tag, it should be `text-generation`
-        if (
-            not hf_model_info.pipeline_tag
-            or hf_model_info.pipeline_tag.lower() != ModelTask.TEXT_GENERATION
+        if not (
+            hf_model_info.pipeline_tag
+            and hf_model_info.pipeline_tag.lower() in ModelTask
         ):
             raise AquaRuntimeError(
                 f"Unsupported pipeline tag for the chosen model: '{hf_model_info.pipeline_tag}'. "

ads/aqua/model/constants.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright (c) 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
@@ -9,6 +8,7 @@
 
 This module contains constants/enums used in Aqua Model.
 """
+
 from ads.common.extended_enum import ExtendedEnumMeta
 
 
@@ -21,6 +21,8 @@ class ModelCustomMetadataFields(str, metaclass=ExtendedEnumMeta):
 
 class ModelTask(str, metaclass=ExtendedEnumMeta):
     TEXT_GENERATION = "text-generation"
+    IMAGE_TEXT_TO_TEXT = "image-text-to-text"
+    IMAGE_TO_TEXT = "image-to-text"
 
 
 class FineTuningMetricCategories(str, metaclass=ExtendedEnumMeta):