Repository: huggingface/huggingface_hub
Pull request #2006 (2024-01-30 22:56:28)

Fix intolerance for new field in TGI stream response: 'index'
Fixes https://github.com/huggingface/huggingface_hub/issues/2005

This PR adds tolerance for an `index` field in the TextGenerationStreamResponse.

The latest version of HuggingFace Text Generation Inference added an `index` field to the streaming response to be compatible with OpenAI clients, but the HuggingFace InferenceClient wasn't updated to accept it.

https://github.com/huggingface/text-generation-inference/commit/0eabc83541225979209ff7183b4b4442e47adf92#diff-c3259a23eb2e9e4e39f9a5226354b99508e42b05a158036dc482397410becd52R501

HuggingFaceDocBuilderDev: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/huggingface_hub/pr_2006). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Patch:
diff --git a/docs/source/en/guides/repository.md b/docs/source/en/guides/repository.md
index e8da4e2b..a9c9d3e9 100644
--- a/docs/source/en/guides/repository.md
+++ b/docs/source/en/guides/repository.md
@@ -108,7 +108,7 @@ You can create new branch and tags using [`create_branch`] and [`create_tag`]:
 >>> create_branch("Matthijs/speecht5-tts-demo", repo_type="space", branch="handle-dog-speaker")
 
 # Create a tag on a Dataset repo from `v0.1-release` branch
