huggingface-hub 0.29.0rc2__py3-none-any.whl → 1.1.3__py3-none-any.whl

huggingface_hub/serialization/_torch.py

@@ -20,7 +20,7 @@ import re
 from collections import defaultdict, namedtuple
 from functools import lru_cache
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, NamedTuple, Optional, Set, Tuple, Union
+from typing import TYPE_CHECKING, Any, Iterable, NamedTuple, Optional, Union
 
 from packaging import version
 
@@ -43,10 +43,10 @@ def save_torch_model(
     filename_pattern: Optional[str] = None,
     force_contiguous: bool = True,
     max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
-    metadata: Optional[Dict[str, str]] = None,
+    metadata: Optional[dict[str, str]] = None,
     safe_serialization: bool = True,
     is_main_process: bool = True,
-    shared_tensors_to_discard: Optional[List[str]] = None,
+    shared_tensors_to_discard: Optional[list[str]] = None,
 ):
     """
     Saves a given torch model to disk, handling sharding and shared tensors issues.
@@ -63,18 +63,12 @@ def save_torch_model(
 
     Before saving the model, the `save_directory` is cleaned from any previous shard files.
 
-    <Tip warning={true}>
+    > [!WARNING]
+    > If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a
+    > size greater than `max_shard_size`.
 
-    If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a
-    size greater than `max_shard_size`.
-
-    </Tip>
-
-    <Tip warning={true}>
-
-    If your model is a `transformers.PreTrainedModel`, you should pass `model._tied_weights_keys` as `shared_tensors_to_discard` to properly handle shared tensors saving. This ensures the correct duplicate tensors are discarded during saving.
-
-    </Tip>
+    > [!WARNING]
+    > If your model is a `transformers.PreTrainedModel`, you should pass `model._tied_weights_keys` as `shared_tensors_to_discard` to properly handle shared tensors saving. This ensures the correct duplicate tensors are discarded during saving.
 
     Args:
         model (`torch.nn.Module`):
@@ -92,7 +86,7 @@ def save_torch_model(
             that reason. Defaults to `True`.
         max_shard_size (`int` or `str`, *optional*):
             The maximum size of each shard, in bytes. Defaults to 5GB.
-        metadata (`Dict[str, str]`, *optional*):
+        metadata (`dict[str, str]`, *optional*):
             Extra information to save along with the model. Some metadata will be added for each dropped tensors.
             This information will not be enough to recover the entire shared structure but might help understanding
             things.
@@ -104,7 +98,7 @@ def save_torch_model(
             Whether the process calling this is the main process or not. Useful when in distributed training like
             TPUs and need to call this function from all processes. In this case, set `is_main_process=True` only on
             the main process to avoid race conditions. Defaults to True.
-        shared_tensors_to_discard (`List[str]`, *optional*):
+        shared_tensors_to_discard (`list[str]`, *optional*):
             List of tensor names to drop when saving shared tensors. If not provided and shared tensors are
             detected, it will drop the first name alphabetically.
 
@@ -137,16 +131,16 @@ def save_torch_model(
 
 
 def save_torch_state_dict(
-    state_dict: Dict[str, "torch.Tensor"],
+    state_dict: dict[str, "torch.Tensor"],
     save_directory: Union[str, Path],
     *,
     filename_pattern: Optional[str] = None,
     force_contiguous: bool = True,
     max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
-    metadata: Optional[Dict[str, str]] = None,
+    metadata: Optional[dict[str, str]] = None,
     safe_serialization: bool = True,
     is_main_process: bool = True,
-    shared_tensors_to_discard: Optional[List[str]] = None,
+    shared_tensors_to_discard: Optional[list[str]] = None,
 ) -> None:
     """
     Save a model state dictionary to the disk, handling sharding and shared tensors issues.
