fix: update payload preparation to merge parameters into the output dictionary (#3160)

* fix: update payload preparation to merge parameters into the output dictionary

* dev change

* Update only feature-extraction task + add test + add specs link

* [dev change]

* Revert "[dev change]"

This reverts commit 50dba39.

---------

Co-authored-by: Lucain Pouget <[email protected]>

Author: 2 people

src/huggingface_hub/inference/_providers/hf_inference.py

@@ -194,6 +194,18 @@ class HFInferenceFeatureExtractionTask(HFInferenceTask):
     def __init__(self):
         super().__init__("feature-extraction")
 
+    def _prepare_payload_as_dict(
+        self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
+    ) -> Optional[Dict]:
+        if isinstance(inputs, bytes):
+            raise ValueError(f"Unexpected binary input for task {self.task}.")
+        if isinstance(inputs, Path):
+            raise ValueError(f"Unexpected path input for task {self.task} (got {inputs})")
+
+        # Parameters are sent at root-level for feature-extraction task
+        # See specs: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/tasks/feature-extraction/spec/input.json
+        return {"inputs": inputs, **filter_none(parameters)}
+
     def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
         if isinstance(response, bytes):
             return _bytes_to_dict(response)

tests/test_inference_providers.py

@@ -1,7 +1,7 @@
 import base64
 import logging
 from typing import Dict
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 from pytest import LogCaptureFixture
@@ -33,6 +33,7 @@
 from huggingface_hub.inference._providers.hf_inference import (
     HFInferenceBinaryInputTask,
     HFInferenceConversational,
+    HFInferenceFeatureExtractionTask,
     HFInferenceTask,
 )
 from huggingface_hub.inference._providers.hyperbolic import HyperbolicTextGenerationTask, HyperbolicTextToImageTask
@@ -654,6 +655,15 @@ def test_prepare_payload_as_dict_conversational(self, mapped_model, parameters,
         assert payload["model"] == expected_model
         assert payload["messages"] == messages
 
+    def test_prepare_payload_feature_extraction(self):
+        helper = HFInferenceFeatureExtractionTask()
+        payload = helper._prepare_payload_as_dict(
+            inputs="This is a test sentence.",
+            parameters={"truncate": True},
+            provider_mapping_info=MagicMock(),
+        )
+        assert payload == {"inputs": "This is a test sentence.", "truncate": True}  # not under "parameters"
+
     @pytest.mark.parametrize(
         "pipeline_tag,tags,task,should_raise",
         [