Repository: huggingface/huggingface_hub
Pull request #1711 (2023-10-05 08:25:19)

Handle TGI error when streaming tokens
Fix https://github.com/huggingface/huggingface_hub/issues/1651.

When streaming tokens from a TGI server, an error might happen and be returned in the stream of tokens. At the moment this raises a pydantic ValidationError as it's not recognized as a proper `TextGenerationStreamResponse`. This PR adds a check to parse correctly the error message + raise a custom exception whenever possible.

cc @osanseviero @OlivierDehaene 

**Note:** I couldn't really test this PR as the problem is hard to reproduce (it happened in the falcon-180b demo Space which has heavy usage -see [thread](https://huggingface.slack.com/archives/C05PZ9K0L6R/p1694091341361899?thread_ts=1694033671.072209&cid=C05PZ9K0L6R) (internal)). Shouldn't be too risky though (error/error_type being optional client-side).

HuggingFaceDocBuilderDev: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/huggingface_hub/pr_1711). All of your documentation changes will be reflected on that endpoint.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1711?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) Report
Attention: `3 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`90c1182`)](https://app.codecov.io/gh/huggingface/huggingface_hub/commit/90c11827d1743b04ab56496bcec74f4f592abf22?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) 82.30% compared to head [(`9d3b8a2`)](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1711?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) 82.40%.

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


```diff
@@            Coverage Diff             @@
##             main    #1711      +/-   ##
==========================================
+ Coverage   82.30%   82.40%   +0.10%     
==========================================
  Files          62       62              
  Lines        7226     7236      +10     
