huggingface-hub 0.22.2__py3-none-any.whl → 0.23.0rc0__py3-none-any.whl

huggingface_hub/hf_file_system.py

@@ -1,4 +1,3 @@
-import copy
 import os
 import re
 import tempfile
@@ -19,6 +18,8 @@ from ._commit_api import CommitOperationCopy, CommitOperationDelete
 from .constants import (
     DEFAULT_REVISION,
     ENDPOINT,
+    HF_HUB_DOWNLOAD_TIMEOUT,
+    HF_HUB_ETAG_TIMEOUT,
     REPO_TYPE_MODEL,
     REPO_TYPES_MAPPING,
     REPO_TYPES_URL_PREFIXES,
@@ -123,7 +124,7 @@ class HfFileSystem(fsspec.AbstractFileSystem):
     ) -> Tuple[bool, Optional[Exception]]:
         if (repo_type, repo_id, revision) not in self._repo_and_revision_exists_cache:
             try:
-                self._api.repo_info(repo_id, revision=revision, repo_type=repo_type)
+                self._api.repo_info(repo_id, revision=revision, repo_type=repo_type, timeout=HF_HUB_ETAG_TIMEOUT)
             except (RepositoryNotFoundError, HFValidationError) as e:
                 self._repo_and_revision_exists_cache[(repo_type, repo_id, revision)] = False, e
                 self._repo_and_revision_exists_cache[(repo_type, repo_id, None)] = False, e
@@ -397,7 +398,7 @@ class HfFileSystem(fsspec.AbstractFileSystem):
                 parent_path = self._parent(cache_path_info["name"])
                 self.dircache.setdefault(parent_path, []).append(cache_path_info)
                 out.append(cache_path_info)
-        return copy.deepcopy(out)  # copy to not let users modify the dircache
+        return out
 
     def glob(self, path, **kwargs):
         # Set expand_info=False by default to get a x10 speed boost
@@ -561,7 +562,7 @@ class HfFileSystem(fsspec.AbstractFileSystem):
                 if not expand_info:
                     out = {k: out[k] for k in ["name", "size", "type"]}
         assert out is not None
-        return copy.deepcopy(out)  # copy to not let users modify the dircache
+        return out
 
     def exists(self, path, **kwargs):
         """Is there a file at the given path"""
@@ -701,7 +702,13 @@ class HfFileSystemFile(fsspec.spec.AbstractBufferedFile):
             repo_type=self.resolved_path.repo_type,
             endpoint=self.fs.endpoint,
         )
-        r = http_backoff("GET", url, headers=headers, retry_on_status_codes=(502, 503, 504))
+        r = http_backoff(
+            "GET",
+            url,
+            headers=headers,
+            retry_on_status_codes=(502, 503, 504),
+            timeout=HF_HUB_DOWNLOAD_TIMEOUT,
+        )
         hf_raise_for_status(r)
         return r.content
 
@@ -800,6 +807,7 @@ class HfFileSystemStreamFile(fsspec.spec.AbstractBufferedFile):
                 headers=self.fs._api._build_hf_headers(),
                 retry_on_status_codes=(502, 503, 504),
                 stream=True,
+                timeout=HF_HUB_DOWNLOAD_TIMEOUT,
             )
             hf_raise_for_status(self.response)
         try:
@@ -821,6 +829,7 @@ class HfFileSystemStreamFile(fsspec.spec.AbstractBufferedFile):
                 headers={"Range": "bytes=%d-" % self.loc, **self.fs._api._build_hf_headers()},
                 retry_on_status_codes=(502, 503, 504),
                 stream=True,
+                timeout=HF_HUB_DOWNLOAD_TIMEOUT,
             )
             hf_raise_for_status(self.response)
             try:

huggingface_hub/hub_mixin.py

@@ -3,7 +3,7 @@ import json
 import os
 from dataclasses import asdict, dataclass, is_dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, TypeVar, Union, get_args
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union, get_args
 
 from .constants import CONFIG_NAME, PYTORCH_WEIGHTS_NAME, SAFETENSORS_SINGLE_FILE
 from .file_download import hf_hub_download
@@ -19,7 +19,6 @@ from .utils import (
     logging,
     validate_hf_hub_args,
 )
-from .utils._deprecation import _deprecate_arguments
 
 
 if TYPE_CHECKING:
@@ -37,6 +36,12 @@ logger = logging.get_logger(__name__)
 
 # Generic variable that is either ModelHubMixin or a subclass thereof
 T = TypeVar("T", bound="ModelHubMixin")
