Fix inference search (#3022)

* Fix inference search

* forward compatible inference_provider_mapping

* styling

Author: Wauplin

docs/source/en/package_reference/hf_api.md

@@ -57,6 +57,10 @@ models = hf_api.list_models()
 
 [[autodoc]] huggingface_hub.hf_api.GitRefs
 
+### InferenceProviderMapping
+
+[[autodoc]] huggingface_hub.hf_api.InferenceProviderMapping
+
 ### LFSFileInfo
 
 [[autodoc]] huggingface_hub.hf_api.LFSFileInfo

src/huggingface_hub/hf_api.py

@@ -28,6 +28,7 @@
 from itertools import islice
 from pathlib import Path
 from typing import (
+    TYPE_CHECKING,
     Any,
     BinaryIO,
     Callable,
@@ -135,6 +136,9 @@
 from .utils.endpoint_helpers import _is_emission_within_threshold
 
 
+if TYPE_CHECKING:
+    from .inference._providers import PROVIDER_T
+
 R = TypeVar("R")  # Return type
 CollectionItemType_T = Literal["model", "dataset", "space", "paper", "collection"]
 
@@ -709,21 +713,26 @@ def __init__(self, **kwargs):
 
 @dataclass
 class InferenceProviderMapping:
-    hf_model_id: str
+    provider: "PROVIDER_T"  # Provider name
+    hf_model_id: str  # ID of the model on the Hugging Face Hub
+    provider_id: str  # ID of the model on the provider's side
     status: Literal["live", "staging"]
-    provider_id: str
     task: str
 
     adapter: Optional[str] = None
     adapter_weights_path: Optional[str] = None
+    type: Optional[Literal["single-model", "tag-filter"]] = None
 
     def __init__(self, **kwargs):
+        self.provider = kwargs.pop("provider")
         self.hf_model_id = kwargs.pop("hf_model_id")
-        self.status = kwargs.pop("status")
         self.provider_id = kwargs.pop("providerId")
+        self.status = kwargs.pop("status")
         self.task = kwargs.pop("task")
+
         self.adapter = kwargs.pop("adapter", None)
         self.adapter_weights_path = kwargs.pop("adapterWeightsPath", None)
+        self.type = kwargs.pop("type", None)
         self.__dict__.update(**kwargs)
 
 
@@ -765,12 +774,10 @@ class ModelInfo:
             If so, whether there is manual or automatic approval.
         gguf (`Dict`, *optional*):
             GGUF information of the model.
-        inference (`Literal["cold", "frozen", "warm"]`, *optional*):
-            Status of the model on the inference API.
-            Warm models are available for immediate use. Cold models will be loaded on first inference call.
-            Frozen models are not available in Inference API.
-        inference_provider_mapping (`Dict`, *optional*):
-            Model's inference provider mapping.
+        inference (`Literal["warm"]`, *optional*):
+            Status of the model on Inference Providers. Warm if the model is served by at least one provider.
+        inference_provider_mapping (`List[InferenceProviderMapping]`, *optional*):
+            A list of [`InferenceProviderMapping`] ordered after the user's provider order.
         likes (`int`):
             Number of likes of the model.
         library_name (`str`, *optional*):
@@ -815,8 +822,8 @@ class ModelInfo:
     downloads_all_time: Optional[int]
     gated: Optional[Literal["auto", "manual", False]]
     gguf: Optional[Dict]
-    inference: Optional[Literal["warm", "cold", "frozen"]]
-    inference_provider_mapping: Optional[Dict[str, InferenceProviderMapping]]
+    inference: Optional[Literal["warm"]]
+    inference_provider_mapping: Optional[List[InferenceProviderMapping]]
     likes: Optional[int]
     library_name: Optional[str]
     tags: Optional[List[str]]
@@ -852,14 +859,25 @@ def __init__(self, **kwargs):
         self.gguf = kwargs.pop("gguf", None)
 
         self.inference = kwargs.pop("inference", None)
-        self.inference_provider_mapping = kwargs.pop("inferenceProviderMapping", None)
-        if self.inference_provider_mapping:
-            self.inference_provider_mapping = {
-                provider: InferenceProviderMapping(
-                    **{**value, "hf_model_id": self.id}
-                )  # little hack to simplify Inference Providers logic
-                for provider, value in self.inference_provider_mapping.items()
-            }
+
+        # little hack to simplify Inference Providers logic and make it backward and forward compatible
+        # right now, API returns a dict on model_info and a list on list_models. Let's harmonize to list.
+        mapping = kwargs.pop("inferenceProviderMapping", None)
+        if isinstance(mapping, list):
+            self.inference_provider_mapping = [
+                InferenceProviderMapping(**{**value, "hf_model_id": self.id}) for value in mapping
+            ]
+        elif isinstance(mapping, dict):
+            self.inference_provider_mapping = [
+                InferenceProviderMapping(**{**value, "hf_model_id": self.id, "provider": provider})
+                for provider, value in mapping.items()
+            ]
+        elif mapping is None:
+            self.inference_provider_mapping = None
+        else:
+            raise ValueError(
+                f"Unexpected type for `inferenceProviderMapping`. Expecting `dict` or `list`. Got {mapping}."
+            )
 
         self.tags = kwargs.pop("tags", None)
         self.pipeline_tag = kwargs.pop("pipeline_tag", None)
@@ -1836,7 +1854,8 @@ def list_models(
         filter: Union[str, Iterable[str], None] = None,
         author: Optional[str] = None,
         gated: Optional[bool] = None,
-        inference: Optional[Literal["cold", "frozen", "warm"]] = None,
+        inference: Optional[Literal["warm"]] = None,
+        inference_provider: Optional[Union[Literal["all"], "PROVIDER_T", List["PROVIDER_T"]]] = None,
         library: Optional[Union[str, List[str]]] = None,
         language: Optional[Union[str, List[str]]] = None,
         model_name: Optional[str] = None,
@@ -1870,10 +1889,11 @@ def list_models(
                 A boolean to filter models on the Hub that are gated or not. By default, all models are returned.
                 If `gated=True` is passed, only gated models are returned.
                 If `gated=False` is passed, only non-gated models are returned.
-            inference (`Literal["cold", "frozen", "warm"]`, *optional*):
-                A string to filter models on the Hub by their state on the Inference API.
-                Warm models are available for immediate use. Cold models will be loaded on first inference call.
-                Frozen models are not available in Inference API.
+            inference (`Literal["warm"]`, *optional*):
+                If "warm", filter models on the Hub currently served by at least one provider.
+            inference_provider (`Literal["all"]` or `str`, *optional*):
+                A string to filter models on the Hub that are served by a specific provider.
+                Pass `"all"` to get all models served by at least one provider.
             library (`str` or `List`, *optional*):
                 A string or list of strings of foundational libraries models were
                 originally trained from, such as pytorch, tensorflow, or allennlp.
@@ -1933,7 +1953,7 @@ def list_models(
         Returns:
             `Iterable[ModelInfo]`: an iterable of [`huggingface_hub.hf_api.ModelInfo`] objects.
 
-        Example usage with the `filter` argument:
+        Example:
 
         ```python
         >>> from huggingface_hub import HfApi
@@ -1943,24 +1963,19 @@ def list_models(
         # List all models
         >>> api.list_models()
 
-        # List only the text classification models
+        # List text classification models
         >>> api.list_models(filter="text-classification")
 
-        # List only models from the AllenNLP library
-        >>> api.list_models(filter="allennlp")
-        ```
-
-        Example usage with the `search` argument:
+        # List models from the KerasHub library
+        >>> api.list_models(filter="keras-hub")
 
-        ```python
-        >>> from huggingface_hub import HfApi
-
-        >>> api = HfApi()
+        # List models served by Cohere
+        >>> api.list_models(inference_provider="cohere")
 
-        # List all models with "bert" in their name
+        # List models with "bert" in their name
         >>> api.list_models(search="bert")
 
-        # List all models with "bert" in their name made by google
+        # List models with "bert" in their name and pushed by google
         >>> api.list_models(search="bert", author="google")
         ```
         """
@@ -2003,6 +2018,8 @@ def list_models(
             params["gated"] = gated
         if inference is not None:
             params["inference"] = inference
+        if inference_provider is not None:
+            params["inference_provider"] = inference_provider
         if pipeline_tag:
             params["pipeline_tag"] = pipeline_tag
         search_list = []

src/huggingface_hub/inference/_providers/__init__.py

@@ -183,7 +183,7 @@ def get_provider_helper(
         if model is None:
             raise ValueError("Specifying a model is required when provider is 'auto'")
         provider_mapping = _fetch_inference_provider_mapping(model)
-        provider = next(iter(provider_mapping))
+        provider = next(iter(provider_mapping)).provider
 
     provider_tasks = PROVIDERS.get(provider)  # type: ignore
     if provider_tasks is None:

src/huggingface_hub/inference/_providers/_common.py

@@ -1,5 +1,5 @@
 from functools import lru_cache
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from huggingface_hub import constants
 from huggingface_hub.hf_api import InferenceProviderMapping
@@ -9,6 +9,7 @@
 
 logger = logging.get_logger(__name__)
 
+
 # Dev purposes only.
 # If you want to try to run inference for a new model locally before it's registered on huggingface.co
 # for a given Inference Provider, you can add it to the following dictionary.
@@ -124,7 +125,12 @@ def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMappin
         if HARDCODED_MODEL_INFERENCE_MAPPING.get(self.provider, {}).get(model):
             return HARDCODED_MODEL_INFERENCE_MAPPING[self.provider][model]
 
-        provider_mapping = _fetch_inference_provider_mapping(model).get(self.provider)
+        provider_mapping = None
+        for mapping in _fetch_inference_provider_mapping(model):
+            if mapping.provider == self.provider:
+                provider_mapping = mapping
+                break
+
         if provider_mapping is None:
             raise ValueError(f"Model {model} is not supported by provider {self.provider}.")
 
@@ -236,7 +242,7 @@ def _prepare_payload_as_dict(
 
 
 @lru_cache(maxsize=None)
-def _fetch_inference_provider_mapping(model: str) -> Dict:
+def _fetch_inference_provider_mapping(model: str) -> List["InferenceProviderMapping"]:
     """
     Fetch provider mappings for a model from the Hub.
     """

src/huggingface_hub/inference/_providers/hf_inference.py

@@ -26,15 +26,19 @@ def _prepare_api_key(self, api_key: Optional[str]) -> str:
 
     def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping:
         if model is not None and model.startswith(("http://", "https://")):
-            return InferenceProviderMapping(providerId=model, hf_model_id=model, task=self.task, status="live")
+            return InferenceProviderMapping(
+                provider="hf-inference", providerId=model, hf_model_id=model, task=self.task, status="live"
+            )
         model_id = model if model is not None else _fetch_recommended_models().get(self.task)
         if model_id is None:
             raise ValueError(
                 f"Task {self.task} has no recommended model for HF Inference. Please specify a model"
                 " explicitly. Visit https://huggingface.co/tasks for more info."
             )
         _check_supported_task(model_id, self.task)
-        return InferenceProviderMapping(providerId=model_id, hf_model_id=model_id, task=self.task, status="live")
+        return InferenceProviderMapping(
+            provider="hf-inference", providerId=model_id, hf_model_id=model_id, task=self.task, status="live"
+        )
 
     def _prepare_url(self, api_key: str, mapped_model: str) -> str:
         # hf-inference provider can handle URLs (e.g. Inference Endpoints or TGI deployment)

src/huggingface_hub/inference/_providers/openai.py

@@ -20,4 +20,6 @@ def _prepare_api_key(self, api_key: Optional[str]) -> str:
     def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping:
         if model is None:
             raise ValueError("Please provide an OpenAI model ID, e.g. `gpt-4o` or `o1`.")
-        return InferenceProviderMapping(providerId=model, task="conversational", status="live", hf_model_id=model)
+        return InferenceProviderMapping(
+            provider="openai", providerId=model, task="conversational", status="live", hf_model_id=model
+        )

tests/test_hf_api.py

@@ -60,6 +60,7 @@
     ExpandModelProperty_T,
     ExpandSpaceProperty_T,
     InferenceEndpoint,
+    InferenceProviderMapping,
     ModelInfo,
     RepoSibling,
     RepoUrl,
@@ -2511,6 +2512,38 @@ def test_not_a_safetensors_file(self) -> None:
                 "HuggingFaceH4/zephyr-7b-beta", "pytorch_model-00001-of-00008.bin"
             )
 
+    def test_inference_provider_mapping_model_info(self):
+        model = self._api.model_info("deepseek-ai/DeepSeek-R1-0528", expand="inferenceProviderMapping")
+        mapping = model.inference_provider_mapping
+        assert isinstance(mapping, list)
+        assert len(mapping) > 0
+        for item in mapping:
+            assert isinstance(item, InferenceProviderMapping)
+            assert item.provider is not None
+            assert item.hf_model_id == "deepseek-ai/DeepSeek-R1-0528"
+            assert item.provider_id is not None
+
+    def test_inference_provider_mapping_list_models(self):
+        models = list(self._api.list_models(author="deepseek-ai", expand="inferenceProviderMapping", limit=1))
+        assert len(models) > 0
+        mapping = models[0].inference_provider_mapping
+        assert isinstance(mapping, list)
+        assert len(mapping) > 0
+        for item in mapping:
+            assert isinstance(item, InferenceProviderMapping)
+            assert item.provider is not None
+            assert item.hf_model_id is not None
+            assert item.provider_id is not None
+
+    def test_filter_models_by_inference_provider(self):
+        models = list(
+            self._api.list_models(inference_provider="hf-inference", expand=["inferenceProviderMapping"], limit=10)
+        )
+        assert len(models) > 0
+        for model in models:
+            assert model.inference_provider_mapping is not None
+            assert any(mapping.provider == "hf-inference" for mapping in model.inference_provider_mapping)
+
 
 class HfApiPrivateTest(HfApiCommonTest):
     def setUp(self) -> None: