Repository: huggingface/huggingface_hub
Pull request #1507 (2023-06-13 17:46:36)

Custom headers/cookies in InferenceClient
Related to https://github.com/huggingface/api-inference-community/pull/287 (and especially https://github.com/huggingface/api-inference-community/pull/287#issuecomment-1584578637) cc @StephenHodgson @Narsil 

With this PR, user can set custom headers and cookies to the `InferenceClient` session. This can prove useful when using custom endpoints. If set, custom headers overwrite the default user-agent and authorization headers.

(+fix documentation => save audio as .flac instead of .wav)

HuggingFaceDocBuilderDev: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/huggingface_hub/pr_1507). All of your documentation changes will be reflected on that endpoint.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1507?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) Report
Patch coverage: **`100.00`**% and project coverage change: **`+0.01`** :tada:
> Comparison is base [(`b4e01b0`)](https://app.codecov.io/gh/huggingface/huggingface_hub/commit/b4e01b054d838a3b3763b66e5ed099260a0f8d29?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) 82.18% compared to head [(`b3db890`)](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1507?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) 82.19%.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #1507      +/-   ##
==========================================
+ Coverage   82.18%   82.19%   +0.01%     
==========================================
  Files          57       57              
  Lines        6089     6093       +4     
==========================================
+ Hits         5004     5008       +4     
  Misses       1085     1085              
```


| [Impacted Files](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1507?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) | Coverage Δ | |
|---|---|---|
| [src/huggingface\_hub/\_inference.py](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1507?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface#diff-c3JjL2h1Z2dpbmdmYWNlX2h1Yi9faW5mZXJlbmNlLnB5) | `77.04% <100.00%> (+0.51%)` | :arrow_up: |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1507?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface).


Patch:
diff --git a/docs/source/guides/inference.mdx b/docs/source/guides/inference.mdx
index 5b7c9aea..5cb59e30 100644
--- a/docs/source/guides/inference.mdx
+++ b/docs/source/guides/inference.mdx
@@ -198,7 +198,7 @@ You can catch it and handle it in your code:
 Some tasks require binary inputs, for example, when dealing with images or audio files. In this case, [`InferenceClient`]
 tries to be as permissive as possible and accept different types:
 - raw `bytes`
-- a file-like object, opened as binary (`with open("audio.wav", "rb") as f: ...`)
+- a file-like object, opened as binary (`with open("audio.flac", "rb") as f: ...`)
 - a path (`str` or `Path`) pointing to a local file
 - a URL (`str`) pointing to a remote file (e.g. `https://...`). In this case, the file will be downloaded locally before
 sending it to the Inference API.
diff --git a/src/huggingface_hub/_inference.py b/src/huggingface_hub/_inference.py
index 4aab6568..5fa3a471 100644
--- a/src/huggingface_hub/_inference.py
+++ b/src/huggingface_hub/_inference.py
@@ -44,6 +44,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, BinaryIO, ContextManager, Dict, Generator, List, Optional, Union, overload
 
 from requests import HTTPError, Response
+from requests.structures import CaseInsensitiveDict
 
 from ._inference_types import ClassificationOutput, ConversationalOutput, ImageSegmentationOutput
 from .constants import INFERENCE_ENDPOINT
@@ -96,17 +97,31 @@ class InferenceClient:
             or a URL to a deployed Inference Endpoint. Defaults to None, in which case a recommended model is
             automatically selected for the task.
         token (`str`, *optional*):
-            Hugging Face token. Will default to the locally saved token.
+            Hugging Face token. Will default to the locally saved token. Pass `token=False` if you don't want to send
+            your token to the server.
         timeout (`float`, `optional`):
             The maximum number of seconds to wait for a response from the server. Loading a new model in Inference
             API can take up to several minutes. Defaults to None, meaning it will loop until the server is available.
+        headers (`Dict[str, str]`, `optional`):
+            Additional headers to send to the server. By default only the authorization and user-agent headers are sent.
+            Values in this dictionary will override the default values.
+        cookies (`Dict[str, str]`, `optional`):
+            Additional cookies to send to the server.
     """
 
     def __init__(
-        self, model: Optional[str] = None, token: Optional[str] = None, timeout: Optional[float] = None
+        self,
+        model: Optional[str] = None,
+        token: Union[str, bool, None] = None,
+        timeout: Optional[float] = None,
+        headers: Optional[Dict[str, str]] = None,
+        cookies: Optional[Dict[str, str]] = None,
     ) -> None:
         self.model: Optional[str] = model
-        self.headers = build_hf_headers(token=token)
+        self.headers = CaseInsensitiveDict(build_hf_headers(token=token))  # contains 'authorization' + 'user-agent'
+        if headers is not None:
+            self.headers.update(headers)
+        self.cookies = cookies
         self.timeout = timeout
 
     def __repr__(self):
@@ -157,7 +172,12 @@ class InferenceClient:
             with _open_as_binary(data) as data_as_binary:
                 try:
                     response = get_session().post(
-                        url, json=json, data=data_as_binary, headers=self.headers, timeout=self.timeout
+                        url,
+                        json=json,
+                        data=data_as_binary,
+                        headers=self.headers,
+                        cookies=self.cookies,
+                        timeout=self.timeout,
                     )
                 except TimeoutError as error:
                     # Convert any `TimeoutError` to a `InferenceTimeoutError`
@@ -214,7 +234,7 @@ class InferenceClient:
         ```py
         >>> from huggingface_hub import InferenceClient
         >>> client = InferenceClient()
-        >>> client.audio_classification("audio.wav")
+        >>> client.audio_classification("audio.flac")
         [{'score': 0.4976358711719513, 'label': 'hap'}, {'score': 0.3677836060523987, 'label': 'neu'},...]
         ```
         """
@@ -250,7 +270,7 @@ class InferenceClient:
         ```py
         >>> from huggingface_hub import InferenceClient
         >>> client = InferenceClient()
-        >>> client.automatic_speech_recognition("hello_world.wav")
+        >>> client.automatic_speech_recognition("hello_world.flac")
         "hello world"
         ```
         """
@@ -760,7 +780,7 @@ class InferenceClient:
         >>> client = InferenceClient()
 
         >>> audio = client.text_to_speech("Hello world")
-        >>> Path("hello_world.wav").write_bytes(audio)
+        >>> Path("hello_world.flac").write_bytes(audio)
         ```
         """
         response = self.post(json={"inputs": text}, model=model, task="text-to-speech")


Test patch:
diff --git a/tests/test_inference_client.py b/tests/test_inference_client.py
index 7bf57d81..8b0b45b7 100644
--- a/tests/test_inference_client.py
+++ b/tests/test_inference_client.py
@@ -14,6 +14,7 @@
 import io
 import unittest
 from pathlib import Path
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
@@ -21,6 +22,7 @@ from PIL import Image
 
 from huggingface_hub import InferenceClient, hf_hub_download
 from huggingface_hub._inference import _open_as_binary
+from huggingface_hub.utils import build_hf_headers
 
 from .testing_utils import with_production_testing
 
@@ -251,3 +253,37 @@ class TestResolveURL(InferenceClientTest):
     def test_unsupported_task(self) -> None:
         with self.assertRaises(NotImplementedError):
             InferenceClient()._resolve_url(task="unknown-task")
+
+
+class TestHeadersAndCookies(unittest.TestCase):
+    def test_headers_and_cookies(self) -> None:
+        client = InferenceClient(headers={"X-My-Header": "foo"}, cookies={"my-cookie": "bar"})
+        self.assertEqual(client.headers["X-My-Header"], "foo")
+        self.assertEqual(client.cookies["my-cookie"], "bar")
+
+    def test_headers_overwrite(self) -> None:
+        # Default user agent
+        self.assertTrue(InferenceClient().headers["user-agent"].startswith("unknown/None;"))
+
+        # Overwritten user-agent
+        self.assertEqual(InferenceClient(headers={"user-agent": "bar"}).headers["user-agent"], "bar")
+
+        # Case-insensitive overwrite
+        self.assertEqual(InferenceClient(headers={"USER-agent": "bar"}).headers["user-agent"], "bar")
+
+    @patch("huggingface_hub._inference.get_session")
+    def test_mocked_post(self, get_session_mock: MagicMock) -> None:
+        """Test that headers and cookies are correctly passed to the request."""
+        client = InferenceClient(headers={"X-My-Header": "foo"}, cookies={"my-cookie": "bar"})
+        response = client.post(data=b"content", model="username/repo_name")
+        self.assertEqual(response, get_session_mock().post.return_value)
+
+        expected_user_agent = build_hf_headers()["user-agent"]
+        get_session_mock().post.assert_called_once_with(
+            "https://api-inference.huggingface.co/models/username/repo_name",
+            json=None,
+            data=b"content",
+            headers={"user-agent": expected_user_agent, "X-My-Header": "foo"},
+            cookies={"my-cookie": "bar"},
+            timeout=None,
+        )