+# Generic variable to represent an args type
+ARGS_T = TypeVar("ARGS_T")
+ENCODER_T = Callable[[ARGS_T], Any]
+DECODER_T = Callable[[Any], ARGS_T]
+CODER_T = Tuple[ENCODER_T, DECODER_T]
+
 
 DEFAULT_MODEL_CARD = """
 ---
@@ -45,16 +50,16 @@ DEFAULT_MODEL_CARD = """
 {{ card_data }}
 ---
 
-This model has been pushed to the Hub using **{{ library_name }}**:
-- Repo: {{ repo_url | default("[More Information Needed]", true) }}
+This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
+- Library: {{ 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
+    model_card_template: str
+    model_card_data: ModelCardData
     repo_url: Optional[str] = None
     docs_url: Optional[str] = None
 
@@ -71,15 +76,37 @@ class ModelHubMixin:
     `__init__` but to the class definition itself. This is useful to define metadata about the library integrating
     [`ModelHubMixin`].
 
+    For more details on how to integrate the mixin with your library, checkout the [integration guide](../guides/integrations).
+
     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.
+        model_card_template (`str`, *optional*):
+            Template of the model card. Used to generate model card. Defaults to a generic template.
+        languages (`List[str]`, *optional*):
+            Languages supported by the library. Used to generate model card.
+        library_name (`str`, *optional*):
+            Name of the library integrating ModelHubMixin. Used to generate model card.
+        license (`str`, *optional*):
+            License of the library integrating ModelHubMixin. Used to generate model card.
+            E.g: "apache-2.0"
+        license_name (`str`, *optional*):
+            Name of the library integrating ModelHubMixin. Used to generate model card.
+            Only used if `license` is set to `other`.
+            E.g: "coqui-public-model-license".
+        license_link (`str`, *optional*):
+            URL to the license of the library integrating ModelHubMixin. Used to generate model card.
+            Only used if `license` is set to `other` and `license_name` is set.
+            E.g: "https://coqui.ai/cpml".
+        pipeline_tag (`str`, *optional*):
+            Tag of the pipeline. Used to generate model card. E.g. "text-classification".
+        tags (`List[str]`, *optional*):
+            Tags to be added to the model card. Used to generate model card. E.g. ["x-custom-tag", "arxiv:2304.12244"]
+        coders (`Dict[Type, Tuple[Callable, Callable]]`, *optional*):
+            Dictionary of custom types and their encoders/decoders. Used to encode/decode arguments that are not
+            jsonable by default. E.g dataclasses, argparse.Namespace, OmegaConf, etc.
 
     Example:
 
@@ -90,7 +117,7 @@ class ModelHubMixin:
     >>> class MyCustomModel(
     ...         ModelHubMixin,
     ...         library_name="my-library",
-    ...         tags=["x-custom-tag"],
+    ...         tags=["x-custom-tag", "arxiv:2304.12244"],
     ...         repo_url="https://github.com/huggingface/my-cool-library",
     ...         docs_url="https://huggingface.co/docs/my-cool-library",
     ...         # ^ optional metadata to generate model card
@@ -110,7 +137,7 @@ class ModelHubMixin:
     ...         pretrained_model_name_or_path: Union[str, Path],
     ...         *,
     ...         force_download: bool = False,
-    ...         resume_download: bool = False,
+    ...         resume_download: Optional[bool] = None,
     ...         proxies: Optional[Dict] = None,
     ...         token: Optional[Union[str, bool]] = None,
     ...         cache_dir: Optional[Union[str, Path]] = None,
@@ -131,8 +158,8 @@ class ModelHubMixin:
 
     # Download and initialize weights from the Hub
     >>> reloaded_model = MyCustomModel.from_pretrained("username/my-awesome-model")
-    >>> reloaded_model._hub_mixin_config
-    {"size": 256, "device": "gpu"}
+    >>> reloaded_model.size
+    256
 
     # Model card has been correctly populated
     >>> from huggingface_hub import ModelCard
@@ -148,18 +175,36 @@ class ModelHubMixin:
     # ^ 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
+    _hub_mixin_inject_config: bool  # whether `_from_pretrained` expects `config` or not
+    _hub_mixin_init_parameters: Dict[str, inspect.Parameter]  # __init__ parameters
+    _hub_mixin_jsonable_default_values: Dict[str, Any]  # default values for __init__ parameters
+    _hub_mixin_jsonable_custom_types: Tuple[Type, ...]  # custom types that can be encoded/decoded
+    _hub_mixin_coders: Dict[Type, CODER_T]  # encoders/decoders for custom types
     # ^ internal values to handle config
 
     def __init_subclass__(
         cls,
         *,
-        library_name: Optional[str] = None,
-        tags: Optional[List[str]] = None,
+        # Generic info for model card
         repo_url: Optional[str] = None,
         docs_url: Optional[str] = None,
+        # Model card template
+        model_card_template: str = DEFAULT_MODEL_CARD,
+        # Model card metadata
+        languages: Optional[List[str]] = None,
+        library_name: Optional[str] = None,
+        license: Optional[str] = None,
+        license_name: Optional[str] = None,
+        license_link: Optional[str] = None,
+        pipeline_tag: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+        # How to encode/decode arguments with custom type into a JSON config?
+        coders: Optional[
+            Dict[Type, CODER_T]
+            # Key is a type.
+            # Value is a tuple (encoder, decoder).
+            # Example: {MyCustomType: (lambda x: x.value, lambda data: MyCustomType(data))}
+        ] = None,
     ) -> None:
         """Inspect __init__ signature only once when subclassing + handle modelcard."""
         super().__init_subclass__()
@@ -168,18 +213,30 @@ class ModelHubMixin:
         tags = tags or []
         tags.append("model_hub_mixin")
         cls._hub_mixin_info = MixinInfo(
-            library_name=library_name,
-            tags=tags,
+            model_card_template=model_card_template,
             repo_url=repo_url,
             docs_url=docs_url,
+            model_card_data=ModelCardData(
+                languages=languages,
+                library_name=library_name,
+                license=license,
+                license_name=license_name,
+                license_link=license_link,
+                pipeline_tag=pipeline_tag,
+                tags=tags,
+            ),
         )
 
+        # Handle encoders/decoders for args
+        cls._hub_mixin_coders = coders or {}
+        cls._hub_mixin_jsonable_custom_types = tuple(cls._hub_mixin_coders.keys())
+
         # 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
+            param.name: cls._encode_arg(param.default)
             for param in cls._hub_mixin_init_parameters.values()
-            if param.default is not inspect.Parameter.empty and is_jsonable(param.default)
+            if param.default is not inspect.Parameter.empty and cls._is_jsonable(param.default)
         }
         cls._hub_mixin_inject_config = "config" in inspect.signature(cls._from_pretrained).parameters
 
@@ -220,7 +277,11 @@ class ModelHubMixin:
             # default values
             **cls._hub_mixin_jsonable_default_values,
             # passed values
-            **{key: value for key, value in passed_values.items() if is_jsonable(value)},
+            **{
+                key: cls._encode_arg(value)  # Encode custom types as jsonable value
+                for key, value in passed_values.items()
+                if instance._is_jsonable(value)  # Only if jsonable or we have a custom encoder
+            },
         }
         init_config.pop("config", {})
 
@@ -234,6 +295,29 @@ class ModelHubMixin:
             instance._hub_mixin_config = init_config
         return instance
 
+    @classmethod
+    def _is_jsonable(cls, value: Any) -> bool:
+        """Check if a value is JSON serializable."""
+        if isinstance(value, cls._hub_mixin_jsonable_custom_types):
+            return True
+        return is_jsonable(value)
+
+    @classmethod
+    def _encode_arg(cls, arg: Any) -> Any:
+        """Encode an argument into a JSON serializable format."""
+        for type_, (encoder, _) in cls._hub_mixin_coders.items():
+            if isinstance(arg, type_):
+                return encoder(arg)
+        return arg
+
+    @classmethod
+    def _decode_arg(cls, expected_type: Type[ARGS_T], value: Any) -> ARGS_T:
+        """Decode a JSON serializable value into an argument."""
+        for type_, (_, decoder) in cls._hub_mixin_coders.items():
+            if issubclass(expected_type, type_):
+                return decoder(value)
+        return value
+
     def save_pretrained(
         self,
         save_directory: Union[str, Path],
@@ -258,6 +342,8 @@ class ModelHubMixin:
                 not provided.
             kwargs:
                 Additional key word arguments passed along to the [`~ModelHubMixin.push_to_hub`] method.
+        Returns:
+            `str` or `None`: url of the commit on the Hub if `push_to_hub=True`, `None` otherwise.
         """
         save_directory = Path(save_directory)
         save_directory.mkdir(parents=True, exist_ok=True)
@@ -314,7 +400,7 @@ class ModelHubMixin:
         pretrained_model_name_or_path: Union[str, Path],
         *,
         force_download: bool = False,
-        resume_download: bool = False,
+        resume_download: Optional[bool] = None,
         proxies: Optional[Dict] = None,
         token: Optional[Union[str, bool]] = None,
         cache_dir: Optional[Union[str, Path]] = None,
@@ -336,8 +422,6 @@ class ModelHubMixin:
             force_download (`bool`, *optional*, defaults to `False`):
                 Whether to force (re-)downloading the model weights and configuration files from the Hub, overriding
                 the existing cache.
-            resume_download (`bool`, *optional*, defaults to `False`):
-                Whether to delete incompletely received files. Will attempt to resume the download if such a file exists.
             proxies (`Dict[str, str]`, *optional*):
                 A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
                 'http://hostname': 'foo.bar:4012'}`. The proxies are used on every request.
@@ -380,6 +464,13 @@ class ModelHubMixin:
             with open(config_file, "r", encoding="utf-8") as f:
                 config = json.load(f)
 
+            # Decode custom types in config
+            for key, value in config.items():
+                if key in cls._hub_mixin_init_parameters:
+                    expected_type = cls._hub_mixin_init_parameters[key].annotation
+                    if expected_type is not inspect.Parameter.empty:
+                        config[key] = cls._decode_arg(expected_type, value)
+
             # 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:
@@ -445,7 +536,7 @@ class ModelHubMixin:
         cache_dir: Optional[Union[str, Path]],
         force_download: bool,
         proxies: Optional[Dict],
-        resume_download: bool,
+        resume_download: Optional[bool],
         local_files_only: bool,
         token: Optional[Union[str, bool]],
         **model_kwargs,
@@ -468,8 +559,6 @@ class ModelHubMixin:
             force_download (`bool`, *optional*, defaults to `False`):
                 Whether to force (re-)downloading the model weights and configuration files from the Hub, overriding
                 the existing cache.
-            resume_download (`bool`, *optional*, defaults to `False`):
-                Whether to delete incompletely received files. Will attempt to resume the download if such a file exists.
             proxies (`Dict[str, str]`, *optional*):
                 A dictionary of proxy servers to use by protocol or endpoint (e.g., `{'http': 'foo.bar:3128',
                 'http://hostname': 'foo.bar:4012'}`).
@@ -485,11 +574,6 @@ 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,
@@ -504,8 +588,6 @@ class ModelHubMixin:
         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.
@@ -523,8 +605,6 @@ class ModelHubMixin:
                 Message to commit while pushing.
             private (`bool`, *optional*, defaults to `False`):
                 Whether the repository created should be private.
-            api_endpoint (`str`, *optional*):
-                The API endpoint to use when pushing the model to the hub.
             token (`str`, *optional*):
                 The token to use as HTTP bearer authorization for remote files. By default, it will use the token
                 cached when running `huggingface-cli login`.
@@ -542,7 +622,7 @@ class ModelHubMixin:
         Returns:
             The url of the commit of your model in the given repository.
         """
-        api = HfApi(endpoint=api_endpoint, token=token)
+        api = HfApi(token=token)
         repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id
 
         # Push the files to the repo in a single commit
@@ -563,8 +643,10 @@ class ModelHubMixin:
 
     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,
+            card_data=self._hub_mixin_info.model_card_data,
+            template_str=self._hub_mixin_info.model_card_template,
+            repo_url=self._hub_mixin_info.repo_url,
+            docs_url=self._hub_mixin_info.docs_url,
         )
         return card
 
@@ -575,6 +657,8 @@ class PyTorchModelHubMixin(ModelHubMixin):
     is set in evaluation mode by default using `model.eval()` (dropout modules are deactivated). To train the model,
     you should first set it back in training mode with `model.train()`.
 
+    See [`ModelHubMixin`] for more details on how to use the mixin.
+
     Example:
 
     ```python
@@ -632,7 +716,7 @@ class PyTorchModelHubMixin(ModelHubMixin):
         cache_dir: Optional[Union[str, Path]],
         force_download: bool,
         proxies: Optional[Dict],
-        resume_download: bool,
+        resume_download: Optional[bool],
         local_files_only: bool,
         token: Union[str, bool, None],
         map_location: str = "cpu",