@@ -163,21 +157,15 @@ def save_torch_state_dict(
 
     Before saving the model, the `save_directory` is cleaned from any previous shard files.
 
-    <Tip warning={true}>
-
-    If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a
-    size greater than `max_shard_size`.
-
-    </Tip>
+    > [!WARNING]
+    > If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a
+    > size greater than `max_shard_size`.
 
-    <Tip warning={true}>
-
-    If your model is a `transformers.PreTrainedModel`, you should pass `model._tied_weights_keys` as `shared_tensors_to_discard` to properly handle shared tensors saving. This ensures the correct duplicate tensors are discarded during saving.
-
-    </Tip>
+    > [!WARNING]
+    > If your model is a `transformers.PreTrainedModel`, you should pass `model._tied_weights_keys` as `shared_tensors_to_discard` to properly handle shared tensors saving. This ensures the correct duplicate tensors are discarded during saving.
 
     Args:
-        state_dict (`Dict[str, torch.Tensor]`):
+        state_dict (`dict[str, torch.Tensor]`):
             The state dictionary to save.
         save_directory (`str` or `Path`):
             The directory in which the model will be saved.
@@ -192,7 +180,7 @@ def save_torch_state_dict(
             that reason. Defaults to `True`.
         max_shard_size (`int` or `str`, *optional*):
             The maximum size of each shard, in bytes. Defaults to 5GB.
-        metadata (`Dict[str, str]`, *optional*):
+        metadata (`dict[str, str]`, *optional*):
             Extra information to save along with the model. Some metadata will be added for each dropped tensors.
             This information will not be enough to recover the entire shared structure but might help understanding
             things.
@@ -204,7 +192,7 @@ def save_torch_state_dict(
             Whether the process calling this is the main process or not. Useful when in distributed training like
             TPUs and need to call this function from all processes. In this case, set `is_main_process=True` only on
             the main process to avoid race conditions. Defaults to True.
-        shared_tensors_to_discard (`List[str]`, *optional*):
+        shared_tensors_to_discard (`list[str]`, *optional*):
             List of tensor names to drop when saving shared tensors. If not provided and shared tensors are
             detected, it will drop the first name alphabetically.
 
@@ -246,7 +234,7 @@ def save_torch_state_dict(
             shared_tensors_to_discard=shared_tensors_to_discard,
         )
     else:
-        from torch import save as save_file_fn  # type: ignore[assignment]
+        from torch import save as save_file_fn  # type: ignore[assignment, no-redef]
 
         logger.warning(
             "You are using unsafe serialization. Due to security reasons, it is recommended not to load "
@@ -278,7 +266,7 @@ def save_torch_state_dict(
     safe_file_kwargs = {"metadata": per_file_metadata} if safe_serialization else {}
     for filename, tensors in state_dict_split.filename_to_tensors.items():
         shard = {tensor: state_dict[tensor] for tensor in tensors}
-        save_file_fn(shard, os.path.join(save_directory, filename), **safe_file_kwargs)
+        save_file_fn(shard, os.path.join(save_directory, filename), **safe_file_kwargs)  # ty: ignore[invalid-argument-type]
         logger.debug(f"Shard saved to {filename}")
 
     # Save the index (if any)
@@ -300,7 +288,7 @@ def save_torch_state_dict(
 
 
 def split_torch_state_dict_into_shards(
-    state_dict: Dict[str, "torch.Tensor"],
+    state_dict: dict[str, "torch.Tensor"],
     *,
     filename_pattern: str = constants.SAFETENSORS_WEIGHTS_FILE_PATTERN,
     max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
@@ -314,22 +302,16 @@ def split_torch_state_dict_into_shards(
     [6+2+2GB], [6+2GB], [6GB].
 
 
-    <Tip>
-
-    To save a model state dictionary to the disk, see [`save_torch_state_dict`]. This helper uses
-    `split_torch_state_dict_into_shards` under the hood.
-
-    </Tip>
-
-    <Tip warning={true}>
+    > [!TIP]
+    > To save a model state dictionary to the disk, see [`save_torch_state_dict`]. This helper uses
+    > `split_torch_state_dict_into_shards` under the hood.
 
-    If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a
-    size greater than `max_shard_size`.
-
-    </Tip>
+    > [!WARNING]
+    > If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a
+    > size greater than `max_shard_size`.
 
     Args:
-        state_dict (`Dict[str, torch.Tensor]`):
+        state_dict (`dict[str, torch.Tensor]`):
             The state dictionary to save.
         filename_pattern (`str`, *optional*):
             The pattern to generate the files names in which the model will be saved. Pattern must be a string that
@@ -348,7 +330,7 @@ def split_torch_state_dict_into_shards(
     >>> from safetensors.torch import save_file as safe_save_file
     >>> from huggingface_hub import split_torch_state_dict_into_shards
 
-    >>> def save_state_dict(state_dict: Dict[str, torch.Tensor], save_directory: str):
+    >>> def save_state_dict(state_dict: dict[str, torch.Tensor], save_directory: str):
     ...     state_dict_split = split_torch_state_dict_into_shards(state_dict)
     ...     for filename, tensors in state_dict_split.filename_to_tensors.items():
     ...         shard = {tensor: state_dict[tensor] for tensor in tensors}
@@ -560,7 +542,7 @@ def load_state_dict_from_file(
     map_location: Optional[Union[str, "torch.device"]] = None,
     weights_only: bool = False,
     mmap: bool = False,
-) -> Union[Dict[str, "torch.Tensor"], Any]:
+) -> Union[dict[str, "torch.Tensor"], Any]:
     """
     Loads a checkpoint file, handling both safetensors and pickle checkpoint formats.
 
@@ -580,7 +562,7 @@ def load_state_dict_from_file(
             loading safetensors files, as the `safetensors` library uses memory mapping by default.
 
     Returns:
-        `Union[Dict[str, "torch.Tensor"], Any]`: The loaded checkpoint.
+        `Union[dict[str, "torch.Tensor"], Any]`: The loaded checkpoint.
             - For safetensors files: always returns a dictionary mapping parameter names to tensors.
             - For pickle files: returns any Python object that was pickled (commonly a state dict, but could be
               an entire model, optimizer state, or any other Python object).
@@ -700,12 +682,21 @@ def _validate_keys_for_strict_loading(
         raise RuntimeError(error_message)
 
 
-def _get_unique_id(tensor: "torch.Tensor") -> Union[int, Tuple[Any, ...]]:
+def _get_unique_id(tensor: "torch.Tensor") -> Union[int, tuple[Any, ...]]:
     """Returns a unique id for plain tensor
     or a (potentially nested) Tuple of unique id for the flattened Tensor
     if the input is a wrapper tensor subclass Tensor
     """
 
+    try:
+        from torch.distributed.tensor import DTensor
+
+        if isinstance(tensor, DTensor):
+            local_tensor = tensor.to_local()
+            return local_tensor.storage().data_ptr()
+    except ImportError:
+        pass
+
     try:
         # for torch 2.1 and above we can also handle tensor subclasses
         from torch.utils._python_dispatch import is_traceable_wrapper_subclass
@@ -715,11 +706,11 @@ def _get_unique_id(tensor: "torch.Tensor") -> Union[int, Tuple[Any, ...]]:
             return tuple(_get_unique_id(getattr(tensor, attr)) for attr in attrs)
 
     except ImportError:
-        # for torch version less than 2.1, we can fallback to original implementation
+        # for torch version less than 2.1, we can fall back to original implementation
         pass
 
     if tensor.device.type == "xla" and is_torch_tpu_available():
-        # NOTE: xla tensors dont have storage
+        # NOTE: xla tensors don't have storage
         # use some other unique id to distinguish.
         # this is a XLA tensor, it must be created using torch_xla's
         # device. So the following import is safe:
@@ -732,7 +723,7 @@ def _get_unique_id(tensor: "torch.Tensor") -> Union[int, Tuple[Any, ...]]:
     return unique_id
 
 
-def get_torch_storage_id(tensor: "torch.Tensor") -> Optional[Tuple["torch.device", Union[int, Tuple[Any, ...]], int]]:
+def get_torch_storage_id(tensor: "torch.Tensor") -> Optional[tuple["torch.device", Union[int, tuple[Any, ...]], int]]:
     """
     Return unique identifier to a tensor storage.
 
@@ -753,6 +744,15 @@ def get_torch_storage_size(tensor: "torch.Tensor") -> int:
     """
     Taken from https://github.com/huggingface/safetensors/blob/08db34094e9e59e2f9218f2df133b7b4aaff5a99/bindings/python/py_src/safetensors/torch.py#L31C1-L41C59
     """
+    try:
+        from torch.distributed.tensor import DTensor
+
+        if isinstance(tensor, DTensor):
+            # this returns the size of the FULL tensor in bytes
+            return tensor.nbytes
+    except ImportError:
+        pass
+
     try:
         # for torch 2.1 and above we can also handle tensor subclasses
         from torch.utils._python_dispatch import is_traceable_wrapper_subclass
@@ -761,7 +761,7 @@ def get_torch_storage_size(tensor: "torch.Tensor") -> int:
             attrs, _ = tensor.__tensor_flatten__()  # type: ignore[attr-defined]
             return sum(get_torch_storage_size(getattr(tensor, attr)) for attr in attrs)
     except ImportError:
-        # for torch version less than 2.1, we can fallback to original implementation
+        # for torch version less than 2.1, we can fall back to original implementation
         pass
 
     try:
@@ -797,7 +797,7 @@ def is_torch_tpu_available(check_device=True):
     return False
 
 
-def storage_ptr(tensor: "torch.Tensor") -> Union[int, Tuple[Any, ...]]:
+def storage_ptr(tensor: "torch.Tensor") -> Union[int, tuple[Any, ...]]:
     """
     Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L11.
     """
@@ -808,7 +808,7 @@ def storage_ptr(tensor: "torch.Tensor") -> Union[int, Tuple[Any, ...]]:
         if is_traceable_wrapper_subclass(tensor):
             return _get_unique_id(tensor)  # type: ignore
     except ImportError:
-        # for torch version less than 2.1, we can fallback to original implementation
+        # for torch version less than 2.1, we can fall back to original implementation
         pass
 
     try:
@@ -823,10 +823,10 @@ def storage_ptr(tensor: "torch.Tensor") -> Union[int, Tuple[Any, ...]]:
 
 
 def _clean_state_dict_for_safetensors(
-    state_dict: Dict[str, "torch.Tensor"],
-    metadata: Dict[str, str],
+    state_dict: dict[str, "torch.Tensor"],
+    metadata: dict[str, str],
     force_contiguous: bool = True,
-    shared_tensors_to_discard: Optional[List[str]] = None,
+    shared_tensors_to_discard: Optional[list[str]] = None,
 ):
     """Remove shared tensors from state_dict and update metadata accordingly (for reloading).
 
@@ -860,7 +860,7 @@ def _end_ptr(tensor: "torch.Tensor") -> int:
     return stop
 
 
-def _filter_shared_not_shared(tensors: List[Set[str]], state_dict: Dict[str, "torch.Tensor"]) -> List[Set[str]]:
+def _filter_shared_not_shared(tensors: list[set[str]], state_dict: dict[str, "torch.Tensor"]) -> list[set[str]]:
     """
     Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L44
     """
@@ -888,7 +888,7 @@ def _filter_shared_not_shared(tensors: List[Set[str]], state_dict: Dict[str, "to
     return filtered_tensors
 
 
-def _find_shared_tensors(state_dict: Dict[str, "torch.Tensor"]) -> List[Set[str]]:
+def _find_shared_tensors(state_dict: dict[str, "torch.Tensor"]) -> list[set[str]]:
     """
     Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L69.
     """
@@ -916,7 +916,7 @@ def _is_complete(tensor: "torch.Tensor") -> bool:
             attrs, _ = tensor.__tensor_flatten__()  # type: ignore[attr-defined]
             return all(_is_complete(getattr(tensor, attr)) for attr in attrs)
     except ImportError:
-        # for torch version less than 2.1, we can fallback to original implementation
+        # for torch version less than 2.1, we can fall back to original implementation
         pass
 
     return tensor.data_ptr() == storage_ptr(tensor) and tensor.nelement() * _get_dtype_size(
@@ -925,11 +925,11 @@ def _is_complete(tensor: "torch.Tensor") -> bool:
 
 
 def _remove_duplicate_names(
-    state_dict: Dict[str, "torch.Tensor"],
+    state_dict: dict[str, "torch.Tensor"],
     *,
-    preferred_names: Optional[List[str]] = None,
-    discard_names: Optional[List[str]] = None,
-) -> Dict[str, List[str]]:
+    preferred_names: Optional[list[str]] = None,
+    discard_names: Optional[list[str]] = None,
+) -> dict[str, list[str]]:
     """
     Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L80
     """

huggingface_hub/utils/__init__.py

@@ -14,7 +14,6 @@
 # limitations under the License
 
 # ruff: noqa: F401
-
 from huggingface_hub.errors import (
     BadRequestError,
     CacheNotFound,
@@ -43,6 +42,7 @@ from ._cache_manager import (
     CachedRevisionInfo,
     DeleteCacheStrategy,
     HFCacheInfo,
+    _format_size,
     scan_cache_dir,
 )
 from ._chunk_utils import chunk_iterable
@@ -51,14 +51,18 @@ from ._experimental import experimental
 from ._fixes import SoftTemporaryDirectory, WeakFileLock, yaml_dump
 from ._git_credential import list_credential_helpers, set_git_credential, unset_git_credential
 from ._headers import build_hf_headers, get_token_to_send
-from ._hf_folder import HfFolder
 from ._http import (
-    configure_http_backend,
+    ASYNC_CLIENT_FACTORY_T,
+    CLIENT_FACTORY_T,
+    close_session,
     fix_hf_endpoint_in_url,
+    get_async_session,
     get_session,
     hf_raise_for_status,
     http_backoff,
-    reset_sessions,
+    http_stream_backoff,
+    set_async_client_factory,
+    set_client_factory,
 )
 from ._pagination import paginate
 from ._paths import DEFAULT_IGNORE_PATTERNS, FORBIDDEN_FOLDERS, filter_repo_objects
@@ -71,7 +75,6 @@ from ._runtime import (
     get_gradio_version,
     get_graphviz_version,
     get_hf_hub_version,
-    get_hf_transfer_version,
     get_jinja_version,
     get_numpy_version,
     get_pillow_version,
@@ -81,6 +84,7 @@ from ._runtime import (
     get_tensorboard_version,
     get_tf_version,
     get_torch_version,
+    installation_method,
     is_aiohttp_available,
     is_colab_enterprise,
     is_fastai_available,
@@ -89,7 +93,6 @@ from ._runtime import (
     is_google_colab,
     is_gradio_available,
     is_graphviz_available,
-    is_hf_transfer_available,
     is_jinja_available,
     is_notebook,
     is_numpy_available,
@@ -105,6 +108,22 @@ from ._runtime import (
 from ._safetensors import SafetensorsFileMetadata, SafetensorsRepoMetadata, TensorInfo
 from ._subprocess import capture_output, run_interactive_subprocess, run_subprocess
 from ._telemetry import send_telemetry
+from ._terminal import ANSI, tabulate
 from ._typing import is_jsonable, is_simple_optional_type, unwrap_simple_optional_type
-from ._validators import smoothly_deprecate_use_auth_token, validate_hf_hub_args, validate_repo_id
-from .tqdm import are_progress_bars_disabled, disable_progress_bars, enable_progress_bars, tqdm, tqdm_stream_file
+from ._validators import validate_hf_hub_args, validate_repo_id
+from ._xet import (
+    XetConnectionInfo,
+    XetFileData,
+    XetTokenType,
+    fetch_xet_connection_info_from_repo_info,
+    parse_xet_file_data_from_response,
+    refresh_xet_connection_info,
+)
+from .tqdm import (
+    are_progress_bars_disabled,
+    disable_progress_bars,
+    enable_progress_bars,
+    is_tqdm_disabled,
+    tqdm,
+    tqdm_stream_file,
+)

huggingface_hub/utils/_auth.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Contains an helper to get the token from machine (env variable, secret or config file)."""
+"""Contains a helper to get the token from machine (env variable, secret or config file)."""
 
 import configparser
 import logging
@@ -19,7 +19,7 @@ import os
 import warnings
 from pathlib import Path
 from threading import Lock
-from typing import Dict, Optional
+from typing import Optional
 
 from .. import constants
 from ._runtime import is_colab_enterprise, is_google_colab
@@ -41,7 +41,7 @@ def get_token() -> Optional[str]:
 
     Token is retrieved in priority from the `HF_TOKEN` environment variable. Otherwise, we read the token file located
     in the Hugging Face home folder. Returns None if user is not logged in. To log in, use [`login`] or
-    `huggingface-cli login`.
+    `hf auth login`.
 
     Returns:
         `str` or `None`: The token, `None` if it doesn't exist.
@@ -125,13 +125,13 @@ def _get_token_from_file() -> Optional[str]:
         return None
 
 
-def get_stored_tokens() -> Dict[str, str]:
+def get_stored_tokens() -> dict[str, str]:
     """
     Returns the parsed INI file containing the access tokens.
     The file is located at `HF_STORED_TOKENS_PATH`, defaulting to `~/.cache/huggingface/stored_tokens`.
     If the file does not exist, an empty dictionary is returned.
 
-    Returns: `Dict[str, str]`
+    Returns: `dict[str, str]`
         Key is the token name and value is the token.
     """
     tokens_path = Path(constants.HF_STORED_TOKENS_PATH)
@@ -147,12 +147,12 @@ def get_stored_tokens() -> Dict[str, str]:
     return stored_tokens
 
 
-def _save_stored_tokens(stored_tokens: Dict[str, str]) -> None:
+def _save_stored_tokens(stored_tokens: dict[str, str]) -> None:
     """
     Saves the given configuration to the stored tokens file.
 
     Args:
-        stored_tokens (`Dict[str, str]`):
+        stored_tokens (`dict[str, str]`):
             The stored tokens to save. Key is the token name and value is the token.
     """
     stored_tokens_path = Path(constants.HF_STORED_TOKENS_PATH)