Repository: huggingface/huggingface_hub
Pull request #1641 (2023-09-05 10:15:16)

Distinguish `_bytes_to_dict` and `_bytes_to_list` + fix issues
This PR fixes an incoherence reported by @martinbrose in https://github.com/huggingface/huggingface_hub/pull/1606#discussion_r1315233260. In the inference client logic, we have a `_bytes_to_dict` helper that is annotated as returning an `"Image"` which is wrong. 

This PR changes the return type of  `_bytes_to_dict` and adds a `_bytes_to_list` helper. Both have the exact same logic to decode JSON-encoded data. The advantage of having 2 helpers is to help devs + mypy understand what is the expected data format from the server. If the data is not in a valid form, no issue is raised at runtime.


Patch:
diff --git a/src/huggingface_hub/inference/_client.py b/src/huggingface_hub/inference/_client.py
index 3e768dbc..308e1739 100644
--- a/src/huggingface_hub/inference/_client.py
+++ b/src/huggingface_hub/inference/_client.py
@@ -62,6 +62,7 @@ from huggingface_hub.inference._common import (
     _b64_to_image,
     _bytes_to_dict,
     _bytes_to_image,
+    _bytes_to_list,
     _get_recommended_model,
     _import_numpy,
     _is_tgi_server,
@@ -284,7 +285,7 @@ class InferenceClient:
         ```
         """
         response = self.post(data=audio, model=model, task="audio-classification")
-        return _bytes_to_dict(response)
+        return _bytes_to_list(response)
 
     def automatic_speech_recognition(
         self,
@@ -381,7 +382,7 @@ class InferenceClient:
         if parameters is not None:
             payload["parameters"] = parameters
         response = self.post(json=payload, model=model, task="conversational")
-        return _bytes_to_dict(response)
+        return _bytes_to_dict(response)  # type: ignore
 
     def feature_extraction(self, text: str, *, model: Optional[str] = None) -> "np.ndarray":
         """
@@ -453,7 +454,7 @@ class InferenceClient:
         ```
         """
         response = self.post(data=image, model=model, task="image-classification")
-        return _bytes_to_dict(response)
+        return _bytes_to_list(response)
 
     def image_segmentation(
         self,
@@ -719,7 +720,7 @@ class InferenceClient:
             model=model,
             task="sentence-similarity",
         )
-        return _bytes_to_dict(response)
+        return _bytes_to_list(response)
 
     def summarization(
         self,
@@ -1285,7 +1286,7 @@ class InferenceClient:
             model=model,
             task="zero-shot-image-classification",
         )
-        return _bytes_to_dict(response)
+        return _bytes_to_list(response)
 
     def _resolve_url(self, model: Optional[str] = None, task: Optional[str] = None) -> str:
         model = model or self.model
diff --git a/src/huggingface_hub/inference/_common.py b/src/huggingface_hub/inference/_common.py
index 354cd6b8..c3124491 100644
--- a/src/huggingface_hub/inference/_common.py
+++ b/src/huggingface_hub/inference/_common.py
@@ -226,10 +226,24 @@ def _b64_to_image(encoded_image: str) -> "Image":
     return Image.open(io.BytesIO(base64.b64decode(encoded_image)))
 
 
-def _bytes_to_dict(content: bytes) -> "Image":
+def _bytes_to_list(content: bytes) -> List:
+    """Parse bytes from a Response object into a Python list.
+
+    Expects the response body to be JSON-encoded data.
+
+    NOTE: This is exactly the same implementation as `_bytes_to_dict` and will not complain if the returned data is a
+    dictionary. The only advantage of having both is to help the user (and mypy) understand what kind of data to expect.
+    """
+    return json.loads(content.decode())
+
+
+def _bytes_to_dict(content: bytes) -> Dict:
     """Parse bytes from a Response object into a Python dictionary.
 
-    Expects the response body to be encoded-JSON data.
+    Expects the response body to be JSON-encoded data.
+
+    NOTE: This is exactly the same implementation as `_bytes_to_list` and will not complain if the returned data is a
+    list. The only advantage of having both is to help the user (and mypy) understand what kind of data to expect.
     """
     return json.loads(content.decode())
 
diff --git a/src/huggingface_hub/inference/_generated/_async_client.py b/src/huggingface_hub/inference/_generated/_async_client.py
index f057958d..90629a99 100644
--- a/src/huggingface_hub/inference/_generated/_async_client.py
+++ b/src/huggingface_hub/inference/_generated/_async_client.py
@@ -47,6 +47,7 @@ from huggingface_hub.inference._common import (
     _b64_to_image,
     _bytes_to_dict,
     _bytes_to_image,
+    _bytes_to_list,
     _get_recommended_model,
     _import_numpy,
     _is_tgi_server,
@@ -281,7 +282,7 @@ class AsyncInferenceClient:
         ```
         """
         response = await self.post(data=audio, model=model, task="audio-classification")
-        return _bytes_to_dict(response)
+        return _bytes_to_list(response)
 
     async def automatic_speech_recognition(
         self,
@@ -380,7 +381,7 @@ class AsyncInferenceClient:
         if parameters is not None:
             payload["parameters"] = parameters
         response = await self.post(json=payload, model=model, task="conversational")
-        return _bytes_to_dict(response)
+        return _bytes_to_dict(response)  # type: ignore
 
     async def feature_extraction(self, text: str, *, model: Optional[str] = None) -> "np.ndarray":
         """
@@ -454,7 +455,7 @@ class AsyncInferenceClient:
         ```
         """
         response = await self.post(data=image, model=model, task="image-classification")
-        return _bytes_to_dict(response)
+        return _bytes_to_list(response)
 
     async def image_segmentation(
         self,
@@ -725,7 +726,7 @@ class AsyncInferenceClient:
             model=model,
             task="sentence-similarity",
         )
-        return _bytes_to_dict(response)
+        return _bytes_to_list(response)
 
     async def summarization(
         self,
@@ -1297,7 +1298,7 @@ class AsyncInferenceClient:
             model=model,
             task="zero-shot-image-classification",
         )
-        return _bytes_to_dict(response)
+        return _bytes_to_list(response)
 
     def _resolve_url(self, model: Optional[str] = None, task: Optional[str] = None) -> str:
         model = model or self.model
diff --git a/src/huggingface_hub/utils/_hf_folder.py b/src/huggingface_hub/utils/_hf_folder.py
index 5c9f07c9..77daced7 100644
--- a/src/huggingface_hub/utils/_hf_folder.py
+++ b/src/huggingface_hub/utils/_hf_folder.py
@@ -65,11 +65,14 @@ class HfFolder:
         # 1. Is it set by environment variable ?
         token: Optional[str] = os.environ.get("HUGGING_FACE_HUB_TOKEN")
         if token is not None:
+            token = token.replace("\r", "").replace("\n", "").strip()
             return token
 
         # 2. Is it set in token path ?
         try:
-            return cls.path_token.read_text()
+            token = cls.path_token.read_text()
+            token = token.replace("\r", "").replace("\n", "").strip()
+            return token
         except FileNotFoundError:
             return None
 


Test patch:
diff --git a/tests/test_utils_hf_folder.py b/tests/test_utils_hf_folder.py
index 7a1bcf19..ce0448ec 100644
--- a/tests/test_utils_hf_folder.py
+++ b/tests/test_utils_hf_folder.py
@@ -83,3 +83,12 @@ class HfFolderTest(unittest.TestCase):
             # Un-patch
             new_patcher.stop()
             old_patcher.stop()
+
+    def test_token_strip(self):
+        """
+        Test the workflow when the token is mistakenly finishing with new-line or space character.
+        """
+        token = _generate_token()
+        HfFolder.save_token(" " + token + "\n")
+        self.assertEqual(HfFolder.get_token(), token)
+        HfFolder.delete_token()