huggingface-hub 0.23.1__py3-none-any.whl → 0.23.2__py3-none-any.whl

huggingface_hub/__init__.py

@@ -46,7 +46,7 @@ import sys
 from typing import TYPE_CHECKING
 
 
-__version__ = "0.23.1"
+__version__ = "0.23.2"
 
 # Alphabetical order of definitions is ensured in tests
 # WARNING: any comment added in this dictionary definition will be lost when

huggingface_hub/serialization/_base.py

@@ -14,7 +14,7 @@
 """Contains helpers to split tensors into shards."""
 
 from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, List, Optional, TypeVar
+from typing import Any, Callable, Dict, List, Optional, TypeVar, Union
 
 from .. import logging
 
@@ -46,7 +46,7 @@ def split_state_dict_into_shards_factory(
     get_tensor_size: TensorSizeFn_T,
     get_storage_id: StorageIDFn_T = lambda tensor: None,
     filename_pattern: str = FILENAME_PATTERN,
-    max_shard_size: int = MAX_SHARD_SIZE,
+    max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
 ) -> StateDictSplit:
     """
     Split a model state dictionary in shards so that each shard is smaller than a given size.
@@ -89,6 +89,9 @@ def split_state_dict_into_shards_factory(
     current_shard_size = 0
     total_size = 0
 
+    if isinstance(max_shard_size, str):
+        max_shard_size = parse_size_to_int(max_shard_size)
+
     for key, tensor in state_dict.items():
         # when bnb serialization is used the weights in the state dict can be strings
         # check: https://github.com/huggingface/transformers/pull/24416 for more details
@@ -167,3 +170,44 @@ def split_state_dict_into_shards_factory(
         filename_to_tensors=filename_to_tensors,
         tensor_to_filename=tensor_name_to_filename,
     )
+
+
+SIZE_UNITS = {
+    "TB": 10**12,
+    "GB": 10**9,
+    "MB": 10**6,
+    "KB": 10**3,
+}
+
+
+def parse_size_to_int(size_as_str: str) -> int:
+    """
+    Parse a size expressed as a string with digits and unit (like `"5MB"`) to an integer (in bytes).
+
+    Supported units are "TB", "GB", "MB", "KB".
+
+    Args:
+        size_as_str (`str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> parse_size_to_int("5MB")
+    5000000
+    ```
+    """
+    size_as_str = size_as_str.strip()
+
+    # Parse unit
+    unit = size_as_str[-2:].upper()
+    if unit not in SIZE_UNITS:
+        raise ValueError(f"Unit '{unit}' not supported. Supported units are TB, GB, MB, KB. Got '{size_as_str}'.")
+    multiplier = SIZE_UNITS[unit]
+
+    # Parse value
+    try:
+        value = float(size_as_str[:-2].strip())
+    except ValueError as e:
+        raise ValueError(f"Could not parse the size value from '{size_as_str}': {e}") from e
+
+    return int(value * multiplier)

huggingface_hub/serialization/_numpy.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Contains numpy-specific helpers."""
 
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING, Dict, Union
 
 from ._base import FILENAME_PATTERN, MAX_SHARD_SIZE, StateDictSplit, split_state_dict_into_shards_factory
 
@@ -26,7 +26,7 @@ def split_numpy_state_dict_into_shards(
     state_dict: Dict[str, "np.ndarray"],
     *,
     filename_pattern: str = FILENAME_PATTERN,
-    max_shard_size: int = MAX_SHARD_SIZE,
+    max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
 ) -> StateDictSplit:
     """
     Split a model state dictionary in shards so that each shard is smaller than a given size.

huggingface_hub/serialization/_tensorflow.py

@@ -15,7 +15,7 @@
 
 import math
 import re
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING, Dict, Union
 
 from ._base import MAX_SHARD_SIZE, StateDictSplit, split_state_dict_into_shards_factory
 
@@ -28,7 +28,7 @@ def split_tf_state_dict_into_shards(
     state_dict: Dict[str, "tf.Tensor"],
     *,
     filename_pattern: str = "tf_model{suffix}.h5",
-    max_shard_size: int = MAX_SHARD_SIZE,
+    max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
 ) -> StateDictSplit:
     """
     Split a model state dictionary in shards so that each shard is smaller than a given size.

huggingface_hub/serialization/_torch.py

@@ -15,7 +15,7 @@
 
 import importlib
 from functools import lru_cache
-from typing import TYPE_CHECKING, Dict, Tuple
+from typing import TYPE_CHECKING, Dict, Tuple, Union
 
 from ._base import FILENAME_PATTERN, MAX_SHARD_SIZE, StateDictSplit, split_state_dict_into_shards_factory
 
@@ -28,7 +28,7 @@ def split_torch_state_dict_into_shards(
     state_dict: Dict[str, "torch.Tensor"],
     *,
     filename_pattern: str = FILENAME_PATTERN,
-    max_shard_size: int = MAX_SHARD_SIZE,
+    max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
 ) -> StateDictSplit:
     """
     Split a model state dictionary in shards so that each shard is smaller than a given size.
@@ -67,7 +67,7 @@ def split_torch_state_dict_into_shards(
 
     >>> 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.values():
+    ...     for filename, tensors in state_dict_split.filename_to_tensors.items():
     ...         shard = {tensor: state_dict[tensor] for tensor in tensors}
     ...         safe_save_file(
     ...             shard,

{huggingface_hub-0.23.1.dist-info → huggingface_hub-0.23.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: huggingface-hub
-Version: 0.23.1
+Version: 0.23.2
 Summary: Client library to download and publish models, datasets and other repos on the huggingface.co hub
 Home-page: https://github.com/huggingface/huggingface_hub
 Author: Hugging Face, Inc.

{huggingface_hub-0.23.1.dist-info → huggingface_hub-0.23.2.dist-info}/RECORD

@@ -1,4 +1,4 @@
-huggingface_hub/__init__.py,sha256=Ve0gXDRzk-i00rJGQFbSMQinDOnEl5g_P86_l9l5zpk,32692
+huggingface_hub/__init__.py,sha256=Q1uTvSPbnh_bO8QcgwbIbiQgB-XsnDbGwa7tpWxNl7w,32692
 huggingface_hub/_commit_api.py,sha256=Z1sQnJx1xWfspsX6vS8eGTmr-9QujIoItjbnJVVyyCQ,29299
 huggingface_hub/_commit_scheduler.py,sha256=nlJS_vnLb8i92NLrRwJX8Mg9QZ7f3kfLbLlQuEd5YjU,13647
 huggingface_hub/_inference_endpoints.py,sha256=rBx6xgnSJq0JtntF1_zphj7NsCmduICqgZfmvscdE_w,15667
@@ -74,10 +74,10 @@ huggingface_hub/inference/_generated/types/zero_shot_classification.py,sha256=u6
 huggingface_hub/inference/_generated/types/zero_shot_image_classification.py,sha256=qVH6Ms0FjF8TraGy4BYiS8lmvGq9xiIDdXqGFynLHMA,1689
 huggingface_hub/inference/_generated/types/zero_shot_object_detection.py,sha256=PU4OOlQ2aAOosW2JlG2Z27MEQpmE6BxcygH_ns3w1KQ,1662
 huggingface_hub/serialization/__init__.py,sha256=W74TaCtYnMfpvGEQr1SS-OBmqPUFnM9AeWT9hTJCG9Y,910
-huggingface_hub/serialization/_base.py,sha256=AgO-16i-vyosbERnLSCFYgaXbVqQDM7xfIne8gsWrLQ,7133
-huggingface_hub/serialization/_numpy.py,sha256=idULJp1js6L6E8o-MiGVqNa4lBfXS2cfAmqivnpsaYs,2671
-huggingface_hub/serialization/_tensorflow.py,sha256=Rf4kw1NYxEaoUXB8aLtQLHrTjgobaEAJdzO0w0kbP58,3559
-huggingface_hub/serialization/_torch.py,sha256=xYR6e_G9laMTroWLiQRABSuloTQuuRSQNyYHdT_rmXU,7687
+huggingface_hub/serialization/_base.py,sha256=2wxdid6ee8RASEKhCkpNdP8Kj9x4dRm6j8h72L1AtFQ,8239
+huggingface_hub/serialization/_numpy.py,sha256=E-boJoUuDdyMTaAMRIiHha7F9GH9YhFNQPEioG4UkaY,2690
+huggingface_hub/serialization/_tensorflow.py,sha256=4Wf_wzmLSzZua9hGGmArfngDzz3yw19PWJMdTT76uxc,3578
+huggingface_hub/serialization/_torch.py,sha256=t-pTq4O3NpAprVJIojtC8Rq-kNJ889IluJtJtoLoqVk,7705
 huggingface_hub/templates/datasetcard_template.md,sha256=W-EMqR6wndbrnZorkVv56URWPG49l7MATGeI015kTvs,5503
 huggingface_hub/templates/modelcard_template.md,sha256=4AqArS3cqdtbit5Bo-DhjcnDFR-pza5hErLLTPM4Yuc,6870
 huggingface_hub/utils/__init__.py,sha256=44yhxTtWsuMGrZcALK-3UuVazGBtc94z9nZwLmLnu8w,3589
@@ -107,9 +107,9 @@ huggingface_hub/utils/insecure_hashlib.py,sha256=OjxlvtSQHpbLp9PWSrXBDJ0wHjxCBU-
 huggingface_hub/utils/logging.py,sha256=Cp03s0uEl3kDM9XHQW9a8GAoExODQ-e7kEtgMt-_To8,4728
 huggingface_hub/utils/sha.py,sha256=QLlIwPCyz46MmUc_4L8xl87KfYoBks9kPgsMZ5JCz-o,902
 huggingface_hub/utils/tqdm.py,sha256=x35PqUA8bBBztPrqhv87Y_TGl5CdlfBs4pe6k1YyDJ8,9390
-huggingface_hub-0.23.1.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-huggingface_hub-0.23.1.dist-info/METADATA,sha256=ZRJ_sSeAnUgh-JttdzKOwQwGw0Biu1SSckrYJanqJro,12994
-huggingface_hub-0.23.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-huggingface_hub-0.23.1.dist-info/entry_points.txt,sha256=Y3Z2L02rBG7va_iE6RPXolIgwOdwUFONyRN3kXMxZ0g,131
-huggingface_hub-0.23.1.dist-info/top_level.txt,sha256=8KzlQJAY4miUvjAssOAJodqKOw3harNzuiwGQ9qLSSk,16
-huggingface_hub-0.23.1.dist-info/RECORD,,
+huggingface_hub-0.23.2.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+huggingface_hub-0.23.2.dist-info/METADATA,sha256=jeheWjcbLyu4hs7FgnFpv_6VOHF6aTNe7X-5TzYEtaA,12994
+huggingface_hub-0.23.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+huggingface_hub-0.23.2.dist-info/entry_points.txt,sha256=Y3Z2L02rBG7va_iE6RPXolIgwOdwUFONyRN3kXMxZ0g,131
+huggingface_hub-0.23.2.dist-info/top_level.txt,sha256=8KzlQJAY4miUvjAssOAJodqKOw3harNzuiwGQ9qLSSk,16
+huggingface_hub-0.23.2.dist-info/RECORD,,