Add Hugging Face model support to Shape Recommender

Author: Aryanag2

ads/aqua/shaperecommend/recommend.py

@@ -3,7 +3,11 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import shutil
-from typing import List, Union
+import os
+import re
+import json
+import requests
+from typing import List, Union, Optional, Dict, Any
 
 from pydantic import ValidationError
 from rich.table import Table
@@ -42,6 +46,45 @@
     OCIDataScienceModelDeployment,
 )
 
+class HuggingFaceModelFetcher:
+    """
+    Utility class to fetch model configurations from HuggingFace.
+    """
+    HUGGINGFACE_CONFIG_URL = "https://huggingface.co/{model_id}/resolve/main/config.json"
+
+    @classmethod
+    def is_huggingface_model_id(cls, model_id: str) -> bool:
+        if model_id.startswith("ocid1."):
+            return False
+        hf_pattern = r'^[a-zA-Z0-9_-]+(/[a-zA-Z0-9_.-]+)?$'
+        return bool(re.match(hf_pattern, model_id))
+
+    @classmethod
+    def get_hf_token(cls) -> Optional[str]:
+        return os.environ.get("HUGGING_FACE_HUB_TOKEN") or os.environ.get("HF_TOKEN")
+
+    @classmethod
+    def fetch_config_only(cls, model_id: str) -> Dict[str, Any]:
+        try:
+            config_url = cls.HUGGINGFACE_CONFIG_URL.format(model_id=model_id)
+            headers = {}
+            token = cls.get_hf_token()
+            if token:
+                headers["Authorization"] = f"Bearer {token}"
+            response = requests.get(config_url, headers=headers, timeout=10)
+            if response.status_code == 401:
+                raise AquaValueError(
+                    f"Model '{model_id}' requires authentication. Please set your HuggingFace token."
+                )
+            elif response.status_code == 404:
+                raise AquaValueError(f"Model '{model_id}' not found on HuggingFace.")
+            elif response.status_code != 200:
+                raise AquaValueError(f"Failed to fetch config for '{model_id}'. Status: {response.status_code}")
+            return response.json()
+        except requests.RequestException as e:
+            raise AquaValueError(f"Network error fetching config for {model_id}: {e}") from e
+        except json.JSONDecodeError as e:
+            raise AquaValueError(f"Invalid config format for model '{model_id}'.") from e
 
 class AquaShapeRecommend:
     """
@@ -91,14 +134,8 @@ def which_shapes(
         """
         try:
             shapes = self.valid_compute_shapes(compartment_id=request.compartment_id)
-
-            ds_model = self._validate_model_ocid(request.model_id)
-            data = self._get_model_config(ds_model)
-
+            data, model_name = self._get_model_config_and_name(request.model_id, request.compartment_id)
             llm_config = LLMConfig.from_raw_config(data)
-
-            model_name = ds_model.display_name if ds_model.display_name else ""
-
             shape_recommendation_report = self._summarize_shapes_for_seq_lens(
                 llm_config, shapes, model_name
             )
