Addressing review comments

Author: kumar-shivam-ranjan

ads/aqua/app.py

@@ -6,7 +6,7 @@
 import os
 import traceback
 from dataclasses import fields
-from typing import Dict, Union
+from typing import Dict, Optional, Union
 
 import oci
 from oci.data_science.models import UpdateModelDetails, UpdateModelProvenanceDetails
@@ -268,7 +268,12 @@ def if_artifact_exist(self, model_id: str, **kwargs) -> bool:
                 logger.info(f"Artifact not found in model {model_id}.")
                 return False
 
-    def get_config(self, model_id: str, config_file_name: str) -> Dict:
+    def get_config(
+        self,
+        model_id: str,
+        config_file_name: str,
+        config_folder: Optional[str] = "config",
+    ) -> Dict:
         """Gets the config for the given Aqua model.
 
         Parameters
@@ -277,6 +282,9 @@ def get_config(self, model_id: str, config_file_name: str) -> Dict:
             The OCID of the Aqua model.
         config_file_name: str
             name of the config file
+        config_folder: Optional[str]
+            subfolder path where config_file_name needs to be searched
+            default value: config
 
         Returns
         -------
@@ -306,11 +314,11 @@ def get_config(self, model_id: str, config_file_name: str) -> Dict:
             )
             base_model = self.ds_client.get_model(base_model_ocid).data
             artifact_path = get_artifact_path(base_model.custom_metadata_list)
-            config_path = f"{os.path.dirname(artifact_path)}/config/"
+            config_path = f"{os.path.dirname(artifact_path)}/{config_folder}/"
         else:
             logger.info(f"Loading {config_file_name} for model {oci_model.id}...")
             artifact_path = get_artifact_path(oci_model.custom_metadata_list)
-            config_path = f"{artifact_path.rstrip('/')}/config/"
+            config_path = f"{artifact_path.rstrip('/')}/{config_folder}/"
 
         if not artifact_path:
             logger.debug(
@@ -340,55 +348,6 @@ def get_config(self, model_id: str, config_file_name: str) -> Dict:
 
         return config
 
-    def get_chat_template(self, model_id):
-        """Gets the default chat template for the given Aqua model.
-
-        Parameters
-        ----------
-        model_id: str
-            The OCID of the Aqua model.
-
-        Returns
-        -------
-        str:
-            Chat template string.
-        """
-        chat_template = ""
-        oci_model = self.ds_client.get_model(model_id).data
-        oci_aqua = (
-            (
-                Tags.AQUA_TAG in oci_model.freeform_tags
-                or Tags.AQUA_TAG.lower() in oci_model.freeform_tags
-            )
-            if oci_model.freeform_tags
-            else False
-        )
-
-        if not oci_aqua:
-            raise AquaRuntimeError(f"Target model {oci_model.id} is not Aqua model.")
-        artifact_path = get_artifact_path(oci_model.custom_metadata_list)
-        if not artifact_path:
-            logger.debug(
-                f"Failed to get artifact path from custom metadata for the model: {model_id}"
-            )
-            return chat_template
-
-        try:
-            tokenizer_path = f"{os.path.dirname(artifact_path)}/artifact"
-            chat_template = load_config(
-                file_path=tokenizer_path, config_file_name="tokenizer_config.json"
-            )
-        except Exception:
-            logger.error(
-                f"Error reading tokenizer_config.json file for the model: {model_id}"
-            )
-
-        if not chat_template:
-            logger.error(
-                f"No default chat template is available for the model: {model_id}."
-            )
-        return {"chat_template": chat_template.get("chat_template")}
-
     @property
     def telemetry(self):
         if not self._telemetry:

ads/aqua/constants.py

@@ -10,6 +10,7 @@
 README = "README.md"
 LICENSE_TXT = "config/LICENSE.txt"
 DEPLOYMENT_CONFIG = "deployment_config.json"
+AQUA_MODEL_TOKENIZER_CONFIG = "tokenizer_config.json"
 COMPARTMENT_MAPPING_KEY = "service-model-compartment"
 CONTAINER_INDEX = "container_index.json"
 EVALUATION_REPORT_JSON = "report.json"

ads/aqua/extension/model_handler.py

@@ -317,7 +317,7 @@ def post(self, *args, **kwargs):  # noqa: ARG002
         )
 
 
-class AquaModelChatTemplateHandler(AquaAPIhandler):
+class AquaModelTokenizerConfigHandler(AquaAPIhandler):
     def get(self, model_id):
         url_parse = urlparse(self.request.path)
         paths = url_parse.path.strip("/")
@@ -327,14 +327,14 @@ def get(self, model_id):
             and is_valid_ocid(path_list[2])
             and path_list[3] == "chat_template"
         ):
-            return self.finish(AquaModelApp().get_chat_template(model_id))
+            return self.finish(AquaModelApp().get_hf_tokenizer_config(model_id))
         else:
             raise HTTPError(400, f"The request {self.request.path} is invalid.")
 
 
 __handlers__ = [
     ("model/?([^/]*)", AquaModelHandler),
     ("model/?([^/]*)/license", AquaModelLicenseHandler),
-    ("model/?([^/]*)/chat_template", AquaModelChatTemplateHandler),
+    ("model/?([^/]*)/chat_template", AquaModelTokenizerConfigHandler),
     ("model/hf/search/?([^/]*)", AquaHuggingFaceHandler),
 ]

ads/aqua/model/model.py

@@ -44,6 +44,7 @@
     AQUA_MODEL_ARTIFACT_CONFIG_MODEL_NAME,
     AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE,
     AQUA_MODEL_ARTIFACT_FILE,
+    AQUA_MODEL_TOKENIZER_CONFIG,
     AQUA_MODEL_TYPE_CUSTOM,
     HF_METADATA_FOLDER,
     LICENSE_TXT,
@@ -568,6 +569,24 @@ def _build_ft_metrics(
             training_final,
         ]
 
+    def get_hf_tokenizer_config(self, model_id):
+        """Gets the default chat template for the given Aqua model.
+
+        Parameters
+        ----------
+        model_id: str
+            The OCID of the Aqua model.
+
+        Returns
+        -------
+        str:
+            Chat template string.
+        """
+        config = self.get_config(model_id, AQUA_MODEL_TOKENIZER_CONFIG, "artifact")
+        if not config:
+            logger.debug(f"Tokenizer config for model: {model_id} is not available.")
+        return config
+
     @staticmethod
     def to_aqua_model(
         model: Union[

tests/unitary/with_extras/aqua/test_model_handler.py

@@ -9,7 +9,7 @@
 import pytest
 from huggingface_hub.hf_api import HfApi, ModelInfo
 from huggingface_hub.utils import GatedRepoError
-from notebook.base.handlers import IPythonHandler
+from notebook.base.handlers import IPythonHandler, HTTPError
 from parameterized import parameterized
 
 from ads.aqua.common.errors import AquaRuntimeError
@@ -18,6 +18,7 @@
     AquaHuggingFaceHandler,
     AquaModelHandler,
     AquaModelLicenseHandler,
+    AquaModelTokenizerConfigHandler,
 )
 from ads.aqua.model import AquaModelApp
 from ads.aqua.model.entities import AquaModel, AquaModelSummary, HFModelSummary
@@ -250,6 +251,41 @@ def test_get(self, mock_load_license):
         mock_load_license.assert_called_with("test_model_id")
 
 
+class ModelTokenizerConfigHandlerTestCase(TestCase):
+    @patch.object(IPythonHandler, "__init__")
+    def setUp(self, ipython_init_mock) -> None:
+        ipython_init_mock.return_value = None
+        self.model_tokenizer_config_handler = AquaModelTokenizerConfigHandler(
+            MagicMock(), MagicMock()
+        )
+        self.model_tokenizer_config_handler.finish = MagicMock()
+        self.model_tokenizer_config_handler.request = MagicMock()
+
+    @patch.object(AquaModelApp, "get_hf_tokenizer_config")
+    @patch("ads.aqua.extension.model_handler.urlparse")
+    def test_get(self, mock_urlparse, mock_get_hf_tokenizer_config):
+        request_path = MagicMock(path="aqua/model/ocid1.xx./chat_template")
+        mock_urlparse.return_value = request_path
+        self.model_tokenizer_config_handler.get(model_id="test_model_id")
+        self.model_tokenizer_config_handler.finish.assert_called_with(
+            mock_get_hf_tokenizer_config.return_value
+        )
+        mock_get_hf_tokenizer_config.assert_called_with("test_model_id")
+
+    @patch.object(AquaModelApp, "get_hf_tokenizer_config")
+    @patch("ads.aqua.extension.model_handler.urlparse")
+    def test_get_invalid_path(self, mock_urlparse, mock_get_hf_tokenizer_config):
+        """Test invalid request path should raise HTTPError(400)"""
+        request_path = MagicMock(path="/invalid/path")
+        mock_urlparse.return_value = request_path
+
+        with self.assertRaises(HTTPError) as context:
+            self.model_tokenizer_config_handler.get(model_id="test_model_id")
+        self.assertEqual(context.exception.status_code, 400)
+        self.model_tokenizer_config_handler.finish.assert_not_called()
+        mock_get_hf_tokenizer_config.assert_not_called()
+
+
 class TestAquaHuggingFaceHandler:
     def setup_method(self):
         with patch.object(IPythonHandler, "__init__"):