->>> create_branch("bigcode/the-stack", repo_type="dataset", revision="v0.1-release", tag="v0.1.1", tag_message="Bump release version.")
+>>> create_tag("bigcode/the-stack", repo_type="dataset", revision="v0.1-release", tag="v0.1.1", tag_message="Bump release version.")
 ```
 
 You can use the [`delete_branch`] and [`delete_tag`] functions in the same way to delete a branch or a tag.
diff --git a/src/huggingface_hub/__init__.py b/src/huggingface_hub/__init__.py
index 3041e611..f4aac31a 100644
--- a/src/huggingface_hub/__init__.py
+++ b/src/huggingface_hub/__init__.py
@@ -245,6 +245,7 @@ _SUBMOD_ATTRS = {
         "HfFileSystem",
         "HfFileSystemFile",
         "HfFileSystemResolvedPath",
+        "HfFileSystemStreamFile",
     ],
     "hub_mixin": [
         "ModelHubMixin",
@@ -592,6 +593,7 @@ if TYPE_CHECKING:  # pragma: no cover
         HfFileSystem,  # noqa: F401
         HfFileSystemFile,  # noqa: F401
         HfFileSystemResolvedPath,  # noqa: F401
+        HfFileSystemStreamFile,  # noqa: F401
     )
     from .hub_mixin import (
         ModelHubMixin,  # noqa: F401
diff --git a/src/huggingface_hub/hf_file_system.py b/src/huggingface_hub/hf_file_system.py
index 2407daf9..a78ab0fd 100644
--- a/src/huggingface_hub/hf_file_system.py
+++ b/src/huggingface_hub/hf_file_system.py
@@ -10,6 +10,7 @@ from typing import Any, Dict, List, NoReturn, Optional, Tuple, Union
 from urllib.parse import quote, unquote
 
 import fsspec
+from requests import Response
 
 from ._commit_api import CommitOperationCopy, CommitOperationDelete
 from .constants import DEFAULT_REVISION, ENDPOINT, REPO_TYPE_MODEL, REPO_TYPES_MAPPING, REPO_TYPES_URL_PREFIXES
@@ -216,11 +217,15 @@ class HfFileSystem(fsspec.AbstractFileSystem):
         path: str,
         mode: str = "rb",
         revision: Optional[str] = None,
+        block_size: Optional[int] = None,
         **kwargs,
     ) -> "HfFileSystemFile":
         if "a" in mode:
             raise NotImplementedError("Appending to remote files is not yet supported.")
-        return HfFileSystemFile(self, path, mode=mode, revision=revision, **kwargs)
+        if block_size == 0:
+            return HfFileSystemStreamFile(self, path, mode=mode, revision=revision, block_size=block_size, **kwargs)
+        else:
+            return HfFileSystemFile(self, path, mode=mode, revision=revision, block_size=block_size, **kwargs)
 
     def _rm(self, path: str, revision: Optional[str] = None, **kwargs) -> None:
         resolved_path = self.resolve_path(path, revision=revision)
@@ -649,6 +654,102 @@ class HfFileSystemFile(fsspec.spec.AbstractBufferedFile):
             )
 
 
+class HfFileSystemStreamFile(fsspec.spec.AbstractBufferedFile):
+    def __init__(
+        self,
+        fs: HfFileSystem,
+        path: str,
+        mode: str = "rb",
+        revision: Optional[str] = None,
+        block_size: int = 0,
+        cache_type: str = "none",
+        **kwargs,
+    ):
+        if block_size != 0:
+            raise ValueError(f"HfFileSystemStreamFile only supports block_size=0 but got {block_size}")
+        if cache_type != "none":
+            raise ValueError(f"HfFileSystemStreamFile only supports cache_type='none' but got {cache_type}")
+        if "w" in mode:
+            raise ValueError(f"HfFileSystemStreamFile only supports reading but got mode='{mode}'")
+        try:
+            self.resolved_path = fs.resolve_path(path, revision=revision)
+        except FileNotFoundError as e:
+            if "w" in kwargs.get("mode", ""):
+                raise FileNotFoundError(
+                    f"{e}.\nMake sure the repository and revision exist before writing data."
+                ) from e
+        # avoid an unecessary .info() call to instantiate .details
+        self.details = {"name": self.resolved_path.unresolve(), "size": None}
+        super().__init__(
+            fs, self.resolved_path.unresolve(), mode=mode, block_size=block_size, cache_type=cache_type, **kwargs
+        )
+        self.response: Optional[Response] = None
+        self.fs: HfFileSystem
+
+    def seek(self, loc: int, whence: int = 0):
+        if loc == 0 and whence == 1:
+            return
+        if loc == self.loc and whence == 0:
+            return
+        raise ValueError("Cannot seek streaming HF file")
+
+    def read(self, length: int = -1):
+        read_args = (length,) if length >= 0 else ()
+        if self.response is None or self.response.raw.isclosed():
+            url = hf_hub_url(
+                repo_id=self.resolved_path.repo_id,
+                revision=self.resolved_path.revision,
+                filename=self.resolved_path.path_in_repo,
+                repo_type=self.resolved_path.repo_type,
+                endpoint=self.fs.endpoint,
+            )
+            self.response = http_backoff(
+                "GET",
+                url,
+                headers=self.fs._api._build_hf_headers(),
+                retry_on_status_codes=(502, 503, 504),
+                stream=True,
+            )
+            hf_raise_for_status(self.response)
+        try:
+            out = self.response.raw.read(*read_args)
+        except Exception:
+            self.response.close()
+
+            # Retry by recreating the connection
+            url = hf_hub_url(
+                repo_id=self.resolved_path.repo_id,
+                revision=self.resolved_path.revision,
+                filename=self.resolved_path.path_in_repo,
+                repo_type=self.resolved_path.repo_type,
+                endpoint=self.fs.endpoint,
+            )
+            self.response = http_backoff(
+                "GET",
+                url,
+                headers={"Range": "bytes=%d-" % self.loc, **self.fs._api._build_hf_headers()},
+                retry_on_status_codes=(502, 503, 504),
+                stream=True,
+            )
+            hf_raise_for_status(self.response)
+            try:
+                out = self.response.raw.read(*read_args)
+            except Exception:
+                self.response.close()
+                raise
+        self.loc += len(out)
+        return out
+
+    def __del__(self):
+        if not hasattr(self, "resolved_path"):
+            # Means that the constructor failed. Nothing to do.
+            return
+        return super().__del__()
+
+    def __reduce__(self):
+        return reopen, (self.fs, self.path, self.mode, self.blocksize, self.cache.name)
+
+
 def safe_revision(revision: str) -> str:
     return revision if SPECIAL_REFS_REVISION_REGEX.match(revision) else safe_quote(revision)
 
@@ -666,3 +767,7 @@ def _raise_file_not_found(path: str, err: Optional[Exception]) -> NoReturn:
     elif isinstance(err, HFValidationError):
         msg = f"{path} (invalid repository id)"
     raise FileNotFoundError(msg) from err
+
+
+def reopen(fs: HfFileSystem, path: str, mode: str, block_size: int, cache_type: str):
+    return fs.open(path, mode=mode, block_size=block_size, cache_type=cache_type)
diff --git a/src/huggingface_hub/inference/_text_generation.py b/src/huggingface_hub/inference/_text_generation.py
index f48b0424..ab421f22 100644
--- a/src/huggingface_hub/inference/_text_generation.py
+++ b/src/huggingface_hub/inference/_text_generation.py
@@ -451,6 +451,8 @@ class TextGenerationStreamResponse:
     Args:
         token (`Token`):
             The generated token.
+        index (`Optional[int]`, *optional*):
+            The token index within the stream. Optional to support older clients that omit it.
         generated_text (`Optional[str]`, *optional*):
             The complete generated text. Only available when the generation is finished.
         details (`Optional[StreamDetails]`, *optional*):
@@ -459,6 +461,9 @@ class TextGenerationStreamResponse:
 
     # Generated token
     token: Token
+    # The token index within the stream
+    # Optional to support  older clients that omit it.
+    index: Optional[int] = None
     # Complete generated text
     # Only available when the generation is finished
     generated_text: Optional[str] = None
diff --git a/src/huggingface_hub/utils/endpoint_helpers.py b/src/huggingface_hub/utils/endpoint_helpers.py
index cbda8938..4061ae8f 100644
--- a/src/huggingface_hub/utils/endpoint_helpers.py
+++ b/src/huggingface_hub/utils/endpoint_helpers.py
@@ -75,7 +75,7 @@ class DatasetFilter:
 
     Args:
         author (`str`, *optional*):
-            A string or list of strings that can be used to identify datasets on
+            A string that can be used to identify datasets on
             the Hub by the original uploader (author or organization), such as
             `facebook` or `huggingface`.
         benchmark (`str` or `List`, *optional*):


Test patch:
diff --git a/tests/test_hf_file_system.py b/tests/test_hf_file_system.py
index 00e389a2..02cf9135 100644
--- a/tests/test_hf_file_system.py
+++ b/tests/test_hf_file_system.py
@@ -1,4 +1,5 @@
 import datetime
+import io
 import unittest
 from typing import Optional
 from unittest.mock import patch
@@ -6,7 +7,7 @@ from unittest.mock import patch
 import fsspec
 import pytest
 
-from huggingface_hub.hf_file_system import HfFileSystem
+from huggingface_hub.hf_file_system import HfFileSystem, HfFileSystemFile, HfFileSystemStreamFile
 from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError
 
 from .testing_constants import ENDPOINT_STAGING, TOKEN
@@ -152,8 +153,24 @@ class HfFileSystemTests(unittest.TestCase):
 
     def test_read_file(self):
         with self.hffs.open(self.hf_path + "/data/text_data.txt", "r") as f:
+            self.assertIsInstance(f, io.TextIOWrapper)
+            self.assertIsInstance(f.buffer, HfFileSystemFile)
             self.assertEqual(f.read(), "dummy text data")
 
+    def test_stream_file(self):
+        with self.hffs.open(self.hf_path + "/data/binary_data.bin", block_size=0) as f:
+            self.assertIsInstance(f, HfFileSystemStreamFile)
+            self.assertEqual(f.read(), b"dummy binary data")
+
+    def test_stream_file_retry(self):
+        with self.hffs.open(self.hf_path + "/data/binary_data.bin", block_size=0) as f:
+            self.assertIsInstance(f, HfFileSystemStreamFile)
+            self.assertEqual(f.read(6), b"dummy ")
+            # Simulate that streaming fails mid-way
+            f.response.raw.read = None
+            self.assertEqual(f.read(6), b"binary")
+            self.assertIsNotNone(f.response.raw.read)  # a new connection has been created
+
     def test_read_file_with_revision(self):
         with self.hffs.open(self.hf_path + "/data/binary_data_for_pr.bin", "rb", revision="refs/pr/1") as f:
             self.assertEqual(f.read(), b"dummy binary data on pr")
diff --git a/tests/test_inference_text_generation.py b/tests/test_inference_text_generation.py
index 6f76be05..88e1afcd 100644
--- a/tests/test_inference_text_generation.py
+++ b/tests/test_inference_text_generation.py
@@ -20,6 +20,7 @@ from huggingface_hub.inference._text_generation import (
     OverloadedError,
     TextGenerationParameters,
     TextGenerationRequest,
+    TextGenerationStreamResponse,
     raise_text_generation_error,
 )
 from huggingface_hub.inference._text_generation import (
@@ -106,6 +107,35 @@ class TestTextGenerationTypes(unittest.TestCase):
                 inputs="test", parameters=TextGenerationParameters(best_of=2, do_sample=True), stream=True
             )
 
+    def test_streaming_response_validation(self):
+        """
+        Regression test for #2005.
+
+        See https://github.com/huggingface/huggingface_hub/issues/2005
+        """
+        json_payload_latest = {
+            "index": 97,
+            "token": {"id": 264, "text": " a", "logprob": -0.0003259182, "special": False},
+            "generated_text": None,
+            "details": None,
+        }
+        TextGenerationStreamResponse(**json_payload_latest)
+
+        json_payload_older_had_no_index = {
+            "token": {"id": 264, "text": " a", "logprob": -0.0003259182, "special": False},
+            "generated_text": None,
+            "details": None,
+        }
+        TextGenerationStreamResponse(**json_payload_older_had_no_index)
+
+        json_payload_without_required_field_token = {
+            "index": 97,
+            "generated_text": None,
+            "details": None,
+        }
+        with self.assertRaises((ValidationError, TypeError)):
+            TextGenerationStreamResponse(**json_payload_without_required_field_token)
+
 
 class TestTextGenerationErrors(unittest.TestCase):
     def test_generation_error(self):