huggingface-hub 0.21.2__py3-none-any.whl → 0.22.0__py3-none-any.whl

huggingface_hub/hf_file_system.py

@@ -6,15 +6,24 @@ from collections import deque
 from dataclasses import dataclass, field
 from datetime import datetime
 from itertools import chain
+from pathlib import Path
 from typing import Any, Dict, List, NoReturn, Optional, Tuple, Union
 from urllib.parse import quote, unquote
 
 import fsspec
+from fsspec.callbacks import _DEFAULT_CALLBACK, NoOpCallback, TqdmCallback
+from fsspec.utils import isfilelike
 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
-from .file_download import hf_hub_url
+from .constants import (
+    DEFAULT_REVISION,
+    ENDPOINT,
+    REPO_TYPE_MODEL,
+    REPO_TYPES_MAPPING,
+    REPO_TYPES_URL_PREFIXES,
+)
+from .file_download import hf_hub_url, http_get
 from .hf_api import HfApi, LastCommitInfo, RepoFile
 from .utils import (
     EntryNotFoundError,
@@ -591,6 +600,58 @@ class HfFileSystem(fsspec.AbstractFileSystem):
             url = url.replace("/resolve/", "/tree/", 1)
         return url
 
+    def get_file(self, rpath, lpath, callback=_DEFAULT_CALLBACK, outfile=None, **kwargs) -> None:
+        """Copy single remote file to local."""
+        revision = kwargs.get("revision")
+        unhandled_kwargs = set(kwargs.keys()) - {"revision"}
+        if not isinstance(callback, (NoOpCallback, TqdmCallback)) or len(unhandled_kwargs) > 0:
+            # for now, let's not handle custom callbacks
+            # and let's not handle custom kwargs
+            return super().get_file(rpath, lpath, callback=callback, outfile=outfile, **kwargs)
+
+        # Taken from https://github.com/fsspec/filesystem_spec/blob/47b445ae4c284a82dd15e0287b1ffc410e8fc470/fsspec/spec.py#L883
+        if isfilelike(lpath):
+            outfile = lpath
+        elif self.isdir(rpath):
+            os.makedirs(lpath, exist_ok=True)
+            return None
+
+        if isinstance(lpath, (str, Path)):  # otherwise, let's assume it's a file-like object
+            os.makedirs(os.path.dirname(lpath), exist_ok=True)
+
+        # Open file if not already open
+        close_file = False
+        if outfile is None:
+            outfile = open(lpath, "wb")
+            close_file = True
+        initial_pos = outfile.tell()
+
+        # Custom implementation of `get_file` to use `http_get`.
+        resolve_remote_path = self.resolve_path(rpath, revision=revision)
+        expected_size = self.info(rpath, revision=revision)["size"]
+        callback.set_size(expected_size)
+        try:
+            http_get(
+                url=hf_hub_url(
+                    repo_id=resolve_remote_path.repo_id,
+                    revision=resolve_remote_path.revision,
+                    filename=resolve_remote_path.path_in_repo,
+                    repo_type=resolve_remote_path.repo_type,
+                    endpoint=self.endpoint,
+                ),
+                temp_file=outfile,
+                displayed_filename=rpath,
+                expected_size=expected_size,
+                resume_size=0,
+                headers=self._api._build_hf_headers(),
+                _tqdm_bar=callback.tqdm if isinstance(callback, TqdmCallback) else None,
+            )
+            outfile.seek(initial_pos)
+        finally:
+            # Close file only if we opened it ourselves
+            if close_file:
+                outfile.close()
+
     @property
     def transaction(self):
         """A context within which files are committed together upon exit
@@ -618,6 +679,7 @@ class HfFileSystemFile(fsspec.spec.AbstractBufferedFile):
                 raise FileNotFoundError(
                     f"{e}.\nMake sure the repository and revision exist before writing data."
                 ) from e
+            raise
         super().__init__(fs, self.resolved_path.unresolve(), **kwargs)
         self.fs: HfFileSystem
 
@@ -667,6 +729,18 @@ class HfFileSystemFile(fsspec.spec.AbstractBufferedFile):
                 path=self.resolved_path.unresolve(),
             )
 
+    def read(self, length=-1):
+        """Read remote file.
+
+        If `length` is not provided or is -1, the entire file is downloaded and read. On POSIX systems and if
+        `hf_transfer` is not enabled, the file is loaded in memory directly. Otherwise, the file is downloaded to a
+        temporary file and read from there.
+        """
+        if self.mode == "rb" and (length is None or length == -1) and self.loc == 0:
+            with self.fs.open(self.path, "rb", block_size=0) as f:  # block_size=0 enables fast streaming
+                return f.read()
+        return super().read(length)
+
     def url(self) -> str:
         return self.fs.url(self.path)
 
@@ -695,7 +769,7 @@ class HfFileSystemStreamFile(fsspec.spec.AbstractBufferedFile):
                 raise FileNotFoundError(
                     f"{e}.\nMake sure the repository and revision exist before writing data."
                 ) from e
-        # avoid an unecessary .info() call to instantiate .details
+        # avoid an unnecessary .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

huggingface_hub/hub_mixin.py

@@ -1,17 +1,19 @@
 import inspect
 import json
 import os
-from dataclasses import asdict, is_dataclass
+from dataclasses import asdict, dataclass, is_dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, List, Optional, Type, TypeVar, Union, get_args
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, TypeVar, Union, get_args
 
 from .constants import CONFIG_NAME, PYTORCH_WEIGHTS_NAME, SAFETENSORS_SINGLE_FILE
 from .file_download import hf_hub_download
 from .hf_api import HfApi
+from .repocard import ModelCard, ModelCardData
 from .utils import (
     EntryNotFoundError,
     HfHubHTTPError,
     SoftTemporaryDirectory,
+    is_jsonable,
     is_safetensors_available,
     is_torch_available,
     logging,
@@ -27,8 +29,8 @@ if is_torch_available():
     import torch  # type: ignore
 
 if is_safetensors_available():
-    from safetensors import safe_open
-    from safetensors.torch import save_file
+    from safetensors.torch import load_model as load_model_as_safetensor
+    from safetensors.torch import save_model as save_model_as_safetensor
 
 
 logger = logging.get_logger(__name__)
@@ -36,6 +38,26 @@ logger = logging.get_logger(__name__)
 # Generic variable that is either ModelHubMixin or a subclass thereof
 T = TypeVar("T", bound="ModelHubMixin")
 
+DEFAULT_MODEL_CARD = """
+---
+# For reference on model card metadata, see the spec: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1
+# Doc / guide: https://huggingface.co/docs/hub/model-cards
+{{ card_data }}
+---
+
+This model has been pushed to the Hub using **{{ library_name }}**:
+- Repo: {{ repo_url | default("[More Information Needed]", true) }}
+- Docs: {{ docs_url | default("[More Information Needed]", true) }}
+"""
+
+
+@dataclass
+class MixinInfo:
+    library_name: Optional[str] = None
+    tags: Optional[List[str]] = None
+    repo_url: Optional[str] = None
+    docs_url: Optional[str] = None
+
 
 class ModelHubMixin:
     """
@@ -45,21 +67,35 @@ class ModelHubMixin:
     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.
 
+    When inheriting from [`ModelHubMixin`], you can define class-level attributes. These attributes are not passed to
+    `__init__` but to the class definition itself. This is useful to define metadata about the library integrating
+    [`ModelHubMixin`].
+
+    Args:
+        library_name (`str`, *optional*):
+            Name of the library integrating ModelHubMixin. Used to generate model card.
+        tags (`List[str]`, *optional*):
+            Tags to be added to the model card. Used to generate model card.
+        repo_url (`str`, *optional*):
+            URL of the library repository. Used to generate model card.
+        docs_url (`str`, *optional*):
+            URL of the library documentation. Used to generate model card.
+
     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):
