safetensors 0.6.0rc0__cp38-abi3-win32.whl → 0.6.1rc0__cp38-abi3-win32.whl

safetensors/__init__.py

@@ -4,6 +4,7 @@ from ._safetensors_rust import (  # noqa: F401
     __version__,
     deserialize,
     safe_open,
+    _safe_open_handle,
     serialize,
     serialize_file,
 )

safetensors/__init__.pyi

@@ -49,7 +49,7 @@ def serialize_file(tensor_dict, filename, metadata=None):
 
     Returns:
         (`NoneType`):
-            On success return None.
+            On success return None
     """
     pass
 
@@ -68,19 +68,21 @@ class safe_open:
         device (`str`, defaults to `"cpu"`):
             The device on which you want the tensors.
     """
-
     def __init__(self, filename, framework, device=...):
         pass
+
     def __enter__(self):
         """
         Start the context manager
         """
         pass
+
     def __exit__(self, _exc_type, _exc_value, _traceback):
         """
         Exits the context manager
         """
         pass
+
     def get_slice(self, name):
         """
         Returns a full slice view object
@@ -102,6 +104,7 @@ class safe_open:
         ```
         """
         pass
+
     def get_tensor(self, name):
         """
         Returns a full tensor
@@ -124,6 +127,7 @@ class safe_open:
         ```
         """
         pass
+
     def keys(self):
         """
         Returns the names of the tensors in the file.
@@ -133,6 +137,7 @@ class safe_open:
                 The name of the tensors contained in that file
         """
         pass
+
     def metadata(self):
         """
         Return the special non tensor information in the header
@@ -143,6 +148,16 @@ class safe_open:
         """
         pass
 
+    def offset_keys(self):
+        """
+        Returns the names of the tensors in the file, ordered by offset.
+
+        Returns:
+            (`List[str]`):
+                The name of the tensors contained in that file
+        """
+        pass
+
 class SafetensorError(Exception):
     """
     Custom Python Exception for Safetensor errors.

safetensors/mlx.py

@@ -7,7 +7,9 @@ import mlx.core as mx
 from safetensors import numpy, safe_open
 
 
-def save(tensors: Dict[str, mx.array], metadata: Optional[Dict[str, str]] = None) -> bytes:
+def save(
+    tensors: Dict[str, mx.array], metadata: Optional[Dict[str, str]] = None
+) -> bytes:
     """
     Saves a dictionary of tensors into raw bytes in safetensors format.
 

safetensors/numpy.py

@@ -13,7 +13,9 @@ def _tobytes(tensor: np.ndarray) -> bytes:
     return tensor.tobytes()
 
 
-def save(tensor_dict: Dict[str, np.ndarray], metadata: Optional[Dict[str, str]] = None) -> bytes:
+def save(
+    tensor_dict: Dict[str, np.ndarray], metadata: Optional[Dict[str, str]] = None
+) -> bytes:
     """
     Saves a dictionary of tensors into raw bytes in safetensors format.
 
@@ -38,14 +40,19 @@ def save(tensor_dict: Dict[str, np.ndarray], metadata: Optional[Dict[str, str]]
     byte_data = save(tensors)
     ```
     """
-    flattened = {k: {"dtype": v.dtype.name, "shape": v.shape, "data": _tobytes(v)} for k, v in tensor_dict.items()}
+    flattened = {
+        k: {"dtype": v.dtype.name, "shape": v.shape, "data": _tobytes(v)}
+        for k, v in tensor_dict.items()
+    }
     serialized = serialize(flattened, metadata=metadata)
     result = bytes(serialized)
     return result
 
 
 def save_file(
-    tensor_dict: Dict[str, np.ndarray], filename: Union[str, os.PathLike], metadata: Optional[Dict[str, str]] = None
+    tensor_dict: Dict[str, np.ndarray],
+    filename: Union[str, os.PathLike],
+    metadata: Optional[Dict[str, str]] = None,
 ) -> None:
     """
     Saves a dictionary of tensors into raw bytes in safetensors format.
@@ -73,7 +80,10 @@ def save_file(
     save_file(tensors, "model.safetensors")
     ```
     """
-    flattened = {k: {"dtype": v.dtype.name, "shape": v.shape, "data": _tobytes(v)} for k, v in tensor_dict.items()}
+    flattened = {
+        k: {"dtype": v.dtype.name, "shape": v.shape, "data": _tobytes(v)}
+        for k, v in tensor_dict.items()
+    }
     serialize_file(flattened, filename, metadata=metadata)
 
 

safetensors/paddle.py

@@ -7,7 +7,9 @@ import paddle
 from safetensors import numpy
 
 