@@ -127,6 +164,39 @@ def which_shapes(
 
         return shape_recommendation_report
 
+    def _get_model_config_and_name(self, model_id: str, compartment_id: str) -> (dict, str):
+        """
+        Loads model configuration, handling OCID and Hugging Face model IDs.
+        """
+        if HuggingFaceModelFetcher.is_huggingface_model_id(model_id):
+            logger.info(f"'{model_id}' identified as a Hugging Face model ID.")
+            ds_model = self._search_model_in_catalog(model_id, compartment_id)
+            if ds_model and ds_model.artifact:
+                logger.info("Loading configuration from existing model catalog artifact.")
+                try:
+                    return load_config(ds_model.artifact, "config.json"), ds_model.display_name
+                except AquaFileNotFoundError:
+                    logger.warning("config.json not found in artifact, fetching from Hugging Face Hub.")
+            return HuggingFaceModelFetcher.fetch_config_only(model_id), model_id
+        else:
+            logger.info(f"'{model_id}' identified as a model OCID.")
+            ds_model = self._validate_model_ocid(model_id)
+            return self._get_model_config(ds_model), ds_model.display_name
+
+    def _search_model_in_catalog(self, model_id: str, compartment_id: str) -> Optional[DataScienceModel]:
+        """
+        Searches for a Hugging Face model in the Data Science model catalog by display name.
+        """
+        try:
+            # This should work since the SDK's list method can filter by display_name.
+            models = DataScienceModel.list(compartment_id=compartment_id, display_name=model_id)
+            if models:
+                logger.info(f"Found model '{model_id}' in the Data Science catalog.")
+                return models[0]
+        except Exception as e:
+            logger.warning(f"Could not search for model '{model_id}' in catalog: {e}")
+        return None
+
     def valid_compute_shapes(self, compartment_id: str) -> List["ComputeShapeSummary"]:
         """
         Returns a filtered list of GPU-only ComputeShapeSummary objects by reading and parsing a JSON file.

ads/aqua/shaperecommend/shape_report.py

@@ -18,7 +18,7 @@ class RequestRecommend(BaseModel):
     """
 
     model_id: str = Field(
-        ..., description="The OCID of the model to recommend feasible compute shapes."
+        ..., description="The OCID or Hugging Face ID of the model to recommend feasible compute shapes."
     )
     generate_table: Optional[bool] = (
         Field(

tests/unitary/with_extras/aqua/test_recommend.py

@@ -8,19 +8,19 @@
 import os
 import re
 from unittest.mock import MagicMock
-
+from unittest.mock import patch
 import pytest
 
 from ads.aqua.common.entities import ComputeShapeSummary
-from ads.aqua.common.errors import AquaRecommendationError
+from ads.aqua.common.errors import AquaRecommendationError, AquaValueError
 from ads.aqua.shaperecommend.estimator import (
     LlamaMemoryEstimator,
     MemoryEstimator,
     MixtureMemoryEstimator,
     get_estimator,
 )
 from ads.aqua.shaperecommend.llm_config import LLMConfig
-from ads.aqua.shaperecommend.recommend import AquaShapeRecommend
+from ads.aqua.shaperecommend.recommend import AquaShapeRecommend, HuggingFaceModelFetcher
 from ads.aqua.shaperecommend.shape_report import (
     DeploymentParams,
     ModelConfig,
@@ -454,3 +454,56 @@ def test_shape_report_pareto_front(self):
         assert c and d in pf
         assert a and b not in pf
         assert len(pf) == 2
+# --- Tests for HuggingFaceModelFetcher ---
+class TestHuggingFaceModelFetcher:
+    @pytest.mark.parametrize("model_id, expected", [
+        ("meta-llama/Llama-2-7b-hf", True),
+        ("mistralai/Mistral-7B-v0.1", True),
+        ("ocid1.datasciencemodel.oc1.iad.xxxxxxxx", False),
+    ])
+    def test_is_huggingface_model_id(self, model_id, expected):
+        assert HuggingFaceModelFetcher.is_huggingface_model_id(model_id) == expected
+
+    @patch('requests.get')
+    def test_fetch_config_only_success(self, mock_get):
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {"model_type": "llama"}
+        mock_get.return_value = mock_response
+
+        config = HuggingFaceModelFetcher.fetch_config_only("some/model")
+        assert config == {"model_type": "llama"}
+        mock_get.assert_called_once()
+
+    @patch('requests.get')
+    def test_fetch_config_only_not_found(self, mock_get):
+        mock_response = MagicMock()
+        mock_response.status_code = 404
+        mock_get.return_value = mock_response
+
+        with pytest.raises(AquaValueError, match="not found on HuggingFace"):
+            HuggingFaceModelFetcher.fetch_config_only("non/existent")
+
+    @patch.dict(os.environ, {"HF_TOKEN": "test_token_123"}, clear=True)
+    def test_get_hf_token(self):
+        assert HuggingFaceModelFetcher.get_hf_token() == "test_token_123"
+        
+    # Add this method inside the TestHuggingFaceModelFetcher class
+
+    @pytest.mark.network
+    def test_fetch_config_only_real_call_success(self):
+        """
+        Tests a real network call to fetch a public model's configuration.
+        This test requires an internet connection.
+        """
+        # Use a well-known, small, public model to minimize test flakiness
+        model_id = "distilbert-base-uncased"
+        
+        try:
+            config = HuggingFaceModelFetcher.fetch_config_only(model_id)
+            # Assert that we got a dictionary with expected keys
+            assert isinstance(config, dict)
+            assert "model_type" in config
+            assert "dim" in config
+        except AquaValueError as e:
+            pytest.fail(f"Real network call to Hugging Face failed: {e}")