modified unit tests

Author: elizjo

ads/aqua/shaperecommend/recommend.py

@@ -4,14 +4,14 @@
 
 import json
 import os
+import re
 import shutil
 from typing import Dict, List, Optional, Tuple, Union
 
-from pydantic import ValidationError
-from rich.table import Table
-
 from huggingface_hub import hf_hub_download
 from huggingface_hub.utils import HfHubHTTPError
+from pydantic import ValidationError
+from rich.table import Table
 
 from ads.aqua.app import logger
 from ads.aqua.common.entities import ComputeShapeSummary
@@ -22,15 +22,15 @@
 )
 from ads.aqua.common.utils import (
     build_pydantic_error_message,
+    format_hf_custom_error_message,
     get_resource_type,
     is_valid_ocid,
     load_config,
     load_gpu_shapes_index,
-    format_hf_custom_error_message,
 )
 from ads.aqua.shaperecommend.constants import (
-    BITSANDBYTES,
     BITS_AND_BYTES_4BIT,
+    BITSANDBYTES,
     SAFETENSORS,
     SHAPE_MAP,
     TEXT_GENERATION,
@@ -98,14 +98,10 @@ def which_shapes(
         """
         try:
             shapes = self.valid_compute_shapes(compartment_id=request.compartment_id)
-            
+
             data, model_name = self._get_model_config_and_name(
                 model_id=request.model_id,
             )
-            llm_config = LLMConfig.from_raw_config(data)
-            shape_recommendation_report = self._summarize_shapes_for_seq_lens(
-                llm_config, shapes, model_name
-            )
 
             if request.deployment_config:
                 shape_recommendation_report = (
@@ -115,16 +111,11 @@ def which_shapes(
                 )
 
             else:
-                ds_model = self._get_data_science_model(request.model_id)
-                
-                data = self._get_model_config(ds_model)
-
                 llm_config = LLMConfig.from_raw_config(data)
 
                 shape_recommendation_report = self._summarize_shapes_for_seq_lens(
                     llm_config, shapes, model_name
                 )
-                
 
             if request.generate_table and shape_recommendation_report.recommendations:
                 shape_recommendation_report = self._rich_diff_table(
@@ -174,8 +165,8 @@ def _get_model_config_and_name(
         """
         if is_valid_ocid(model_id):
             logger.info(f"Detected OCID: Fetching OCI model config for '{model_id}'.")
-            ds_model = self._validate_model_ocid(model_id)
-            config = self._fetch_hf_config(model_id)
+            ds_model = self._get_data_science_model(model_id)
+            config = self._get_model_config(ds_model)
             model_name = ds_model.display_name
         else:
             logger.info(
@@ -403,6 +394,7 @@ def _get_model_config(model: DataScienceModel):
         """
 
         model_task = model.freeform_tags.get("task", "").lower()
+        model_task =  re.sub(r"-", "_", model_task)
         model_format = model.freeform_tags.get("model_format", "").lower()
 
         logger.info(f"Current model task type: {model_task}")

tests/unitary/with_extras/aqua/test_recommend.py

@@ -436,18 +436,19 @@ def test_which_shapes_valid_from_file(
             )[1],
         )
 
-        raw = load_config(config_file)
+        mock_raw_config = load_config(config_file)
+        mock_ds_model_name = mock_model.display_name
 
         if service_managed_model:
-            config = AquaDeploymentConfig(**raw)
+            config = AquaDeploymentConfig(**mock_raw_config)
 
             request = RequestRecommend(
                 model_id="ocid1.datasciencemodel.oc1.TEST",
                 generate_table=False,
                 deployment_config=config,
             )
         else:
-            monkeypatch.setattr(app, "_get_model_config", lambda _: raw)
+            monkeypatch.setattr(app, "_get_model_config_and_name", lambda _: (mock_ds_model_name, mock_raw_config))
 
             request = RequestRecommend(
                 model_id="ocid1.datasciencemodel.oc1.TEST", generate_table=False