-def save(tensors: Dict[str, paddle.Tensor], metadata: Optional[Dict[str, str]] = None) -> bytes:
+def save(
+    tensors: Dict[str, paddle.Tensor], metadata: Optional[Dict[str, str]] = None
+) -> bytes:
     """
     Saves a dictionary of tensors into raw bytes in safetensors format.
 
@@ -98,7 +100,9 @@ def load(data: bytes, device: str = "cpu") -> Dict[str, paddle.Tensor]:
     return _np2paddle(flat, device)
 
 
-def load_file(filename: Union[str, os.PathLike], device="cpu") -> Dict[str, paddle.Tensor]:
+def load_file(
+    filename: Union[str, os.PathLike], device="cpu"
+) -> Dict[str, paddle.Tensor]:
     """
     Loads a safetensors file into paddle format.
 
@@ -126,7 +130,9 @@ def load_file(filename: Union[str, os.PathLike], device="cpu") -> Dict[str, padd
     return output
 
 
-def _np2paddle(numpy_dict: Dict[str, np.ndarray], device: str = "cpu") -> Dict[str, paddle.Tensor]:
+def _np2paddle(
+    numpy_dict: Dict[str, np.ndarray], device: str = "cpu"
+) -> Dict[str, paddle.Tensor]:
     for k, v in numpy_dict.items():
         numpy_dict[k] = paddle.to_tensor(v, place=device)
     return numpy_dict

safetensors/tensorflow.py

@@ -7,7 +7,9 @@ import tensorflow as tf
 from safetensors import numpy, safe_open
 
 
-def save(tensors: Dict[str, tf.Tensor], metadata: Optional[Dict[str, str]] = None) -> bytes:
+def save(
+    tensors: Dict[str, tf.Tensor], metadata: Optional[Dict[str, str]] = None
+) -> bytes:
     """
     Saves a dictionary of tensors into raw bytes in safetensors format.
 

safetensors/torch.py

@@ -2,6 +2,7 @@ import os
 import sys
 from collections import defaultdict
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
+from packaging.version import Version
 
 import torch
 
@@ -41,7 +42,9 @@ def storage_size(tensor: torch.Tensor) -> int:
             return tensor.nelement() * _SIZE[tensor.dtype]
 
 
-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]]:
     filtered_tensors = []
     for shared in tensors:
         if len(shared) < 2:
@@ -69,7 +72,11 @@ def _filter_shared_not_shared(tensors: List[Set[str]], state_dict: Dict[str, tor
 def _find_shared_tensors(state_dict: Dict[str, torch.Tensor]) -> List[Set[str]]:
     tensors = defaultdict(set)
     for k, v in state_dict.items():
-        if v.device != torch.device("meta") and storage_ptr(v) != 0 and storage_size(v) != 0:
+        if (
+            v.device != torch.device("meta")
+            and storage_ptr(v) != 0
+            and storage_size(v) != 0
+        ):
             # Need to add device as key because of multiple GPU.
             tensors[(v.device, storage_ptr(v), storage_size(v))].add(k)
     tensors = list(sorted(tensors.values()))
@@ -78,7 +85,9 @@ def _find_shared_tensors(state_dict: Dict[str, torch.Tensor]) -> List[Set[str]]:
 
 
 def _is_complete(tensor: torch.Tensor) -> bool:
-    return tensor.data_ptr() == storage_ptr(tensor) and tensor.nelement() * _SIZE[tensor.dtype] == storage_size(tensor)
+    return tensor.data_ptr() == storage_ptr(tensor) and tensor.nelement() * _SIZE[
+        tensor.dtype
+    ] == storage_size(tensor)
 
 
 def _remove_duplicate_names(
@@ -97,7 +106,9 @@ def _remove_duplicate_names(
     shareds = _find_shared_tensors(state_dict)
     to_remove = defaultdict(list)
     for shared in shareds:
-        complete_names = set([name for name in shared if _is_complete(state_dict[name])])
+        complete_names = set(
+            [name for name in shared if _is_complete(state_dict[name])]
+        )
         if not complete_names:
             raise RuntimeError(
                 "Error while trying to find names to remove to save state dict, but found no suitable name to keep"
@@ -207,7 +218,9 @@ def load_model(
     """
     state_dict = load_file(filename, device=device)
     model_state_dict = model.state_dict()
-    to_removes = _remove_duplicate_names(model_state_dict, preferred_names=state_dict.keys())
+    to_removes = _remove_duplicate_names(
+        model_state_dict, preferred_names=state_dict.keys()
+    )
 
     reverse_to_remove = {}
     for key, to_remove_group in to_removes.items():
@@ -273,7 +286,9 @@ def load_model(
     return missing, unexpected
 
 
-def save(tensors: Dict[str, torch.Tensor], metadata: Optional[Dict[str, str]] = None) -> bytes:
+def save(
+    tensors: Dict[str, torch.Tensor], metadata: Optional[Dict[str, str]] = None
+) -> bytes:
     """
     Saves a dictionary of tensors into raw bytes in safetensors format.
 
@@ -337,7 +352,9 @@ def save_file(
     serialize_file(_flatten(tensors), filename, metadata=metadata)
 
 
-def load_file(filename: Union[str, os.PathLike], device: Union[str, int] = "cpu") -> Dict[str, torch.Tensor]:
+def load_file(
+    filename: Union[str, os.PathLike], device: Union[str, int] = "cpu"
+) -> Dict[str, torch.Tensor]:
     """
     Loads a safetensors file into torch format.
 
@@ -416,6 +433,14 @@ _SIZE = {
     _float8_e8m0: 1,
     _float4_e2m1_x2: 1,
 }
+if Version(torch.__version__) > Version("2.0.0"):
+    _SIZE.update(
+        {
+            torch.uint64: 8,
+            torch.uint32: 4,
+            torch.uint16: 2,
+        }
+    )
 
 _TYPES = {
     "F64": torch.float64,
@@ -423,17 +448,22 @@ _TYPES = {
     "F16": torch.float16,
     "BF16": torch.bfloat16,
     "I64": torch.int64,
-    # "U64": torch.uint64,
     "I32": torch.int32,
-    # "U32": torch.uint32,
     "I16": torch.int16,
-    # "U16": torch.uint16,
     "I8": torch.int8,
     "U8": torch.uint8,
     "BOOL": torch.bool,
     "F8_E4M3": _float8_e4m3fn,
     "F8_E5M2": _float8_e5m2,
 }
+if Version(torch.__version__) > Version("2.0.0"):
+    _TYPES.update(
+        {
+            "U64": torch.uint64,
+            "U32": torch.uint32,
+            "U16": torch.uint16,
+        }
+    )
 
 
 def _getdtype(dtype_str: str) -> torch.dtype:
@@ -517,12 +547,16 @@ def _tobytes(tensor: torch.Tensor, name: str) -> bytes:
 
 def _flatten(tensors: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, Any]]:
     if not isinstance(tensors, dict):
-        raise ValueError(f"Expected a dict of [str, torch.Tensor] but received {type(tensors)}")
+        raise ValueError(
+            f"Expected a dict of [str, torch.Tensor] but received {type(tensors)}"
+        )
 
     invalid_tensors = []
     for k, v in tensors.items():
         if not isinstance(v, torch.Tensor):
-            raise ValueError(f"Key `{k}` is invalid, expected torch.Tensor but received {type(v)}")
+            raise ValueError(
+                f"Key `{k}` is invalid, expected torch.Tensor but received {type(v)}"
+            )
 
         if v.layout != torch.strided:
             invalid_tensors.append(k)

{safetensors-0.6.0rc0.dist-info → safetensors-0.6.1rc0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: safetensors
-Version: 0.6.0rc0
+Version: 0.6.1rc0
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Education
@@ -28,10 +28,7 @@ Requires-Dist: jaxlib>=0.3.25 ; extra == 'jax'
 Requires-Dist: mlx>=0.0.9 ; extra == 'mlx'
 Requires-Dist: safetensors[numpy] ; extra == 'paddlepaddle'
 Requires-Dist: paddlepaddle>=2.4.1 ; extra == 'paddlepaddle'
-Requires-Dist: black==22.3 ; extra == 'quality'
-Requires-Dist: click==8.0.4 ; extra == 'quality'
-Requires-Dist: isort>=5.5.4 ; extra == 'quality'
-Requires-Dist: flake8>=3.8.3 ; extra == 'quality'
+Requires-Dist: ruff ; extra == 'quality'
 Requires-Dist: safetensors[numpy] ; extra == 'testing'
 Requires-Dist: h5py>=3.7.0 ; extra == 'testing'
 Requires-Dist: huggingface-hub>=0.12.1 ; extra == 'testing'
@@ -39,6 +36,12 @@ Requires-Dist: setuptools-rust>=1.5.2 ; extra == 'testing'
 Requires-Dist: pytest>=7.2.0 ; extra == 'testing'
 Requires-Dist: pytest-benchmark>=4.0.0 ; extra == 'testing'
 Requires-Dist: hypothesis>=6.70.2 ; extra == 'testing'
+Requires-Dist: safetensors[numpy] ; extra == 'testingfree'
+Requires-Dist: huggingface-hub>=0.12.1 ; extra == 'testingfree'
+Requires-Dist: setuptools-rust>=1.5.2 ; extra == 'testingfree'
+Requires-Dist: pytest>=7.2.0 ; extra == 'testingfree'
+Requires-Dist: pytest-benchmark>=4.0.0 ; extra == 'testingfree'
+Requires-Dist: hypothesis>=6.70.2 ; extra == 'testingfree'
 Requires-Dist: safetensors[torch] ; extra == 'all'
 Requires-Dist: safetensors[numpy] ; extra == 'all'
 Requires-Dist: safetensors[pinned-tf] ; extra == 'all'
@@ -56,6 +59,7 @@ Provides-Extra: mlx
 Provides-Extra: paddlepaddle
 Provides-Extra: quality
 Provides-Extra: testing
+Provides-Extra: testingfree
 Provides-Extra: all
 Provides-Extra: dev
 License-File: LICENSE

safetensors-0.6.1rc0.dist-info/RECORD

@@ -0,0 +1,14 @@
+safetensors-0.6.1rc0.dist-info/METADATA,sha256=b9UhFet-V41dU1DjVtpEBwxjqd4KdYo1TXZI1aTCeDI,4144
+safetensors-0.6.1rc0.dist-info/WHEEL,sha256=f_z1-UCPQDLldsOrrwsrPzP3Lu45fZVQJi3w5k9L9Kw,90
+safetensors-0.6.1rc0.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+safetensors/__init__.py,sha256=HYY5VVsb3b-cxhZBwhNx53ZKqSIB4M14nIXLTOAM1Wc,204
+safetensors/__init__.pyi,sha256=tnVaPqYbh8ggFbOZdYKUC4ArqitiWDfrIQt1BNJ377k,4183
+safetensors/_safetensors_rust.pyd,sha256=Ka0Sw1QC34OT7VG3HnygBtlgL76WcpGmATrh9PHdmxg,611328
+safetensors/flax.py,sha256=SnuiGojmth0eCFIWoKEvAfh95nZP9uCZ9E-S4NndrbU,3991
+safetensors/mlx.py,sha256=KvfTWusLSx1hSPWQgg99iL-z9VoD6zQ8l4-RAsCe7P8,3990
+safetensors/numpy.py,sha256=MaUhU4V3J4nDjLmoy0OdeTc6JC8Dq2PBHmnxjQU2bfQ,5189
+safetensors/paddle.py,sha256=B8TLF5MFeqeipUxynSOB_NqPNypL_dE1C3vihtEWj0A,4337
+safetensors/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+safetensors/tensorflow.py,sha256=DajI3qkz00Zy2h7jublSAvTaD51QOPdaIgKQIeSiCRs,4042
+safetensors/torch.py,sha256=vpMhMmierfW5Os6BLlwYgYbb7akoPTBiWPD8fW6nwSE,20849
+safetensors-0.6.1rc0.dist-info/RECORD,,

{safetensors-0.6.0rc0.dist-info → safetensors-0.6.1rc0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: maturin (1.8.7)
+Generator: maturin (1.9.3)
 Root-Is-Purelib: false
 Tag: cp38-abi3-win32

safetensors-0.6.0rc0.dist-info/RECORD

@@ -1,14 +0,0 @@
-safetensors-0.6.0rc0.dist-info/METADATA,sha256=JOcXGydzbxQyLCKSV_L6kQ3DYJdXoKD5R1dJZJRzdCM,3908
-safetensors-0.6.0rc0.dist-info/WHEEL,sha256=SDCbBFz5TSmn0QRHfGLgEdqCqxx9FJdkEZPReKjVInM,90
-safetensors-0.6.0rc0.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-safetensors/__init__.py,sha256=xxEn5gc4JWHjA0wnafb0Wwwq1m2QXwlbm5nsItGRqO0,180
-safetensors/__init__.pyi,sha256=ItiQgMqqoYAuwZH8EWkXh3ApUwGkDTRiLkDoMBIMR3s,3918
-safetensors/_safetensors_rust.pyd,sha256=J8rqSSSTCpRV2DOOg3Peya6Wgu-dvvHg4gGg7Zaf5oA,603648
-safetensors/flax.py,sha256=SnuiGojmth0eCFIWoKEvAfh95nZP9uCZ9E-S4NndrbU,3991
-safetensors/mlx.py,sha256=Bcb3g1LiEbdY_h3r2v_USXgFzdQdwMAI-dlSY4cBGNY,3982
-safetensors/numpy.py,sha256=JKxjDTfr3Z0bqri-WO7LxSwstY58DUP4cAFOcGuoDHY,5120
-safetensors/paddle.py,sha256=ZKyerzx1bwVb-hB10sKQnpSI0z_U-iDknmBEqxVzFRI,4313
-safetensors/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-safetensors/tensorflow.py,sha256=zcg2HyMfSwb3XHtsSQsFm3VgwCoIeRgs73xlPA72Jzc,4034
-safetensors/torch.py,sha256=MhfQdczycyNM0AMmuDNkSwdBkvYbCFPQPfkp-5fWVVc,20315
-safetensors-0.6.0rc0.dist-info/RECORD,,