HF look ahead search

Author: kumar-shivam-ranjan

ads/aqua/common/utils.py

@@ -1064,4 +1064,12 @@ def get_hf_model_info(repo_id: str) -> ModelInfo:
         raise format_hf_custom_error_message(err) from err
 
 
+@cached(cache=TTLCache(maxsize=1, ttl=timedelta(hours=5), timer=datetime.now))
+def list_hf_models(query:str) -> List[str]:
+    try:
+        models= HfApi().list_models(model_name=query,task="text-generation")
+        return [model.id for model in models if model.disabled is None]
+    except HfHubHTTPError as err:
+        raise format_hf_custom_error_message(err) from err
+
 

ads/aqua/extension/model_handler.py

@@ -9,7 +9,7 @@
 
 from ads.aqua.common.decorator import handle_exceptions
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
-from ads.aqua.common.utils import get_hf_model_info
+from ads.aqua.common.utils import get_hf_model_info, list_hf_models
 from ads.aqua.extension.base_handler import AquaAPIhandler
 from ads.aqua.extension.errors import Errors
 from ads.aqua.model import AquaModelApp
@@ -177,6 +177,31 @@ def _find_matching_aqua_model(model_id: str) -> Optional[AquaModelSummary]:
 
         return None
 
+
+
+    @handle_exceptions
+    def get(self,*args, **kwargs):
+        """
+        Finds a list of matching models from hugging face based on query string provided from users.
+
+        Parameters
+        ----------
+        query (str): The Hugging Face model name to search for.
+
+        Returns
+        -------
+        List[AquaModelSummary]
+            Returns the matching AquaModelSummary object if found, else 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})
+
+
+
     @handle_exceptions
     def post(self, *args, **kwargs):
         """Handles post request for the HF Models APIs