[Inference Providers] sambanova supports feature extraction (#3037)

* add feature extraction for sambanova

* nit

* update table

* hf inference feature extraction

* fix linter

* fix

Author: hanouticelina

docs/source/en/guides/inference.md

@@ -191,7 +191,7 @@ For more details, refer to the [Inference Providers pricing documentation](https
 | [`~InferenceClient.automatic_speech_recognition`]   | ❌                 | ❌        | ❌      | ✅      | ❌            | ✅            | ❌          | ❌                | ❌         | ❌         | ❌         | ❌        |
 | [`~InferenceClient.chat_completion`]                | ❌                 | ✅        | ✅      | ❌      | ✅            | ✅            | ✅          | ✅                | ✅         | ❌         | ✅         | ✅        |
 | [`~InferenceClient.document_question_answering`]    | ❌                 | ❌        | ❌      | ❌      | ❌            | ✅            | ❌          | ❌                | ❌         | ❌         | ❌         | ❌        |
-| [`~InferenceClient.feature_extraction`]             | ❌                 | ❌        | ❌      | ❌      | ❌            | ✅            | ❌          | ❌                | ❌         | ❌         | ❌         | ❌        |
+| [`~InferenceClient.feature_extraction`]             | ❌                 | ❌        | ❌      | ❌      | ❌            | ✅            | ❌          | ❌                | ❌         | ❌         | ✅         | ❌        |
 | [`~InferenceClient.fill_mask`]                      | ❌                 | ❌        | ❌      | ❌      | ❌            | ✅            | ❌          | ❌                | ❌         | ❌         | ❌         | ❌        |
 | [`~InferenceClient.image_classification`]           | ❌                 | ❌        | ❌      | ❌      | ❌            | ✅            | ❌          | ❌                | ❌         | ❌         | ❌         | ❌        |
 | [`~InferenceClient.image_segmentation`]             | ❌                 | ❌        | ❌      | ❌      | ❌            | ✅            | ❌          | ❌                | ❌         | ❌         | ❌         | ❌        |

src/huggingface_hub/inference/_client.py

@@ -1080,7 +1080,7 @@ def feature_extraction(
         )
         response = self._inner_post(request_parameters)
         np = _import_numpy()
-        return np.array(_bytes_to_dict(response), dtype="float32")
+        return np.array(provider_helper.get_response(response), dtype="float32")
 
     def fill_mask(
         self,

src/huggingface_hub/inference/_generated/_async_client.py

@@ -1122,7 +1122,7 @@ async def feature_extraction(
         )
         response = await self._inner_post(request_parameters)
         np = _import_numpy()
-        return np.array(_bytes_to_dict(response), dtype="float32")
+        return np.array(provider_helper.get_response(response), dtype="float32")
 
     async def fill_mask(
         self,

src/huggingface_hub/inference/_providers/__init__.py

@@ -13,13 +13,18 @@
     FalAITextToVideoTask,
 )
 from .fireworks_ai import FireworksAIConversationalTask
-from .hf_inference import HFInferenceBinaryInputTask, HFInferenceConversational, HFInferenceTask
+from .hf_inference import (
+    HFInferenceBinaryInputTask,
+    HFInferenceConversational,
+    HFInferenceFeatureExtractionTask,
+    HFInferenceTask,
+)
 from .hyperbolic import HyperbolicTextGenerationTask, HyperbolicTextToImageTask
 from .nebius import NebiusConversationalTask, NebiusTextGenerationTask, NebiusTextToImageTask
 from .novita import NovitaConversationalTask, NovitaTextGenerationTask, NovitaTextToVideoTask
 from .openai import OpenAIConversationalTask
 from .replicate import ReplicateTask, ReplicateTextToSpeechTask
-from .sambanova import SambanovaConversationalTask
+from .sambanova import SambanovaConversationalTask, SambanovaFeatureExtractionTask
 from .together import TogetherConversationalTask, TogetherTextGenerationTask, TogetherTextToImageTask
 
 
@@ -72,7 +77,7 @@
         "audio-classification": HFInferenceBinaryInputTask("audio-classification"),
         "automatic-speech-recognition": HFInferenceBinaryInputTask("automatic-speech-recognition"),
         "fill-mask": HFInferenceTask("fill-mask"),
-        "feature-extraction": HFInferenceTask("feature-extraction"),
+        "feature-extraction": HFInferenceFeatureExtractionTask(),
         "image-classification": HFInferenceBinaryInputTask("image-classification"),
         "image-segmentation": HFInferenceBinaryInputTask("image-segmentation"),
         "document-question-answering": HFInferenceTask("document-question-answering"),
@@ -116,6 +121,7 @@
     },
     "sambanova": {
         "conversational": SambanovaConversationalTask(),
+        "feature-extraction": SambanovaFeatureExtractionTask(),
     },
     "together": {
         "text-to-image": TogetherTextToImageTask(),

src/huggingface_hub/inference/_providers/hf_inference.py

@@ -1,11 +1,11 @@
 import json
 from functools import lru_cache
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Union
 
 from huggingface_hub import constants
 from huggingface_hub.hf_api import InferenceProviderMapping
-from huggingface_hub.inference._common import _b64_encode, _open_as_binary
+from huggingface_hub.inference._common import RequestParameters, _b64_encode, _bytes_to_dict, _open_as_binary
 from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none
 from huggingface_hub.utils import build_hf_headers, get_session, get_token, hf_raise_for_status
 
@@ -177,3 +177,13 @@ def _check_supported_task(model: str, task: str) -> None:
             f"Model '{model}' doesn't support task '{task}'. Supported tasks: '{pipeline_tag}', got: '{task}'"
         )
     return
+
+
+class HFInferenceFeatureExtractionTask(HFInferenceTask):
+    def __init__(self):
+        super().__init__("feature-extraction")
+
+    def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
+        if isinstance(response, bytes):
+            return _bytes_to_dict(response)
+        return response

src/huggingface_hub/inference/_providers/sambanova.py

@@ -1,6 +1,28 @@
-from huggingface_hub.inference._providers._common import BaseConversationalTask
+from typing import Any, Dict, Optional, Union
+
+from huggingface_hub.hf_api import InferenceProviderMapping
+from huggingface_hub.inference._common import RequestParameters, _as_dict
+from huggingface_hub.inference._providers._common import BaseConversationalTask, TaskProviderHelper, filter_none
 
 
 class SambanovaConversationalTask(BaseConversationalTask):
     def __init__(self):
         super().__init__(provider="sambanova", base_url="https://api.sambanova.ai")
+
+
+class SambanovaFeatureExtractionTask(TaskProviderHelper):
+    def __init__(self):
+        super().__init__(provider="sambanova", base_url="https://api.sambanova.ai", task="feature-extraction")
+
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
+        return "/v1/embeddings"
+
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[Dict]:
+        parameters = filter_none(parameters)
+        return {"input": inputs, "model": provider_mapping_info.provider_id, **parameters}
+
+    def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
+        embeddings = _as_dict(response)["data"]
+        return [embedding["embedding"] for embedding in embeddings]

tests/test_inference_providers.py

@@ -35,7 +35,7 @@
 from huggingface_hub.inference._providers.novita import NovitaConversationalTask, NovitaTextGenerationTask
 from huggingface_hub.inference._providers.openai import OpenAIConversationalTask
 from huggingface_hub.inference._providers.replicate import ReplicateTask, ReplicateTextToSpeechTask
-from huggingface_hub.inference._providers.sambanova import SambanovaConversationalTask
+from huggingface_hub.inference._providers.sambanova import SambanovaConversationalTask, SambanovaFeatureExtractionTask
 from huggingface_hub.inference._providers.together import TogetherTextToImageTask
 
 from .testing_utils import assert_in_logs
@@ -903,13 +903,34 @@ def test_get_response_single_output(self, mocker):
 
 
 class TestSambanovaProvider:
-    def test_prepare_url(self):
+    def test_prepare_url_conversational(self):
         helper = SambanovaConversationalTask()
         assert (
             helper._prepare_url("sambanova_token", "username/repo_name")
             == "https://api.sambanova.ai/v1/chat/completions"
         )
 
+    def test_prepare_payload_as_dict_feature_extraction(self):
+        helper = SambanovaFeatureExtractionTask()
+        payload = helper._prepare_payload_as_dict(
+            "Hello world",
+            {"truncate": True},
+            InferenceProviderMapping(
+                hf_model_id="username/repo_name",
+                providerId="provider-id",
+                task="feature-extraction",
+                status="live",
+            ),
+        )
+        assert payload == {"input": "Hello world", "model": "provider-id", "truncate": True}
+
+    def test_prepare_url_feature_extraction(self):
+        helper = SambanovaFeatureExtractionTask()
+        assert (
+            helper._prepare_url("hf_token", "username/repo_name")
+            == "https://router.huggingface.co/sambanova/v1/embeddings"
+        )
+
 
 class TestTogetherProvider:
     def test_prepare_route_text_to_image(self):