Fixes the validation of acceptable model types for registration in AQUA. (#958)

Author: mrDzurb

ads/aqua/extension/model_handler.py

@@ -13,7 +13,6 @@
 from ads.aqua.extension.base_handler import AquaAPIhandler
 from ads.aqua.extension.errors import Errors
 from ads.aqua.model import AquaModelApp
-from ads.aqua.model.constants import ModelTask
 from ads.aqua.model.entities import AquaModelSummary, HFModelSummary
 from ads.aqua.ui import ModelFormat
 
@@ -177,10 +176,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 +191,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):
@@ -234,16 +229,17 @@ def post(self, *args, **kwargs):
                 "Please verify the model's status on the Hugging Face Model Hub or select a different model."
             )
 
-        # 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
-        ):
-            raise AquaRuntimeError(
-                f"Unsupported pipeline tag for the chosen model: '{hf_model_info.pipeline_tag}'. "
-                f"AQUA currently supports the following tasks only: {', '.join(ModelTask.values())}. "
-                "Please select a model with a compatible pipeline tag."
-            )
+        # Commented the validation below to let users to register any model type.
+        # # Check pipeline_tag, it should be `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}'. "
+        #         f"AQUA currently supports the following tasks only: {', '.join(ModelTask.values())}. "
+        #         "Please select a model with a compatible pipeline tag."
+        #     )
 
         # Check if it is a service/verified model
         aqua_model_info: AquaModelSummary = self._find_matching_aqua_model(

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):

tests/unitary/with_extras/aqua/test_model_handler.py

@@ -12,14 +12,15 @@
 from notebook.base.handlers import IPythonHandler
 
 from ads.aqua.common.errors import AquaRuntimeError
+from ads.aqua.common.utils import get_hf_model_info
 from ads.aqua.extension.model_handler import (
+    AquaHuggingFaceHandler,
     AquaModelHandler,
     AquaModelLicenseHandler,
-    AquaHuggingFaceHandler,
 )
 from ads.aqua.model import AquaModelApp
+from ads.aqua.model.constants import ModelTask
 from ads.aqua.model.entities import AquaModel, AquaModelSummary, HFModelSummary
-from ads.aqua.common.utils import get_hf_model_info
 
 
 class ModelHandlerTestCase(TestCase):
@@ -263,21 +264,22 @@ def test_post_negative(self, mock_uuid, mock_format_hf_custom_error_message):
             )
         get_hf_model_info.cache_clear()
 
-        # case 6
-        self.mock_handler.get_json_body = MagicMock(
-            return_value={"model_id": "test_model_id"}
-        )
-        with patch.object(HfApi, "model_info") as mock_model_info:
-            mock_model_info.return_value = MagicMock(
-                disabled=False, id="test_model_id", pipeline_tag="not-text-generation"
-            )
-            self.mock_handler.post()
-            self.mock_handler.finish.assert_called_with(
-                '{"status": 400, "message": "Something went wrong with your request.", '
-                '"service_payload": {}, "reason": "Unsupported pipeline tag for the chosen '
-                "model: 'not-text-generation'. AQUA currently supports the following tasks only: "
-                'text-generation. Please select a model with a compatible pipeline tag.", "request_id": "###"}'
-            )
+        # # case 6 pipeline Tag
+        # self.mock_handler.get_json_body = MagicMock(
+        #     return_value={"model_id": "test_model_id"}
+        # )
+        # with patch.object(HfApi, "model_info") as mock_model_info:
+        #     mock_model_info.return_value = MagicMock(
+        #         disabled=False, id="test_model_id", pipeline_tag="not-text-generation"
+        #     )
+        #     self.mock_handler.post()
+        #     self.mock_handler.finish.assert_called_with(
+        #         '{"status": 400, "message": "Something went wrong with your request.", '
+        #         '"service_payload": {}, "reason": "Unsupported pipeline tag for the chosen '
+        #         "model: 'not-text-generation'. AQUA currently supports the following tasks only: "
+        #         f'{", ".join(ModelTask.values())}. '
+        #         'Please select a model with a compatible pipeline tag.", "request_id": "###"}'
+        #     )
         get_hf_model_info.cache_clear()
 
     @patch("uuid.uuid4")