huggingface-hub 0.20.2__py3-none-any.whl → 0.21.0__py3-none-any.whl

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)
@@ -244,9 +249,8 @@ class HfFileSystem(fsspec.AbstractFileSystem):
         **kwargs,
     ) -> None:
         resolved_path = self.resolve_path(path, revision=revision)
-        root_path = REPO_TYPES_URL_PREFIXES.get(resolved_path.repo_type, "") + resolved_path.repo_id
         paths = self.expand_path(path, recursive=recursive, maxdepth=maxdepth, revision=revision)
-        paths_in_repo = [path[len(root_path) + 1 :] for path in paths if not self.isdir(path)]
+        paths_in_repo = [self.resolve_path(path).path_in_repo for path in paths if not self.isdir(path)]
         operations = [CommitOperationDelete(path_in_repo=path_in_repo) for path_in_repo in paths_in_repo]
         commit_message = f"Delete {path} "
         commit_message += "recursively " if recursive else ""
@@ -439,7 +443,7 @@ class HfFileSystem(fsspec.AbstractFileSystem):
             resolved_path1.repo_type == resolved_path2.repo_type and resolved_path1.repo_id == resolved_path2.repo_id
         )
 
-        if same_repo and self.info(path1, revision=resolved_path1.revision)["lfs"] is not None:
+        if same_repo:
             commit_message = f"Copy {path1} to {path2}"
             self._api.create_commit(
                 repo_id=resolved_path1.repo_id,
@@ -573,6 +577,20 @@ class HfFileSystem(fsspec.AbstractFileSystem):
         except:  # noqa: E722
             return False
 
+    def url(self, path: str) -> str:
+        """Get the HTTP URL of the given path"""
+        resolved_path = self.resolve_path(path)
+        url = hf_hub_url(
+            resolved_path.repo_id,
+            resolved_path.path_in_repo,
+            repo_type=resolved_path.repo_type,
+            revision=resolved_path.revision,
+            endpoint=self.endpoint,
+        )
+        if self.isdir(path):
+            url = url.replace("/resolve/", "/tree/", 1)
+        return url
+
     @property
     def transaction(self):
         """A context within which files are committed together upon exit
@@ -593,9 +611,6 @@ class HfFileSystem(fsspec.AbstractFileSystem):
 
 class HfFileSystemFile(fsspec.spec.AbstractBufferedFile):
     def __init__(self, fs: HfFileSystem, path: str, revision: Optional[str] = None, **kwargs):
-        super().__init__(fs, path, **kwargs)
-        self.fs: HfFileSystem
-
         try:
             self.resolved_path = fs.resolve_path(path, revision=revision)
         except FileNotFoundError as e:
@@ -603,6 +618,8 @@ class HfFileSystemFile(fsspec.spec.AbstractBufferedFile):
                 raise FileNotFoundError(
                     f"{e}.\nMake sure the repository and revision exist before writing data."
                 ) from e
+        super().__init__(fs, self.resolved_path.unresolve(), **kwargs)
+        self.fs: HfFileSystem
 
     def __del__(self):
         if not hasattr(self, "resolved_path"):
@@ -622,7 +639,7 @@ class HfFileSystemFile(fsspec.spec.AbstractBufferedFile):
             repo_type=self.resolved_path.repo_type,
             endpoint=self.fs.endpoint,
         )
-        r = http_backoff("GET", url, headers=headers)
+        r = http_backoff("GET", url, headers=headers, retry_on_status_codes=(502, 503, 504))
         hf_raise_for_status(r)
         return r.content
 
@@ -650,6 +667,108 @@ class HfFileSystemFile(fsspec.spec.AbstractBufferedFile):
                 path=self.resolved_path.unresolve(),
             )
 
+    def url(self) -> str:
+        return self.fs.url(self.path)
+
+
+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 url(self) -> str:
+        return self.fs.url(self.path)
+
+    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)
@@ -668,3 +787,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)

huggingface_hub/hub_mixin.py

@@ -1,17 +1,36 @@
+import inspect
 import json
 import os
+from dataclasses import asdict, is_dataclass
 from pathlib import Path
-from typing import Dict, List, Optional, Type, TypeVar, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Type, TypeVar, Union, get_args
 
-from .constants import CONFIG_NAME, PYTORCH_WEIGHTS_NAME
-from .file_download import hf_hub_download, is_torch_available
+from .constants import CONFIG_NAME, PYTORCH_WEIGHTS_NAME, SAFETENSORS_SINGLE_FILE
+from .file_download import hf_hub_download
 from .hf_api import HfApi
-from .utils import HfHubHTTPError, SoftTemporaryDirectory, logging, validate_hf_hub_args
+from .utils import (
+    EntryNotFoundError,
+    HfHubHTTPError,
+    SoftTemporaryDirectory,
+    is_safetensors_available,
+    is_torch_available,
+    logging,
+    validate_hf_hub_args,
+)
+from .utils._deprecation import _deprecate_arguments
 
 
+if TYPE_CHECKING:
+    from _typeshed import DataclassInstance
+
 if is_torch_available():
     import torch  # type: ignore
 
+if is_safetensors_available():
+    from safetensors import safe_open
+    from safetensors.torch import save_file
+
+
 logger = logging.get_logger(__name__)
 
 # Generic variable that is either ModelHubMixin or a subclass thereof
@@ -25,16 +44,89 @@ class ModelHubMixin:
     To integrate your framework, your model class must inherit from this class. Custom logic for saving/loading models
     have to be overwritten in  [`_from_pretrained`] and [`_save_pretrained`]. [`PyTorchModelHubMixin`] is a good example
     of mixin integration with the Hub. Check out our [integration guide](../guides/integrations) for more instructions.
+
+    Example:
+
+    ```python
+    >>> from dataclasses import dataclass
+    >>> from huggingface_hub import ModelHubMixin
+
+    # Define your model configuration (optional)
+    >>> @dataclass
+    ... class Config:
+    ...     foo: int = 512
+    ...     bar: str = "cpu"
+
+    # Inherit from ModelHubMixin (and optionally from your framework's model class)
+    >>> class MyCustomModel(ModelHubMixin):
+    ...     def __init__(self, config: Config):
+    ...         # define how to initialize your model
+    ...         super().__init__()
+    ...         ...
+    ...
+    ...     def _save_pretrained(self, save_directory: Path) -> None:
+    ...         # define how to serialize your model
+    ...         ...
+    ...
+    ...     @classmethod
+    ...     def from_pretrained(
+    ...         cls: Type[T],
+    ...         pretrained_model_name_or_path: Union[str, Path],
+    ...         *,
+    ...         force_download: bool = False,
+    ...         resume_download: bool = False,
+    ...         proxies: Optional[Dict] = None,
+    ...         token: Optional[Union[str, bool]] = None,
+    ...         cache_dir: Optional[Union[str, Path]] = None,
+    ...         local_files_only: bool = False,
+    ...         revision: Optional[str] = None,
+    ...         **model_kwargs,
+    ...     ) -> T:
+    ...         # define how to deserialize your model
+    ...         ...
+
+    >>> model = MyCustomModel(config=Config(foo=256, bar="gpu"))
+
+    # Save model weights to local directory
+    >>> model.save_pretrained("my-awesome-model")
+
+    # Push model weights to the Hub
+    >>> model.push_to_hub("my-awesome-model")
+
+    # Download and initialize weights from the Hub
+    >>> reloaded_model = MyCustomModel.from_pretrained("username/my-awesome-model")
+    >>> reloaded_model.config
+    Config(foo=256, bar="gpu")
+    ```
     """
 
+    config: Optional[Union[dict, "DataclassInstance"]] = None
+    # ^ optional config attribute automatically set in `from_pretrained` (if not already set by the subclass)
+
+    def __new__(cls, *args, **kwargs) -> "ModelHubMixin":
+        instance = super().__new__(cls)
+
+        # Set `config` attribute if not already set by the subclass
+        if instance.config is None:
+            if "config" in kwargs:
+                instance.config = kwargs["config"]
+            elif len(args) > 0:
+                sig = inspect.signature(cls.__init__)
+                parameters = list(sig.parameters)[1:]  # remove `self`
+                for key, value in zip(parameters, args):
+                    if key == "config":
+                        instance.config = value
+                        break
+        return instance
+
     def save_pretrained(
         self,
         save_directory: Union[str, Path],
         *,
-        config: Optional[dict] = None,
+        config: Optional[Union[dict, "DataclassInstance"]] = None,
         repo_id: Optional[str] = None,
         push_to_hub: bool = False,
-        **kwargs,
+        **push_to_hub_kwargs,
     ) -> Optional[str]:
         """
         Save weights in local directory.
@@ -42,8 +134,8 @@ class ModelHubMixin:
         Args:
             save_directory (`str` or `Path`):
                 Path to directory in which the model weights and configuration will be saved.
-            config (`dict`, *optional*):
-                Model configuration specified as a key/value dictionary.
+            config (`dict` or `DataclassInstance`, *optional*):
+                Model configuration specified as a key/value dictionary or a dataclass instance.
             push_to_hub (`bool`, *optional*, defaults to `False`):
                 Whether or not to push your model to the Huggingface Hub after saving it.
             repo_id (`str`, *optional*):
@@ -55,15 +147,20 @@ class ModelHubMixin:
         save_directory = Path(save_directory)
         save_directory.mkdir(parents=True, exist_ok=True)
 
-        # saving model weights/files
+        # save model weights/files (framework-specific)
         self._save_pretrained(save_directory)
 
-        # saving config
-        if isinstance(config, dict):
-            (save_directory / CONFIG_NAME).write_text(json.dumps(config))
+        # save config (if provided)
+        if config is None:
+            config = self.config
+        if config is not None:
+            if is_dataclass(config):
+                config = asdict(config)  # type: ignore[arg-type]
+            (save_directory / CONFIG_NAME).write_text(json.dumps(config, indent=2))
 
+        # push to the Hub if required
         if push_to_hub:
-            kwargs = kwargs.copy()  # soft-copy to avoid mutating input
+            kwargs = push_to_hub_kwargs.copy()  # soft-copy to avoid mutating input
             if config is not None:  # kwarg for `push_to_hub`
                 kwargs["config"] = config
             if repo_id is None:
@@ -126,17 +223,17 @@ class ModelHubMixin:
             model_kwargs (`Dict`, *optional*):
                 Additional kwargs to pass to the model during initialization.
         """
-        model_id = pretrained_model_name_or_path
+        model_id = str(pretrained_model_name_or_path)
         config_file: Optional[str] = None
         if os.path.isdir(model_id):
             if CONFIG_NAME in os.listdir(model_id):
                 config_file = os.path.join(model_id, CONFIG_NAME)
             else:
                 logger.warning(f"{CONFIG_NAME} not found in {Path(model_id).resolve()}")
-        elif isinstance(model_id, str):
+        else:
             try:
                 config_file = hf_hub_download(
-                    repo_id=str(model_id),
+                    repo_id=model_id,
                     filename=CONFIG_NAME,
                     revision=revision,
                     cache_dir=cache_dir,
@@ -146,15 +243,35 @@ class ModelHubMixin:
                     token=token,
                     local_files_only=local_files_only,
                 )
-            except HfHubHTTPError:
-                logger.info(f"{CONFIG_NAME} not found in HuggingFace Hub.")
+            except HfHubHTTPError as e:
+                logger.info(f"{CONFIG_NAME} not found on the HuggingFace Hub: {str(e)}")
 
+        config = None
         if config_file is not None:
+            # Read config
             with open(config_file, "r", encoding="utf-8") as f:
                 config = json.load(f)
-            model_kwargs.update({"config": config})
 
-        return cls._from_pretrained(
+            # Check if class expect a `config` argument
+            init_parameters = inspect.signature(cls.__init__).parameters
+            if "config" in init_parameters:
+                # Check if `config` argument is a dataclass
+                config_annotation = init_parameters["config"].annotation
+                if config_annotation is inspect.Parameter.empty:
+                    pass  # no annotation
+                elif is_dataclass(config_annotation):
+                    config = config_annotation(**config)  # expect a dataclass
+                else:
+                    # if Optional/Union annotation => check if a dataclass is in the Union
+                    for _sub_annotation in get_args(config_annotation):
+                        if is_dataclass(_sub_annotation):
+                            config = _sub_annotation(**config)
+                            break
+
+                # Forward config to model initialization
+                model_kwargs["config"] = config
+
+        instance = cls._from_pretrained(
             model_id=str(model_id),
             revision=revision,
             cache_dir=cache_dir,
@@ -166,6 +283,13 @@ class ModelHubMixin:
             **model_kwargs,
         )
 
+        # Implicitly set the config as instance attribute if not already set by the class
+        # This way `config` will be available when calling `save_pretrained` or `push_to_hub`.
+        if config is not None and instance.config is None:
+            instance.config = config
+
+        return instance
+
     @classmethod
     def _from_pretrained(
         cls: Type[T],
@@ -215,21 +339,27 @@ class ModelHubMixin:
         """
         raise NotImplementedError
 
+    @_deprecate_arguments(
+        version="0.23.0",
+        deprecated_args=["api_endpoint"],
+        custom_message="Use `HF_ENDPOINT` environment variable instead.",
+    )
     @validate_hf_hub_args
     def push_to_hub(
         self,
         repo_id: str,
         *,
-        config: Optional[dict] = None,
+        config: Optional[Union[dict, "DataclassInstance"]] = None,
         commit_message: str = "Push model using huggingface_hub.",
         private: bool = False,
-        api_endpoint: Optional[str] = None,
         token: Optional[str] = None,
         branch: Optional[str] = None,
         create_pr: Optional[bool] = None,
         allow_patterns: Optional[Union[List[str], str]] = None,
         ignore_patterns: Optional[Union[List[str], str]] = None,
         delete_patterns: Optional[Union[List[str], str]] = None,
+        # TODO: remove once deprecated
+        api_endpoint: Optional[str] = None,
     ) -> str:
         """
         Upload model checkpoint to the Hub.
@@ -238,12 +368,11 @@ class ModelHubMixin:
         `delete_patterns` to delete existing remote files in the same commit. See [`upload_folder`] reference for more
         details.
 
-
         Args:
             repo_id (`str`):
                 ID of the repository to push to (example: `"username/my-model"`).
-            config (`dict`, *optional*):
-                Configuration object to be saved alongside the model weights.
+            config (`dict` or `DataclassInstance`, *optional*):
+                Model configuration specified as a key/value dictionary or a dataclass instance.
             commit_message (`str`, *optional*):
                 Message to commit while pushing.
             private (`bool`, *optional*, defaults to `False`):
@@ -296,16 +425,22 @@ class PyTorchModelHubMixin(ModelHubMixin):
     Example:
 
     ```python
+    >>> from dataclasses import dataclass
     >>> import torch
     >>> import torch.nn as nn
     >>> from huggingface_hub import PyTorchModelHubMixin
 
+    >>> @dataclass
+    ... class Config:
+    ...     hidden_size: int = 512
+    ...     vocab_size: int = 30000
+    ...     output_size: int = 4
 
     >>> class MyModel(nn.Module, PyTorchModelHubMixin):
-    ...     def __init__(self):
+    ...     def __init__(self, config: Config):
     ...         super().__init__()
-    ...         self.param = nn.Parameter(torch.rand(3, 4))
-    ...         self.linear = nn.Linear(4, 5)
+    ...         self.param = nn.Parameter(torch.rand(config.hidden_size, config.vocab_size))
+    ...         self.linear = nn.Linear(config.output_size, config.vocab_size)
 
     ...     def forward(self, x):
     ...         return self.linear(x + self.param)
@@ -325,7 +460,7 @@ class PyTorchModelHubMixin(ModelHubMixin):
     def _save_pretrained(self, save_directory: Path) -> None:
         """Save weights from a Pytorch model to a local directory."""
         model_to_save = self.module if hasattr(self, "module") else self  # type: ignore
-        torch.save(model_to_save.state_dict(), save_directory / PYTORCH_WEIGHTS_NAME)
+        save_file(model_to_save.state_dict(), save_directory / SAFETENSORS_SINGLE_FILE)
 
     @classmethod
     def _from_pretrained(
@@ -344,25 +479,52 @@ class PyTorchModelHubMixin(ModelHubMixin):
         **model_kwargs,
     ):
         """Load Pytorch pretrained weights and return the loaded model."""
+        model = cls(**model_kwargs)
         if os.path.isdir(model_id):
             print("Loading weights from local directory")
-            model_file = os.path.join(model_id, PYTORCH_WEIGHTS_NAME)
+            model_file = os.path.join(model_id, SAFETENSORS_SINGLE_FILE)
+            return cls._load_as_safetensor(model, model_file, map_location, strict)
         else:
-            model_file = hf_hub_download(
-                repo_id=model_id,
-                filename=PYTORCH_WEIGHTS_NAME,
-                revision=revision,
-                cache_dir=cache_dir,
-                force_download=force_download,
-                proxies=proxies,
-                resume_download=resume_download,
-                token=token,
-                local_files_only=local_files_only,
-            )
-        model = cls(**model_kwargs)
+            try:
+                model_file = hf_hub_download(
+                    repo_id=model_id,
+                    filename=SAFETENSORS_SINGLE_FILE,
+                    revision=revision,
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    proxies=proxies,
+                    resume_download=resume_download,
+                    token=token,
+                    local_files_only=local_files_only,
+                )
+                return cls._load_as_safetensor(model, model_file, map_location, strict)
+            except EntryNotFoundError:
+                model_file = hf_hub_download(
+                    repo_id=model_id,
+                    filename=PYTORCH_WEIGHTS_NAME,
+                    revision=revision,
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    proxies=proxies,
+                    resume_download=resume_download,
+                    token=token,
+                    local_files_only=local_files_only,
+                )
+                return cls._load_as_pickle(model, model_file, map_location, strict)
 
+    @classmethod
+    def _load_as_pickle(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
         state_dict = torch.load(model_file, map_location=torch.device(map_location))
         model.load_state_dict(state_dict, strict=strict)  # type: ignore
         model.eval()  # type: ignore
+        return model
 
+    @classmethod
+    def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
+        state_dict = {}
+        with safe_open(model_file, framework="pt", device=map_location) as f:  # type: ignore [attr-defined]
+            for k in f.keys():
+                state_dict[k] = f.get_tensor(k)
+        model.load_state_dict(state_dict, strict=strict)  # type: ignore
+        model.eval()  # type: ignore
         return model