==========================================
+ Hits         5947     5963      +16     
+ Misses       1279     1273       -6     
```


| [Files](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1711?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/\_text\_generation.py](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1711?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface#diff-c3JjL2h1Z2dpbmdmYWNlX2h1Yi9pbmZlcmVuY2UvX3RleHRfZ2VuZXJhdGlvbi5weQ==) | `95.78% <93.33%> (-0.47%)` | :arrow_down: |
| [src/huggingface\_hub/inference/\_common.py](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1711?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface#diff-c3JjL2h1Z2dpbmdmYWNlX2h1Yi9pbmZlcmVuY2UvX2NvbW1vbi5weQ==) | `91.40% <60.00%> (-1.34%)` | :arrow_down: |

... and [3 files with indirect coverage changes](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1711/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1711?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface).   
:loudspeaker: Have feedback on the report? [Share it here](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/src/huggingface_hub/hf_file_system.py b/src/huggingface_hub/hf_file_system.py
index 16b58ac7..f9018bda 100644
--- a/src/huggingface_hub/hf_file_system.py
+++ b/src/huggingface_hub/hf_file_system.py
@@ -1,5 +1,6 @@
 import itertools
 import os
+import re
 import tempfile
 from dataclasses import dataclass
 from datetime import datetime
@@ -23,6 +24,17 @@ from .utils import (
 )
 
 
+# Regex used to match special revisions with "/" in them (see #1710)
+SPECIAL_REFS_REVISION_REGEX = re.compile(
+    r"""
+    (^refs\/convert\/parquet) # `refs/convert/parquet` revisions
+    |
+    (^refs\/pr\/\d+)          # PR revisions
+    """,
+    re.VERBOSE,
+)
+
+
 @dataclass
 class HfFileSystemResolvedPath:
     """Data structure containing information about a resolved Hugging Face file system path."""
@@ -34,7 +46,7 @@ class HfFileSystemResolvedPath:
 
     def unresolve(self) -> str:
         return (
-            f"{REPO_TYPES_URL_PREFIXES.get(self.repo_type, '') + self.repo_id}@{safe_quote(self.revision)}/{self.path_in_repo}"
+            f"{REPO_TYPES_URL_PREFIXES.get(self.repo_type, '') + self.repo_id}@{safe_revision(self.revision)}/{self.path_in_repo}"
             .rstrip("/")
         )
 
@@ -156,7 +168,13 @@ class HfFileSystem(fsspec.AbstractFileSystem):
             if "@" in path:
                 repo_id, revision_in_path = path.split("@", 1)
                 if "/" in revision_in_path:
-                    revision_in_path, path_in_repo = revision_in_path.split("/", 1)
+                    match = SPECIAL_REFS_REVISION_REGEX.search(revision_in_path)
+                    if match is not None and revision in (None, match.group()):
+                        # Handle `refs/convert/parquet` and PR revisions separately
+                        path_in_repo = SPECIAL_REFS_REVISION_REGEX.sub("", revision_in_path).lstrip("/")
+                        revision_in_path = match.group()
+                    else:
+                        revision_in_path, path_in_repo = revision_in_path.split("/", 1)
                 else:
                     path_in_repo = ""
                 revision_in_path = unquote(revision_in_path)
@@ -262,7 +280,7 @@ class HfFileSystem(fsspec.AbstractFileSystem):
     ) -> List[Union[str, Dict[str, Any]]]:
         """List the contents of a directory."""
         resolved_path = self.resolve_path(path, revision=revision)
-        revision_in_path = "@" + safe_quote(resolved_path.revision)
+        revision_in_path = "@" + safe_revision(resolved_path.revision)
         has_revision_in_path = revision_in_path in path
         path = resolved_path.unresolve()
         if path not in self.dircache or refresh:
@@ -367,7 +385,7 @@ class HfFileSystem(fsspec.AbstractFileSystem):
     def info(self, path: str, **kwargs) -> Dict[str, Any]:
         resolved_path = self.resolve_path(path)
         if not resolved_path.path_in_repo:
-            revision_in_path = "@" + safe_quote(resolved_path.revision)
+            revision_in_path = "@" + safe_revision(resolved_path.revision)
             has_revision_in_path = revision_in_path in path
             name = resolved_path.unresolve()
             name = name.replace(revision_in_path, "", 1) if not has_revision_in_path else name
@@ -418,5 +436,9 @@ class HfFileSystemFile(fsspec.spec.AbstractBufferedFile):
             )
 
 
+def safe_revision(revision: str) -> str:
+    return revision if SPECIAL_REFS_REVISION_REGEX.match(revision) else safe_quote(revision)
+
+
 def safe_quote(s: str) -> str:
     return quote(s, safe="")
diff --git a/src/huggingface_hub/inference/_common.py b/src/huggingface_hub/inference/_common.py
index c3124491..5973083a 100644
--- a/src/huggingface_hub/inference/_common.py
+++ b/src/huggingface_hub/inference/_common.py
@@ -48,9 +48,7 @@ from ..utils import (
     is_numpy_available,
     is_pillow_available,
 )
-from ._text_generation import (
-    TextGenerationStreamResponse,
-)
+from ._text_generation import TextGenerationStreamResponse, _parse_text_generation_error
 
 
 if TYPE_CHECKING:
@@ -275,7 +273,10 @@ def _stream_text_generation_response(
         if payload.startswith("data:"):
             # Decode payload
             json_payload = json.loads(payload.lstrip("data:").rstrip("/n"))
-            # Parse payload
+            # Either an error as being returned
+            if json_payload.get("error") is not None:
+                raise _parse_text_generation_error(json_payload["error"], json_payload.get("error_type"))
+            # Or parse token payload
             output = TextGenerationStreamResponse(**json_payload)
             yield output.token.text if not details else output
 
@@ -295,7 +296,10 @@ async def _async_stream_text_generation_response(
         if payload.startswith("data:"):
             # Decode payload
             json_payload = json.loads(payload.lstrip("data:").rstrip("/n"))
-            # Parse payload
+            # Either an error as being returned
+            if json_payload.get("error") is not None:
+                raise _parse_text_generation_error(json_payload["error"], json_payload.get("error_type"))
+            # Or parse token payload
             output = TextGenerationStreamResponse(**json_payload)
             yield output.token.text if not details else output
 
diff --git a/src/huggingface_hub/inference/_text_generation.py b/src/huggingface_hub/inference/_text_generation.py
index aba07834..a67c127b 100644
--- a/src/huggingface_hub/inference/_text_generation.py
+++ b/src/huggingface_hub/inference/_text_generation.py
@@ -447,6 +447,10 @@ class IncompleteGenerationError(TextGenerationError):
     pass
 
 
+class UnknownError(TextGenerationError):
+    pass
+
+
 def raise_text_generation_error(http_error: HTTPError) -> NoReturn:
     """
     Try to parse text-generation-inference error message and raise HTTPError in any case.
@@ -460,21 +464,27 @@ def raise_text_generation_error(http_error: HTTPError) -> NoReturn:
     try:
         # Hacky way to retrieve payload in case of aiohttp error
         payload = getattr(http_error, "response_error_payload", None) or http_error.response.json()
-        message = payload.get("error")
+        error = payload.get("error")
         error_type = payload.get("error_type")
     except Exception:  # no payload
         raise http_error
 
     # If error_type => more information than `hf_raise_for_status`
     if error_type is not None:
-        if error_type == "generation":
-            raise GenerationError(message) from http_error  # type: ignore
-        if error_type == "incomplete_generation":
-            raise IncompleteGenerationError(message) from http_error  # type: ignore
-        if error_type == "overloaded":
-            raise OverloadedError(message) from http_error  # type: ignore
-        if error_type == "validation":
-            raise ValidationError(message) from http_error  # type: ignore
+        exception = _parse_text_generation_error(error, error_type)
+        raise exception from http_error
 
     # Otherwise, fallback to default error
     raise http_error
+
+
+def _parse_text_generation_error(error: Optional[str], error_type: Optional[str]) -> TextGenerationError:
+    if error_type == "generation":
+        return GenerationError(error)  # type: ignore
+    if error_type == "incomplete_generation":
+        return IncompleteGenerationError(error)  # type: ignore
+    if error_type == "overloaded":
+        return OverloadedError(error)  # type: ignore
+    if error_type == "validation":
+        return ValidationError(error)  # type: ignore
+    return UnknownError(error)  # type: ignore


Test patch:
diff --git a/tests/test_hf_file_system.py b/tests/test_hf_file_system.py
index 9e2ca9ca..5c660cbe 100644
--- a/tests/test_hf_file_system.py
+++ b/tests/test_hf_file_system.py
@@ -213,7 +213,7 @@ class HfFileSystemTests(unittest.TestCase):
                 self.assertTrue(files[0]["name"].endswith("/data/binary_data_for_pr.bin"))  # PR file
 
 
-@pytest.mark.parametrize("path_in_repo", ["", "foo"])
+@pytest.mark.parametrize("path_in_repo", ["", "file.txt", "path/to/file"])
 @pytest.mark.parametrize(
     "root_path,revision,repo_type,repo_id,resolved_revision",
     [
@@ -238,6 +238,19 @@ class HfFileSystemTests(unittest.TestCase):
         ("hf://datasets/squad", None, "dataset", "squad", "main"),
         ("hf://datasets/squad", "dev", "dataset", "squad", "dev"),
         ("hf://datasets/squad@dev", None, "dataset", "squad", "dev"),
+        # Parse with `refs/convert/parquet` and `refs/pr/(\d)+` revisions.
+        # Regression tests for https://github.com/huggingface/huggingface_hub/issues/1710.
+        ("datasets/squad@refs/convert/parquet", None, "dataset", "squad", "refs/convert/parquet"),
+        (
+            "hf://datasets/username/my_dataset@refs/convert/parquet",
+            None,
+            "dataset",
+            "username/my_dataset",
+            "refs/convert/parquet",
+        ),
+        ("gpt2@refs/pr/2", None, "model", "gpt2", "refs/pr/2"),
+        ("hf://username/my_model@refs/pr/10", None, "model", "username/my_model", "refs/pr/10"),
+        ("hf://username/my_model@refs/pr/10", "refs/pr/10", "model", "username/my_model", "refs/pr/10"),
     ],
 )
 def test_resolve_path(
@@ -251,13 +264,7 @@ def test_resolve_path(
     fs = HfFileSystem()
     path = root_path + "/" + path_in_repo if path_in_repo else root_path
 
-    def mock_repo_info(repo_id: str, *, revision: str, repo_type: str, **kwargs):
-        if repo_id not in ["gpt2", "squad", "username/my_dataset", "username/my_model"]:
-            raise RepositoryNotFoundError(repo_id)
-        if revision is not None and revision not in ["main", "dev"]:
-            raise RevisionNotFoundError(revision)
-
-    with patch.object(fs._api, "repo_info", mock_repo_info):
+    with mock_repo_info(fs):
         resolved_path = fs.resolve_path(path, revision=revision)
         assert (
             resolved_path.repo_type,
@@ -267,6 +274,49 @@ def test_resolve_path(
         ) == (repo_type, repo_id, resolved_revision, path_in_repo)
 
 
+@pytest.mark.parametrize("path_in_repo", ["", "file.txt", "path/to/file"])
+@pytest.mark.parametrize(
+    "path,revision,expected_path",
+    [
+        ("hf://datasets/squad@dev", None, "datasets/squad@dev"),
+        ("datasets/squad@refs/convert/parquet", None, "datasets/squad@refs/convert/parquet"),
+        ("hf://username/my_model@refs/pr/10", None, "username/my_model@refs/pr/10"),
+        ("username/my_model", "refs/weirdo", "username/my_model@refs%2Fweirdo"),  # not a "special revision" -> encode
+    ],
+)
+def test_unresolve_path(path: str, revision: Optional[str], expected_path: str, path_in_repo: str) -> None:
+    fs = HfFileSystem()
+    path = path + "/" + path_in_repo if path_in_repo else path
+    expected_path = expected_path + "/" + path_in_repo if path_in_repo else expected_path
+
+    with mock_repo_info(fs):
+        assert fs.resolve_path(path, revision=revision).unresolve() == expected_path
+
+
+def test_resolve_path_with_refs_revision() -> None:
+    """
+    Testing a very specific edge case where a user has a repo with a revisions named "refs" and a file/directory
+    named "pr/10". We can still process them but the user has to use the `revision` argument to disambiguate between
+    the two.
+    """
+    fs = HfFileSystem()
+    with mock_repo_info(fs):
+        resolved = fs.resolve_path("hf://username/my_model@refs/pr/10", revision="refs")
+        assert resolved.revision == "refs"
+        assert resolved.path_in_repo == "pr/10"
+        assert resolved.unresolve() == "username/my_model@refs/pr/10"
+
+
+def mock_repo_info(fs: HfFileSystem):
+    def _inner(repo_id: str, *, revision: str, repo_type: str, **kwargs):
+        if repo_id not in ["gpt2", "squad", "username/my_dataset", "username/my_model"]:
+            raise RepositoryNotFoundError(repo_id)
+        if revision is not None and revision not in ["main", "dev", "refs"] and not revision.startswith("refs/"):
+            raise RevisionNotFoundError(revision)
+
+    return patch.object(fs._api, "repo_info", _inner)
+
+
 def test_resolve_path_with_non_matching_revisions():
     fs = HfFileSystem()
     with pytest.raises(ValueError):