+    # Inherit from ModelHubMixin
+    >>> class MyCustomModel(
+    ...         ModelHubMixin,
+    ...         library_name="my-library",
+    ...         tags=["x-custom-tag"],
+    ...         repo_url="https://github.com/huggingface/my-cool-library",
+    ...         docs_url="https://huggingface.co/docs/my-cool-library",
+    ...         # ^ optional metadata to generate model card
+    ...     ):
+    ...     def __init__(self, size: int = 512, device: str = "cpu"):
     ...         # define how to initialize your model
     ...         super().__init__()
     ...         ...
@@ -85,7 +121,7 @@ class ModelHubMixin:
     ...         # define how to deserialize your model
     ...         ...
 
-    >>> model = MyCustomModel(config=Config(foo=256, bar="gpu"))
+    >>> model = MyCustomModel(size=256, device="gpu")
 
     # Save model weights to local directory
     >>> model.save_pretrained("my-awesome-model")
@@ -95,28 +131,107 @@ class ModelHubMixin:
 
     # Download and initialize weights from the Hub
     >>> reloaded_model = MyCustomModel.from_pretrained("username/my-awesome-model")
-    >>> reloaded_model.config
-    Config(foo=256, bar="gpu")
+    >>> reloaded_model._hub_mixin_config
+    {"size": 256, "device": "gpu"}
+
+    # Model card has been correctly populated
+    >>> from huggingface_hub import ModelCard
+    >>> card = ModelCard.load("username/my-awesome-model")
+    >>> card.data.tags
+    ["x-custom-tag", "pytorch_model_hub_mixin", "model_hub_mixin"]
+    >>> card.data.library_name
+    "my-library"
     ```
     """
 
-    config: Optional[Union[dict, "DataclassInstance"]] = None
-    # ^ optional config attribute automatically set in `from_pretrained` (if not already set by the subclass)
+    _hub_mixin_config: Optional[Union[dict, "DataclassInstance"]] = None
+    # ^ optional config attribute automatically set in `from_pretrained`
+    _hub_mixin_info: MixinInfo
+    # ^ information about the library integrating ModelHubMixin (used to generate model card)
+    _hub_mixin_init_parameters: Dict[str, inspect.Parameter]
+    _hub_mixin_jsonable_default_values: Dict[str, Any]
+    _hub_mixin_inject_config: bool
+    # ^ internal values to handle config
+
+    def __init_subclass__(
+        cls,
+        *,
+        library_name: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+        repo_url: Optional[str] = None,
+        docs_url: Optional[str] = None,
+    ) -> None:
+        """Inspect __init__ signature only once when subclassing + handle modelcard."""
+        super().__init_subclass__()
+
+        # Will be reused when creating modelcard
+        tags = tags or []
+        tags.append("model_hub_mixin")
+        cls._hub_mixin_info = MixinInfo(
+            library_name=library_name,
+            tags=tags,
+            repo_url=repo_url,
+            docs_url=docs_url,
+        )
+
+        # Inspect __init__ signature to handle config
+        cls._hub_mixin_init_parameters = dict(inspect.signature(cls.__init__).parameters)
+        cls._hub_mixin_jsonable_default_values = {
+            param.name: param.default
+            for param in cls._hub_mixin_init_parameters.values()
+            if param.default is not inspect.Parameter.empty and is_jsonable(param.default)
+        }
+        cls._hub_mixin_inject_config = "config" in inspect.signature(cls._from_pretrained).parameters
 
     def __new__(cls, *args, **kwargs) -> "ModelHubMixin":
+        """Create a new instance of the class and handle config.
+
+        3 cases:
+        - If `self._hub_mixin_config` is already set, do nothing.
+        - If `config` is passed as a dataclass, set it as `self._hub_mixin_config`.
+        - Otherwise, build `self._hub_mixin_config` from default values and passed values.
+        """
         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
+        # If `config` is already set, return early
+        if instance._hub_mixin_config is not None:
+            return instance
+
+        # Infer passed values
+        passed_values = {
+            **{
+                key: value
+                for key, value in zip(
+                    # [1:] to skip `self` parameter
+                    list(cls._hub_mixin_init_parameters)[1:],
+                    args,
+                )
+            },
+            **kwargs,
+        }
+
+        # If config passed as dataclass => set it and return early
+        if is_dataclass(passed_values.get("config")):
+            instance._hub_mixin_config = passed_values["config"]
+            return instance
+
+        # Otherwise, build config from default + passed values
+        init_config = {
+            # default values
+            **cls._hub_mixin_jsonable_default_values,
+            # passed values
+            **{key: value for key, value in passed_values.items() if is_jsonable(value)},
+        }
+        init_config.pop("config", {})
+
+        # Populate `init_config` with provided config
+        provided_config = passed_values.get("config")
+        if isinstance(provided_config, dict):
+            init_config.update(provided_config)
+
+        # Set `config` attribute and return
+        if init_config != {}:
+            instance._hub_mixin_config = init_config
         return instance
 
     def save_pretrained(
@@ -147,16 +262,29 @@ class ModelHubMixin:
         save_directory = Path(save_directory)
         save_directory.mkdir(parents=True, exist_ok=True)
 
+        # Remove config.json if already exists. After `_save_pretrained` we don't want to overwrite config.json
+        # as it might have been saved by the custom `_save_pretrained` already. However we do want to overwrite
+        # an existing config.json if it was not saved by `_save_pretrained`.
+        config_path = save_directory / CONFIG_NAME
+        config_path.unlink(missing_ok=True)
+
         # save model weights/files (framework-specific)
         self._save_pretrained(save_directory)
 
-        # save config (if provided)
+        # save config (if provided and if not serialized yet in `_save_pretrained`)
         if config is None:
-            config = self.config
+            config = self._hub_mixin_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))
+            if not config_path.exists():
+                config_str = json.dumps(config, sort_keys=True, indent=2)
+                config_path.write_text(config_str)
+
+        # save model card
+        model_card_path = save_directory / "README.md"
+        if not model_card_path.exists():  # do not overwrite if already exists
+            self.generate_model_card().save(save_directory / "README.md")
 
         # push to the Hub if required
         if push_to_hub:
@@ -246,31 +374,44 @@ class ModelHubMixin:
             except HfHubHTTPError as e:
                 logger.info(f"{CONFIG_NAME} not found on the HuggingFace Hub: {str(e)}")
 
+        # Read config
         config = None
         if config_file is not None:
-            # Read config
             with open(config_file, "r", encoding="utf-8") as f:
                 config = json.load(f)
 
-            # Check if class expect a `config` argument
-            init_parameters = inspect.signature(cls.__init__).parameters
-            if "config" in init_parameters:
+            # Populate model_kwargs from config
+            for param in cls._hub_mixin_init_parameters.values():
+                if param.name not in model_kwargs and param.name in config:
+                    model_kwargs[param.name] = config[param.name]
+
+            # Check if `config` argument was passed at init
+            if "config" in cls._hub_mixin_init_parameters:
                 # Check if `config` argument is a dataclass
-                config_annotation = init_parameters["config"].annotation
+                config_annotation = cls._hub_mixin_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
+                    config = _load_dataclass(config_annotation, config)
                 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)
+                            config = _load_dataclass(_sub_annotation, config)
                             break
 
                 # Forward config to model initialization
                 model_kwargs["config"] = config
 
+            if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in cls._hub_mixin_init_parameters.values()):
+                for key, value in config.items():
+                    if key not in model_kwargs:
+                        model_kwargs[key] = value
+
+            # Finally, also inject if `_from_pretrained` expects it
+            if cls._hub_mixin_inject_config:
+                model_kwargs["config"] = config
+
         instance = cls._from_pretrained(
             model_id=str(model_id),
             revision=revision,
@@ -285,8 +426,8 @@ class ModelHubMixin:
 
         # 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
+        if config is not None and (getattr(instance, "_hub_mixin_config", None) in (None, {})):
+            instance._hub_mixin_config = config
 
         return instance
 
@@ -415,6 +556,13 @@ class ModelHubMixin:
                 delete_patterns=delete_patterns,
             )
 
+    def generate_model_card(self, *args, **kwargs) -> ModelCard:
+        card = ModelCard.from_template(
+            card_data=ModelCardData(**asdict(self._hub_mixin_info)),
+            template_str=DEFAULT_MODEL_CARD,
+        )
+        return card
+
 
 class PyTorchModelHubMixin(ModelHubMixin):
     """
@@ -425,26 +573,26 @@ 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, config: Config):
+    >>> class MyModel(
+    ...         nn.Module,
+    ...         PyTorchModelHubMixin,
+    ...         library_name="keras-nlp",
+    ...         repo_url="https://github.com/keras-team/keras-nlp",
+    ...         docs_url="https://keras.io/keras_nlp/",
+    ...         # ^ optional metadata to generate model card
+    ...     ):
+    ...     def __init__(self, hidden_size: int = 512, vocab_size: int = 30000, output_size: int = 4):
     ...         super().__init__()
-    ...         self.param = nn.Parameter(torch.rand(config.hidden_size, config.vocab_size))
-    ...         self.linear = nn.Linear(config.output_size, config.vocab_size)
+    ...         self.param = nn.Parameter(torch.rand(hidden_size, vocab_size))
+    ...         self.linear = nn.Linear(output_size, vocab_size)
 
     ...     def forward(self, x):
     ...         return self.linear(x + self.param)
-    >>> model = MyModel()
+    >>> model = MyModel(hidden_size=256)
 
     # Save model weights to local directory
     >>> model.save_pretrained("my-awesome-model")
@@ -454,13 +602,21 @@ class PyTorchModelHubMixin(ModelHubMixin):
 
     # Download and initialize weights from the Hub
     >>> model = MyModel.from_pretrained("username/my-awesome-model")
+    >>> model.hidden_size
+    256
     ```
     """
 
+    def __init_subclass__(cls, *args, tags: Optional[List[str]] = None, **kwargs) -> None:
+        tags = tags or []
+        tags.append("pytorch_model_hub_mixin")
+        kwargs["tags"] = tags
+        return super().__init_subclass__(*args, **kwargs)
+
     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
-        save_file(model_to_save.state_dict(), save_directory / SAFETENSORS_SINGLE_FILE)
+        save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE))
 
     @classmethod
     def _from_pretrained(
@@ -521,10 +677,23 @@ class PyTorchModelHubMixin(ModelHubMixin):
 
     @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
+        load_model_as_safetensor(model, model_file, strict=strict)  # type: ignore [arg-type]
+        if map_location != "cpu":
+            # TODO: remove this once https://github.com/huggingface/safetensors/pull/449 is merged.
+            logger.warning(
+                "Loading model weights on other devices than 'cpu' is not supported natively."
+                " This means that the model is loaded on 'cpu' first and then copied to the device."
+                " This leads to a slower loading time."
+                " Support for loading directly on other devices is planned to be added in future releases."
+                " See https://github.com/huggingface/huggingface_hub/pull/2086 for more details."
+            )
+            model.to(map_location)  # type: ignore [attr-defined]
         return model
+
+
+def _load_dataclass(datacls: Type["DataclassInstance"], data: dict) -> "DataclassInstance":
+    """Load a dataclass instance from a dictionary.
+
+    Fields not expected by the dataclass are ignored.
+    """
+    return datacls(**{k: v for k, v in data.items() if k in datacls.__dataclass_fields__})